fork_safe_greenlet.py

import gc
import gevent
import greenlet
import multiprocessing
import os

"""
Gevent doesn't cope well with forking. The child process after the
fork ends up with all the same greenlets running as the parent
process, which is almost certainly not what you wanted.

This module works around that by providing two functions that are
called before and after you fork (in the _parent_ process). The
before_fork() function sets a greenlet tracer (yes, it'll temporarily
disable any greenlet tracer you have already set, but it'll restore it
in after_fork()) which detects when gevent attempts to switch to
another greenlet after the fork, and kills all the greenlets other
than the one in which you called before_fork(), i.e., the one in which
you're doing the fork.

Then, the after_fork() function disables the tracer in the parent and
puts back any tracer you had set before.

_Don't call after_fork() in the child process._

You can also use start_process rather than multiprocessing.Process. It
will create and start the process and return it, handling the
before_fork() and after_fork() for you.

"""

old_tracer = None


def before_fork():
    global old_tracer
    current = gevent.getcurrent()
    parent_pid = os.getpid()

    def tracer(*args):
        if os.getpid() == parent_pid:
            return
        for ob in gc.get_objects():
            if isinstance(ob, greenlet.greenlet) and ob != current \
                    and not isinstance(ob, gevent.hub.Hub):
                gevent.kill(ob)
        greenlet.settrace(None)

    old_tracer = greenlet.settrace(tracer)


def after_fork():
    greenlet.settrace(old_tracer)


def start_process(*args, **kwargs):
    p = multiprocessing.Process(*args, **kwargs)
    before_fork()
    p.start()
    after_fork()
    return p