Use _interpreters.call().

Author: ericsnowcurrently

Lib/concurrent/futures/interpreter.py

@@ -45,12 +45,8 @@ def resolve_task(fn, args, kwargs):
                 # XXX Circle back to this later.
                 raise TypeError('scripts not supported')
             else:
-                # Functions defined in the __main__ module can't be pickled,
-                # so they can't be used here.  In the future, we could possibly
-                # borrow from multiprocessing to work around this.
                 task = (fn, args, kwargs)
-                data = pickle.dumps(task)
-            return data
+            return task
 
         if initializer is not None:
             try:
@@ -65,35 +61,6 @@ def create_context():
             return cls(initdata, shared)
         return create_context, resolve_task
 
-    @classmethod
-    @contextlib.contextmanager
-    def _capture_exc(cls, resultsid):
-        try:
-            yield
-        except BaseException as exc:
-            # Send the captured exception out on the results queue,
-            # but still leave it unhandled for the interpreter to handle.
-            _interpqueues.put(resultsid, (None, exc))
-            raise  # re-raise
-
-    @classmethod
-    def _send_script_result(cls, resultsid):
-        _interpqueues.put(resultsid, (None, None))
-
-    @classmethod
-    def _call(cls, func, args, kwargs, resultsid):
-        with cls._capture_exc(resultsid):
-            res = func(*args or (), **kwargs or {})
-        # Send the result back.
-        with cls._capture_exc(resultsid):
-            _interpqueues.put(resultsid, (res, None))
-
-    @classmethod
-    def _call_pickled(cls, pickled, resultsid):
-        with cls._capture_exc(resultsid):
-            fn, args, kwargs = pickle.loads(pickled)
-        cls._call(fn, args, kwargs, resultsid)
-
     def __init__(self, initdata, shared=None):
         self.initdata = initdata
         self.shared = dict(shared) if shared else None
@@ -104,11 +71,56 @@ def __del__(self):
         if self.interpid is not None:
             self.finalize()
 
