Repository: mars-project/mars

Issue:
[BUG] Reference cycle exists in promise module
**Describe the bug**
In ``promise`` module, we wrap functions to handle succeeding promises with ``Promise._wrap_hander``. The promise object itself is contained in the ``__closure__`` of wrapper function. Hence a reference cycle exists:

![image](https://user-images.githubusercontent.com/8284922/59827929-6ca4c700-936c-11e9-8134-11b62b3870fc.png)

If some large objects are referenced in the wrapped function, they will be held in memory until GC condition is triggered, which may cause unexpected memory condition.


Discussion:

Resolving patch:
diff --git a/mars/promise.py b/mars/promise.py
--- a/mars/promise.py
+++ b/mars/promise.py
@@ -38,8 +38,6 @@ def __init__(self, resolve=None, reject=None, done=False, failed=False,
                  args=None, kwargs=None):
         # use random promise id
         self._id = struct.pack('<Q', id(self)) + np.random.bytes(32)
-        # register in global pool to reject gc collection
-        _promise_pool[self._id] = self
 
         self._accept_handler = self._wrap_handler(resolve)
         self._reject_handler = self._wrap_handler(reject)
@@ -62,6 +60,9 @@ def __init__(self, resolve=None, reject=None, done=False, failed=False,
             self._accepted = False
         else:
             self._accepted = None
+            # register in global pool to reject gc collection when calls not finished
+            _promise_pool[self._id] = self
+
         self._args = args or ()
         self._kwargs = kwargs or {}
 
@@ -69,6 +70,7 @@ def __init__(self, resolve=None, reject=None, done=False, failed=False,
 
     def __del__(self):
         self.pre_destroy()
+        self._clear_result_cache()
 
     @property
     def id(self):
@@ -89,26 +91,35 @@ def _wrap_handler(self, func):
 
         @wraps(func)
         def _wrapped(*args, **kwargs):
+            _promise_pool.pop(self.id, None)
+
+            result = None
             try:
                 result = func(*args, **kwargs)
                 if isinstance(result, Promise):
                     # the function itself returns a promise object
                     # bind returned promise object to current promise
+                    assert result._next_item is None
                     result._bind_item = self
                     if result._accepted is not None:
                         # promise already done, we move next
                         args = result._args or ()
                         kwargs = result._kwargs or {}
+                        result._clear_result_cache()
+
                         kwargs['_accept'] = result._accepted
-                        result.step_next(*args, **kwargs)
+                        return result._internal_step_next, args, kwargs
                 else:
                     # return non-promise result, we just step next
-                    self.step_next(result)
-            except:
+                    return self._internal_step_next, (result,), {}
+            except:  # noqa: E722
                 # error occurred when executing func, we reject with exc_info
                 logger.exception('Exception met in executing promise.')
                 exc = sys.exc_info()
-                self.step_next(*exc, _accept=False)
+                self._clear_result_cache()
+                return self._internal_step_next, exc, dict(_accept=False)
+            finally:
+                del result
         return _wrapped
 
     def _get_bind_root(self):
@@ -119,9 +130,10 @@ def _get_bind_root(self):
         """
         target = self
         while target._bind_item is not None:
-            if target and target.id in _promise_pool:
+            if target.id in _promise_pool:
                 # remove binder that will not be used later
-                del _promise_pool[target.id]
+                _promise_pool.pop(target.id, None)
+            target._clear_result_cache()
             target = target._bind_item
         return target
 
@@ -134,11 +146,12 @@ def _get_handling_promise(p, handler_attr):
         :rtype: Promise
         """
         while getattr(p, handler_attr) is None:
+            p._accept_handler = p._reject_handler = None
+
             p = p._get_bind_root()  # type: Promise
             if p and p._next_item is not None:
-                if p.id in _promise_pool:
-                    # remove predecessor that will not be used later
-                    del _promise_pool[p.id]
+                _promise_pool.pop(p.id, None)
+                p._clear_result_cache()
                 p = p._next_item
             else:
                 break
@@ -149,11 +162,16 @@ def _log_unexpected_error(args):
         if args and len(args) == 3 and issubclass(args[0], Exception):
             logger.exception('Unhandled exception in promise', exc_info=args)
 
-    def step_next(self, *args, **kwargs):
+    def _clear_result_cache(self):
+        self._args = ()
+        self._kwargs = {}
+
+    def _internal_step_next(self, *args, **kwargs):
         """
-        Step into next promise with given args and kwargs
+        Actual call to step promise call into the next step.
+        If there are further steps, the function will return it
+        as a (func, args, kwargs) tuple to reduce memory footprint.
         """
-
         accept = kwargs.pop('_accept', True)
         target_promise = self
 
@@ -162,45 +180,74 @@ def step_next(self, *args, **kwargs):
         try:
             root_promise = self._get_bind_root()
 
-            if root_promise and root_promise.id in _promise_pool:
-                del _promise_pool[root_promise.id]
+            if root_promise:
+                _promise_pool.pop(root_promise.id, None)
 
             target_promise = root_promise._next_item
-            root_promise._accepted = self._accepted
-            root_promise._args = args
-            root_promise._kwargs = kwargs
+            root_promise._accepted = accept
             if not target_promise:
+                root_promise._args = args
+                root_promise._kwargs = kwargs
                 if not accept:
                     self._log_unexpected_error(args)
                 return
+            else:
+                root_promise._clear_result_cache()
 
+            next_call = None
             if accept:
                 acceptor = self._get_handling_promise(target_promise, '_accept_handler')
                 if acceptor and acceptor._accept_handler:
-                    acceptor._accept_handler(*args, **kwargs)
+                    # remove the handler in promise in case that
+                    # function closure is not freed
+                    handler, acceptor._accept_handler = acceptor._accept_handler, None
+                    next_call = handler(*args, **kwargs)
                 else:
                     acceptor._accepted = accept
                     acceptor._args = args
                     acceptor._kwargs = kwargs
+                    _promise_pool.pop(acceptor.id, None)
             else:
                 rejecter = self._get_handling_promise(target_promise, '_reject_handler')
                 if rejecter and rejecter._reject_handler:
-                    rejecter._reject_handler(*args, **kwargs)
+                    # remove the handler in promise in case that
+                    # function closure is not freed
+                    handler, rejecter._reject_handler = rejecter._reject_handler, None
+                    next_call = handler(*args, **kwargs)
                 else:
                     rejecter._accepted = accept
                     rejecter._args = args
                     rejecter._kwargs = kwargs
+                    _promise_pool.pop(rejecter.id, None)
                     self._log_unexpected_error(args)
+            return next_call
         finally:
-            if target_promise and target_promise.id in _promise_pool:
-                del _promise_pool[target_promise.id]
+            del args, kwargs
+            if target_promise:
+                _promise_pool.pop(target_promise.id, None)
+
+    def step_next(self, args_and_kwargs=None):
+        """
+        Step into next promise with given args and kwargs
+        """
+        if args_and_kwargs is None:
+            args_and_kwargs = [(), {}]
+        step_call = (self._internal_step_next,) + tuple(args_and_kwargs)
+        del args_and_kwargs[:]
+
+        while step_call is not None:
+            func, args, kwargs = step_call
+            step_call = func(*args, **kwargs)
 
     def then(self, on_fulfilled, on_rejected=None):
         promise = Promise(on_fulfilled, on_rejected)
+        assert self._bind_item is None
         self._next_item = promise
         if self._accepted is not None:
-            self._kwargs['_accept'] = self._accepted
-            self.step_next(*self._args, **self._kwargs)
+            args_and_kwargs = [self._args, self._kwargs]
+            args_and_kwargs[1]['_accept'] = self._accepted
+            self._clear_result_cache()
+            self.step_next(args_and_kwargs)
         return promise
 
     def catch(self, on_rejected):
@@ -303,7 +350,7 @@ def _wrapped(*args, **kwargs):
 
         try:
             return func(*args, **kwargs)
-        except:
+        except:  # noqa: E722
             actor = args[0]
             logger.exception('Unhandled exception in promise call')
             if callback:
@@ -317,15 +364,18 @@ class PromiseActor(FunctionActor):
     """
     Actor class providing promise functionality
     """
-    def promise_ref(self, *args, **kwargs):
-        """
-        Wraps an existing ActorRef into a promise ref
-        """
+    def _prepare_promise_registration(self):
         if not hasattr(self, '_promises'):
             self._promises = dict()
             self._promise_ref_keys = dict()
             self._ref_key_promises = dict()
 
+    def promise_ref(self, *args, **kwargs):
+        """
+        Wraps an existing ActorRef into a promise ref
+        """
+        self._prepare_promise_registration()
+
         if not args and not kwargs:
             ref = self.ref()
         elif args and isinstance(args[0], ActorRef):
@@ -334,6 +384,39 @@ def promise_ref(self, *args, **kwargs):
             ref = self.ctx.actor_ref(*args, **kwargs)
         return PromiseRefWrapper(ref, self)
 
+    def spawn_promised(self, func, *args, **kwargs):
+        """
+        Run func asynchronously in a pool and returns a promise.
+        The running of the function does not block current actor.
+
+        :param func: function to run
+        :return: promise
+        """
+        self._prepare_promise_registration()
+
+        if not hasattr(self, '_async_group'):
+            self._async_group = self.ctx.asyncpool()
+
+        p = Promise()
+        ref = self.ref()
+        self.register_promise(p, ref)
+
+        def _wrapped(*a, **kw):
+            try:
+                result = func(*a, **kw)
+            except:  # noqa: E722
+                ref.handle_promise(p.id, *sys.exc_info(), **dict(_accept=False, _tell=True))
+            else:
+                ref.handle_promise(p.id, result, _tell=True)
+            finally:
+                del a, kw
+
+        try:
+            self._async_group.spawn(_wrapped, *args, **kwargs)
+        finally:
+            del args, kwargs
+        return p
+
     def register_promise(self, promise, ref):
         """
         Register a promise into the actor with referrer info
@@ -399,7 +482,7 @@ def reject_promise_refs(self, refs, *args, **kwargs):
                 p = self.get_promise(promise_id)
                 if p is None:
                     continue
-                p.step_next(*args, **kwargs)
+                p.step_next([args, kwargs])
         return handled_refs
 
     def tell_promise(self, callback, *args, **kwargs):
@@ -418,9 +501,9 @@ def handle_promise(self, promise_id, *args, **kwargs):
         """
         p = self.get_promise(promise_id)
         if p is None:
-            logger.warning('Promise %r reentered in %s', promise_id, self.uid)
+            logger.warning('Promise %r reentered or not registered in %s', promise_id, self.uid)
             return
-        self.get_promise(promise_id).step_next(*args, **kwargs)
+        self.get_promise(promise_id).step_next([args, kwargs])
         self.delete_promise(promise_id)
 
     def handle_promise_timeout(self, promise_id):
@@ -429,12 +512,12 @@ def handle_promise_timeout(self, promise_id):
         :param promise_id: promise key
         """
         p = self.get_promise(promise_id)
-        if not p or p._accepted is not None:
+        if p is None or p._accepted is not None:
             # skip promises that are already finished
             return
 
         self.delete_promise(promise_id)
-        p.step_next(*build_exc_info(PromiseTimeout), **dict(_accept=False))
+        p.step_next([build_exc_info(PromiseTimeout), dict(_accept=False)])
 
 
 def all_(promises):
@@ -444,7 +527,7 @@ def all_(promises):
     :param promises: collection of promises
     :return: the new promise
     """
-    promises = list(promises)
+    promises = [p for p in promises if isinstance(p, Promise)]
     new_promise = Promise()
     finish_set = set()
 
@@ -459,16 +542,12 @@ def _handle_reject(*args, **kw):
         if new_promise._accepted is not None:
             return
         for p in promises:
-            try:
-                del _promise_pool[p.id]
-            except KeyError:
-                pass
+            _promise_pool.pop(p.id, None)
         kw['_accept'] = False
-        new_promise.step_next(*args, **kw)
+        new_promise.step_next([args, kw])
 
     for p in promises:
-        if isinstance(p, Promise):
-            p.then(_build_then(p), _handle_reject)
+        p.then(_build_then(p), _handle_reject)
 
     if promises:
         return new_promise
@@ -477,6 +556,10 @@ def _handle_reject(*args, **kw):
         return new_promise
 
 
+def get_active_promise_count():
+    return len(_promise_pool)
+
+
 def finished(*args, **kwargs):
     """
     Create a promise already finished


Test patch:
diff --git a/mars/tests/test_promise.py b/mars/tests/test_promise.py
--- a/mars/tests/test_promise.py
+++ b/mars/tests/test_promise.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import functools
 import sys
 import threading
@@ -33,7 +32,7 @@ def __init__(self):
 
     @promise.reject_on_exception
     def serve(self, value, delay=None, accept=True, raises=False, callback=None):
-        self.ctx.sleep(delay or 0.1)
+        self.ctx.sleep(delay if delay is not None else 0.1)
         if raises:
             raise ValueError('User-induced error')
         self._result_list.append(value)
@@ -55,6 +54,9 @@ def __init__(self):
     def get_finished(self):
         return self._finished
 
+    def reset_finished(self):
+        self._finished = False
+
     def test_normal(self):
         self._finished = False
 
@@ -79,6 +81,18 @@ def test_error_raise(self):
             .catch(lambda *_: ref.serve(-1, _promise=True)) \
             .then(lambda *_: setattr(self, '_finished', True))
 
+    def test_spawn(self, raises=False):
+        def _task():
+            self.ctx.sleep(0.2)
+            if raises:
+                raise SystemError
+
+        ref = self.promise_ref('ServeActor')
+        promise.all_([self.spawn_promised(_task) for _ in range(4)]) \
+            .then(lambda *_: ref.serve(0, delay=0, _promise=True)) \
+            .catch(lambda *exc: ref.serve(exc[0].__name__, delay=0, _promise=True)) \
+            .then(lambda *_: setattr(self, '_finished', True))
+
     def test_all_promise(self):
         self._finished = False
 
@@ -150,6 +164,23 @@ def _rejecter(*exc):
             .then(lambda *_: setattr(self, '_finished', True))
         self.reject_dead_endpoints([self.address], *build_exc_info(WorkerDead))
 
+    def test_closure_refcount(self, content=''):
+        ref = self.promise_ref('ServeActor', address=self.address)
+
+        class Intermediate(object):
+            def __init__(self, s):
+                self.str = s
+
+        new_content = Intermediate('Content: %s' % (content,))
+
+        def _acceptor(*_):
+            ref.serve(weakref.ref(new_content))
+            return 'Processed: %s' % new_content.str
+
+        ref.serve(0, delay=0.5, _promise=True) \
+            .then(_acceptor) \
+            .then(lambda *_: setattr(self, '_finished', True))
+
 
 def _raise_exception(exc):
     raise exc
@@ -159,7 +190,7 @@ def wait_test_actor_result(ref, timeout):
     import gevent
     t = time.time()
     while not ref.get_finished():
-        gevent.sleep(0.1)
+        gevent.sleep(0.01)
         if time.time() > t + timeout:
             raise TimeoutError
 
@@ -180,7 +211,7 @@ def test_thread_body():
                     break
                 value_list.append(('thread_body', v))
                 time.sleep(time_unit)
-                promises[idx].step_next(v, _accept=success)
+                promises[idx].step_next([(v,), dict(_accept=success)])
 
         try:
             thread = threading.Thread(target=test_thread_body)
@@ -391,7 +422,6 @@ def gen_promise(value, accept=True):
                  ('catch', 2)]
             )
         finally:
-            gc.collect()
             self.assertDictEqual(promise._promise_pool, {})
             req_queue.put((None, None, None))
 
@@ -406,12 +436,13 @@ def testPromiseActor(self):
                 self.assertListEqual(serve_ref.get_result(), list(range(11)))
 
                 serve_ref.clear_result()
+                test_ref.reset_finished()
 
                 test_ref.test_error_raise()
                 wait_test_actor_result(test_ref, 10)
                 self.assertListEqual(serve_ref.get_result(), [-1])
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
 
     def testAll(self):
         promises = weakref.WeakValueDictionary()
@@ -427,7 +458,7 @@ def test_thread_body():
                     break
                 value_list.append(('thread_body', v))
                 time.sleep(time_unit)
-                promises[idx].step_next(v, _accept=success)
+                promises[idx].step_next([(v,), dict(_accept=success)])
 
         def gen_promise(value, accept=True):
             p = promise.Promise()
@@ -484,9 +515,38 @@ def _gen_all_promise(*_):
             self.assertListEqual(value_list[:len(expected)], expected)
             time.sleep(0.5)
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
             req_queue.put((None, None, None))
 
+    def testSpawnPromisedActor(self):
+        try:
+            with create_actor_pool(n_process=1) as pool:
+                serve_ref = pool.create_actor(ServeActor, uid='ServeActor')
+                test_ref = pool.create_actor(PromiseTestActor)
+
+                start_time = time.time()
+                test_ref.test_spawn()
+                self.assertLess(time.time() - start_time, 0.1)
+                wait_test_actor_result(test_ref, 30)
+                self.assertEqual(serve_ref.get_result(), [0])
+
+                self.assertGreater(time.time() - start_time, 0.2)
+                self.assertLess(time.time() - start_time, 0.3)
+
+                serve_ref.clear_result()
+                test_ref.reset_finished()
+
+                start_time = time.time()
+                test_ref.test_spawn(raises=True)
+                self.assertLess(time.time() - start_time, 0.1)
+                wait_test_actor_result(test_ref, 30)
+                self.assertEqual(serve_ref.get_result(), ['SystemError'])
+
+                self.assertGreater(time.time() - start_time, 0.2)
+                self.assertLess(time.time() - start_time, 0.3)
+        finally:
+            self.assertEqual(promise.get_active_promise_count(), 0)
+
     def testAllActor(self):
         try:
             with create_actor_pool(n_process=1) as pool:
@@ -494,14 +554,14 @@ def testAllActor(self):
                 test_ref = pool.create_actor(PromiseTestActor)
 
                 test_ref.test_all_promise()
-                gc.collect()
+
                 wait_test_actor_result(test_ref, 30)
                 self.assertListEqual(
                     serve_ref.get_result(),
                     [-128] + list(range(0, 20, 2)) + list(range(1, 20, 2)) + [127]
                 )
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
 
     def testTimeoutActor(self):
         try:
@@ -510,11 +570,11 @@ def testTimeoutActor(self):
                 test_ref = pool.create_actor(PromiseTestActor)
 
                 test_ref.test_timeout()
-                gc.collect()
+
                 wait_test_actor_result(test_ref, 30)
                 self.assertListEqual(serve_ref.get_result(), [0, 'PromiseTimeout'])
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
 
     def testNoTimeoutActor(self):
         try:
@@ -523,12 +583,12 @@ def testNoTimeoutActor(self):
                 test_ref = pool.create_actor(PromiseTestActor)
 
                 test_ref.test_no_timeout()
-                gc.collect()
+
                 wait_test_actor_result(test_ref, 30)
 
                 self.assertListEqual(serve_ref.get_result(), [0])
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
 
     def testRefReject(self):
         try:
@@ -537,11 +597,11 @@ def testRefReject(self):
                 test_ref = pool.create_actor(PromiseTestActor)
 
                 test_ref.test_ref_reject()
-                gc.collect()
+
                 wait_test_actor_result(test_ref, 30)
                 self.assertListEqual(serve_ref.get_result(), [0, 'WorkerProcessStopped'])
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
 
     def testAddrReject(self):
         try:
@@ -550,8 +610,20 @@ def testAddrReject(self):
                 test_ref = pool.create_actor(PromiseTestActor)
 
                 test_ref.test_addr_reject()
-                gc.collect()
+
                 wait_test_actor_result(test_ref, 30)
                 self.assertListEqual(serve_ref.get_result(), [0, 'WorkerDead'])
         finally:
-            self.assertDictEqual(promise._promise_pool, {})
+            self.assertEqual(promise.get_active_promise_count(), 0)
+
+    def testClosureRefcount(self):
+        try:
+            with create_actor_pool(n_process=1) as pool:
+                serve_ref = pool.create_actor(ServeActor, uid='ServeActor')
+                test_ref = pool.create_actor(PromiseTestActor)
+
+                test_ref.test_closure_refcount()
+                wait_test_actor_result(test_ref, 30)
+                self.assertIsNone(serve_ref.get_result()[-1]())
+        finally:
+            self.assertEqual(promise.get_active_promise_count(), 0)