Offer a prepare_subprocess_call to make a func multiprocessing safe

When multiprocessing uses another start method different of 'fork', it
requires that the target function and its arguments can be pickled.

In particular, multiprocessing needs to pickle them in the parent
process and un-pickle them in the child process.

Due how pickling works in Python, pickling a function only involves
storing the information needed to reload it: the bytecode is never
stored.

This makes pickling particular tricky in byexample: if a
we want to call a function that it is from one of the byexample
modules/plugins, the pickling will fail.

It will not fail when multiprocessing serialize it (dumps) but when it
deserialize it (loads) because in the child process, the byexample
modules/plugins will not be loaded in sys.modules and because they are
not in the sys.path (in principle), Python will not be able to find
them.

_prepare_subprocess_call() can wrap the target function and its arguments
and replacing the target by a _subprocess_trampoline function that will
call the former.

This _subprocess_trampoline will do all the bootstraping needed in the
child process, including the (re)loading of the modules/plugins, to make
the un-pickling work.

Most of these details are hidden from the user (plugin developer). He/she
is only required to call prepare_subprocess_call() and use the returned
target/arguments in replace of his/her.

The prepare_subprocess_call() is a partial bound function of
_prepare_subprocess_call(). The former can be obtained optionally from
the extension constructor (__init__ method of ExampleParser,
ExampleFinder, ExampleRunner, ZoneDelimiter and Concern).

The function is thread-safe and it will accessible both in the main
byexample process and in each worker thread.

Author: eldipa

byexample/cfg.py

@@ -19,6 +19,9 @@ class Config(collections.abc.Mapping):
 
         The rest of the keys-values must be of immutable types.
 
+        A special key 'prepare_subprocess_call' is expected to be a partial
+        function, which it is assumed to have immutable bound parameters.
+
         To use Config in a multithread/multiprocess environment,
         you need to call copy() to get an independent copy to work with.
 
@@ -39,7 +42,10 @@ def __len__(self):
 
     def _ensure_cfg_is_constant(self):
         const_types = (int, frozenset, str, bool, bytes, type(None))
-        exception_keys = ('options', 'output', 'registry', 'namespaces')
+        exception_keys = (
+            'options', 'output', 'registry', 'namespaces',
+            'prepare_subprocess_call'
+        )
         for k, v in self._d.items():
             if k in exception_keys:
                 continue
@@ -58,6 +64,9 @@ def copy(self, patch={}):
             Assuming that all the values are constant (immutable),
             the copy does not incur in the costs of a real copy.
 
+            This includes 'prepare_subprocess_call' for which only
+            a reference is copied and not a full copy.
+
             Exceptions to this are:
              - options: which a real copy is made (see Options.copy)
              - output: a file which it is NOT copied

byexample/init.py

@@ -1,5 +1,5 @@
 from __future__ import unicode_literals
-import sys, pkgutil, inspect, pprint, os, operator, traceback
+import sys, pkgutil, inspect, pprint, os, operator, traceback, functools
 import importlib.util
 
 from itertools import chain as chain_iters
@@ -627,6 +627,21 @@ def init_byexample(args, sharer):
     cfg['use_colors'] &= are_tty_colors_supported(cfg['output'])
     cfg['selected_languages'] = frozenset(args.languages)
 
+    # Make a partial application of _prepare_subprocess_call(), binding
+    # all the necessary parameters to import and register the byexample
+    # modules (again) in a subprocess.
+    #
+    # The bound parameters are constant so the function, despite having
+    # state, it is actually stateless (its state is constant, immutable)
+    # Moreover, _prepare_subprocess_call() is thread-safe so the resulting
+    # partial-bound function is thread-safe too.
+    #
+    # See _prepare_subprocess_call().
+    prepare_subprocess_call = functools.partial(
+        _prepare_subprocess_call, dirnames=tuple(args.modules_dirs)
+    )
+    cfg['prepare_subprocess_call'] = prepare_subprocess_call
+
     registry, namespaces_by_class = load_modules(args.modules_dirs, cfg)
 
     allowed_languages = get_allowed_languages(registry, args.languages)
@@ -692,6 +707,70 @@ def init_byexample(args, sharer):
     return testfiles, Config(cfg)
 
 
+def _subprocess_trampoline(
+    dirnames, serialized_func, serialized_args, serialized_kwargs
+):
+    # All of this happens in the *child* process
+    # We reload the modules if they weren't loaded yet
+    # and only then we deserialize the target function and we
+    # call it.
+    #
+    # If _subprocess_trampoline is called in a fresh subprocess,
+    # we are sure that no module was loaded yet however, it is
+    # possible that the user runs a subprocess using forking
+    # which makes a copy of the python process (parent) and therefore
+    # it will have the modules loaded already.
+    #
+    # By the moment it is unclear if in addition to the loading we want
+    # to do more like the instantiation of the plugins.
+    from .init import import_and_register_modules_iter
+    _ = list(import_and_register_modules_iter(dirnames))
+
+    import multiprocessing.reduction
+    fpickler = multiprocessing.reduction.ForkingPickler
+    target = fpickler.loads(serialized_func)
+    args = fpickler.loads(serialized_args)
+    kwargs = fpickler.loads(serialized_kwargs)
+
+    return target(*args, **kwargs)
+
+
+def _prepare_subprocess_call(target, dirnames, *, args=(), kwargs={}):
+    ''' Prepare the given target function to be executable in a separated
+        process (child process).
+
+        The preparation includes the (re)import and (re)registration of
+        the modules found in <dirnames>, once loaded by byexample in the parent
+        process.
+
+        This re-import and re-registration within the child process
+        is needed because the child may be an independent fresh Python
+        process without any idea of how to load byexample modules.
+
+        _prepare_subprocess_call returns a dictionary with keys 'target'
+        and 'args' suitable to call multiprocessing.Process.
+
+        Note: no user code should call _prepare_subprocess_call directly.
+        Instead, call a partial bound function from the Config cfg object
+        given to each extension (ExampleFinder, ExampleParser, Concern, ...).
+        This partial function will not require the <dirnames> argument.
+    '''
+    # Implementation note: this function must be thread-safe because it
+    # may be called from different threads.
+    import multiprocessing.reduction
+    fpickler = multiprocessing.reduction.ForkingPickler
+
+    serialized_func = bytes(fpickler.dumps(target))
+    serialized_args = bytes(fpickler.dumps(args))
+    serialized_kwargs = bytes(fpickler.dumps(kwargs))
+
+    trampoline_args = (
+        dirnames, serialized_func, serialized_args, serialized_kwargs
+    )
+
+    return {'target': _subprocess_trampoline, 'args': trampoline_args}
+
+
 @profile
 def init_worker(cfg, job_num):
     ''' Initialize a worker with worker/job number is passed