-    def _exec(self, script):
-        assert self.interpid is not None
-        excinfo = _interpreters.exec(self.interpid, script, restrict=True)
+    def _call(self, fn, args, kwargs):
+        def do_call(resultsid, func, *args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except BaseException as exc:
+                # Avoid relying on globals.
+                import _interpreters
+                import _interpqueues
+                # Send the captured exception out on the results queue,
+                # but still leave it unhandled for the interpreter to handle.
+                try:
+                    _interpqueues.put(resultsid, exc)
+                except _interpreters.NotShareableError:
+                    # The exception is not shareable.
+                    import sys
+                    import traceback
+                    print('exception is not shareable:', file=sys.stderr)
+                    traceback.print_exception(exc)
+                    _interpqueues.put(resultsid, None)
+                raise  # re-raise
+
+        args = (self.resultsid, fn, *args)
+        res, excinfo = _interpreters.call(self.interpid, do_call, args, kwargs)
         if excinfo is not None:
             raise ExecutionFailed(excinfo)
+        return res
+
+    def _get_exception(self):
+        # Wait for the exception data to show up.
+        while True:
+            try:
+                excdata = _interpqueues.get(self.resultsid)
+            except _interpqueues.QueueNotFoundError:
+                raise  # re-raise
+            except _interpqueues.QueueError as exc:
+                if exc.__cause__ is not None or exc.__context__ is not None:
+                    raise  # re-raise
+                if str(exc).endswith(' is empty'):
+                    continue
+                else:
+                    raise  # re-raise
+            except ModuleNotFoundError:
+                # interpreters.queues doesn't exist, which means
+                # QueueEmpty doesn't.  Act as though it does.
+                continue
+            else:
+                break
+        exc, unboundop = excdata
+        assert unboundop is None, unboundop
+        return exc
 
     def initialize(self):
         assert self.interpid is None, self.interpid
@@ -119,8 +131,6 @@ def initialize(self):
             maxsize = 0
             self.resultsid = _interpqueues.create(maxsize)
 
-            self._exec(f'from {__name__} import WorkerContext')
-
             if self.shared:
                 _interpreters.set___main___attrs(
                                     self.interpid, self.shared, restrict=True)
@@ -148,37 +158,15 @@ def finalize(self):
                 pass
 
     def run(self, task):
-        data = task
-        script = f'WorkerContext._call_pickled({data!r}, {self.resultsid})'
-
+        fn, args, kwargs = task
         try:
-            self._exec(script)
-        except ExecutionFailed as exc:
-            exc_wrapper = exc
-        else:
-            exc_wrapper = None
-
-        # Return the result, or raise the exception.
-        while True:
-            try:
-                obj = _interpqueues.get(self.resultsid)
-            except _interpqueues.QueueNotFoundError:
+            return self._call(fn, args, kwargs)
+        except ExecutionFailed as wrapper:
+            exc = self._get_exception()
+            if exc is None:
+                # The exception must have been not shareable.
                 raise  # re-raise
-            except _interpqueues.QueueError:
-                continue
-            except ModuleNotFoundError:
-                # interpreters.queues doesn't exist, which means
-                # QueueEmpty doesn't.  Act as though it does.
-                continue
-            else:
-                break
-        (res, exc), unboundop = obj
-        assert unboundop is None, unboundop
-        if exc is not None:
-            assert res is None, res
-            assert exc_wrapper is not None
-            raise exc from exc_wrapper
-        return res
+            raise exc from wrapper
 
 
 class BrokenInterpreterPool(_thread.BrokenThreadPool):

Lib/test/test_concurrent_futures/test_interpreter_pool.py

@@ -2,7 +2,7 @@
 import contextlib
 import io
 import os
-import pickle
+import select
 import time
 import unittest
 from concurrent.futures.interpreter import (
@@ -22,10 +22,14 @@ def noop():
 
 
 def write_msg(fd, msg):
+    import os
     os.write(fd, msg + b'\0')
 
 
-def read_msg(fd):
+def read_msg(fd, timeout=10.0):
+    r, _, _ = select.select([fd], [], [], timeout)
+    if fd not in r:
+        raise TimeoutError('nothing to read')
     msg = b''
     while ch := os.read(fd, 1):
         if ch == b'\0':
@@ -121,19 +125,32 @@ def init2():
             nonlocal count
             count += 1
 
-        with self.assertRaises(pickle.PicklingError):
-            self.executor_type(initializer=init1)
-        with self.assertRaises(pickle.PicklingError):
-            self.executor_type(initializer=init2)
+        with contextlib.redirect_stderr(io.StringIO()) as stderr:
+            with self.executor_type(initializer=init1) as executor:
+                fut = executor.submit(lambda: None)
+        self.assertIn('NotShareableError', stderr.getvalue())
+        with self.assertRaises(BrokenInterpreterPool):
+            fut.result()
+
+        with contextlib.redirect_stderr(io.StringIO()) as stderr:
+            with self.executor_type(initializer=init2) as executor:
+                fut = executor.submit(lambda: None)
+        self.assertIn('NotShareableError', stderr.getvalue())
+        with self.assertRaises(BrokenInterpreterPool):
+            fut.result()
 
     def test_init_instance_method(self):
         class Spam:
             def initializer(self):
                 raise NotImplementedError
         spam = Spam()
 
-        with self.assertRaises(pickle.PicklingError):
-            self.executor_type(initializer=spam.initializer)
+        with contextlib.redirect_stderr(io.StringIO()) as stderr:
+            with self.executor_type(initializer=spam.initializer) as executor:
+                fut = executor.submit(lambda: None)
+        self.assertIn('NotShareableError', stderr.getvalue())
+        with self.assertRaises(BrokenInterpreterPool):
+            fut.result()
 
     def test_init_shared(self):
         msg = b'eggs'
@@ -178,8 +195,6 @@ def test_init_exception_in_func(self):
         stderr = stderr.getvalue()
         self.assertIn('ExecutionFailed: Exception: spam', stderr)
         self.assertIn('Uncaught in the interpreter:', stderr)
-        self.assertIn('The above exception was the direct cause of the following exception:',
-                      stderr)
 
     @unittest.expectedFailure
     def test_submit_script(self):
@@ -208,19 +223,24 @@ def task2():
             return spam
 
         executor = self.executor_type()
-        with self.assertRaises(pickle.PicklingError):
-            executor.submit(task1)
-        with self.assertRaises(pickle.PicklingError):
-            executor.submit(task2)
+
+        fut = executor.submit(task1)
+        with self.assertRaises(_interpreters.NotShareableError):
+            fut.result()
+
+        fut = executor.submit(task2)
+        with self.assertRaises(_interpreters.NotShareableError):
+            fut.result()
 
     def test_submit_local_instance(self):
         class Spam:
             def __init__(self):
                 self.value = True
 
         executor = self.executor_type()
-        with self.assertRaises(pickle.PicklingError):
-            executor.submit(Spam)
+        fut = executor.submit(Spam)
+        with self.assertRaises(_interpreters.NotShareableError):
+            fut.result()
 
     def test_submit_instance_method(self):
         class Spam:
@@ -229,8 +249,9 @@ def run(self):
         spam = Spam()
 
         executor = self.executor_type()
-        with self.assertRaises(pickle.PicklingError):
-            executor.submit(spam.run)
+        fut = executor.submit(spam.run)
+        with self.assertRaises(_interpreters.NotShareableError):
+            fut.result()
 
     def test_submit_func_globals(self):
         executor = self.executor_type()
@@ -242,6 +263,7 @@ def test_submit_func_globals(self):
 
     @unittest.expectedFailure
     def test_submit_exception_in_script(self):
+        # Scripts are not supported currently.
         fut = self.executor.submit('raise Exception("spam")')
         with self.assertRaises(Exception) as captured:
             fut.result()
@@ -289,13 +311,20 @@ def test_idle_thread_reuse(self):
         executor.shutdown(wait=True)
 
     def test_pickle_errors_propagate(self):
-        # GH-125864: Pickle errors happen before the script tries to execute, so the
-        # queue used to wait infinitely.
-
+        # GH-125864: Pickle errors happen before the script tries to execute,
+        # so the queue used to wait infinitely.
         fut = self.executor.submit(PickleShenanigans(0))
         with self.assertRaisesRegex(RuntimeError, "gotcha"):
             fut.result()
 
+    def test_no_stale_references(self):
+        # Weak references don't cross between interpreters.
+        raise unittest.SkipTest('not applicable')
+
+    def test_free_reference(self):
+        # Weak references don't cross between interpreters.
+        raise unittest.SkipTest('not applicable')
+
 
 class AsyncioTest(InterpretersMixin, testasyncio_utils.TestCase):