EventLoop and tasklets refactor. (#73)

It was found, through working query iterators, that when a TaskletFuture
calls ``_advance_tasklet`` in its done callback, if it calls it
directly, we can get a really deep call stack that eventually reaches
the limit for maximum recursion. This refactors that to add the
``_advance_tasklet`` call to the eventloop to be run soon, rather than
calling it directly, which fixes this issue by keeping the call stack
shallow. The legacy NDB code was actually already using this
indirection and I had changed it to a direct call because I couldn't
figure out why the legacy code didn't just do that, already. I knew
there was a chance that the reason would reveal itself through trial and
error, and it eventually did.

I also performed a minor refactor on the eventloop itself, adding the
``call_soon`` method to handle the case that was previously handled by
passing ``None`` for the ``delay`` argument to ``queue_call``. This just
makes that API a little bit cleaner.

Author: Chris Rossi

packages/google-cloud-ndb/src/google/cloud/ndb/_eventloop.py

@@ -25,6 +25,7 @@
 
 __all__ = [
     "add_idle",
+    "call_soon",
     "EventLoop",
     "get_event_loop",
     "queue_call",
@@ -80,25 +81,14 @@ class EventLoop:
     other futures were waiting on those results and results derived from those
     results.
 
-    This is somewhat of a work in progress. Initially this was ported (cargo
-    culted) from legacy NDB without a clear understanding of how all the pieces
-    would fit together or what all the different features were actually for. As
-    we've been forced to do some things a little differently with the rewrite,
-    it's not entirely clear that all of the features here have a purpose in the
-    rewrite, but it's still early to say definitively.
-
     Currently, these are the seperate queues used by the event loop in the
     order they are checked by :meth:`~EventLoop.run1`. For each call to
     :meth:`~EventLoop.run1`, the first thing it finds is called:
 
-        current: These callbacks are called first, if there are any. In legacy
-            NDB, these were used by tasklets to queue calls to
-            ``_help_tasklet_along`` when a result from a yielded future was
-            ready. With the rewrite, I haven't seen any reason not to just go
-            ahead and call :meth:`~tasklets.TaskletFuture._advance_tasklet`
-            immediately when a result is available. If a good reason becomes
-            apparent in the course of the rewrite, this is subject to change.
-            Currently, nothing uses this.
+        current: These callbacks are called first, if there are any. Currently
+            this is used to schedule calls to
+            :meth:`tasklets.TaskletFuture._advance_tasklet` when it's time to
+            send a tasklet a value that it was previously waiting on.
 
         idlers: Effectively, these are the same as ``current``, but just get
             called afterwards. These currently are used for batching certain
@@ -113,8 +103,7 @@ class EventLoop:
             time.
 
         queue: These are callbacks that are supposed to be run at (or after) a
-            certain time. Nothing uses these currently. It's not clear, yet,
-            what the use case was in legacy NDB.
+            certain time. This is used by :function:`tasklets.sleep`.
 
         rpcs: If all other queues are empty, and we are waiting on results of a
             gRPC call, then we'll call :method:`queue.Queue.get` on the
@@ -125,7 +114,8 @@ class EventLoop:
 
     Atrributes:
         current (deque): a FIFO list of (callback, args, kwds). These callbacks
-            run immediately when the eventloop runs. Not currently used.
+            run immediately when the eventloop runs. Used by tasklets to
+            schedule calls to :meth:`tasklets.TaskletFuture._advance_tasklet`.
         idlers (deque): a FIFO list of (callback, args, kwds). Thes callbacks
             run only when no other RPCs need to be fired first. Used for
             batching calls to the Datastore back end.
@@ -205,6 +195,16 @@ def insort_event_right(self, event):
                 low = mid + 1
         queue.insert(low, event)
 
+    def call_soon(self, callback, *args, **kwargs):
+        """Schedule a function to be called soon, without a delay.
+
+        Arguments:
+            callback (callable): The function to eventually call.
+            *args: Positional arguments to be passed to callback.
+            **kwargs: Keyword arguments to be passed to callback.
+        """
+        self.current.append((callback, args, kwargs))
+
     def queue_call(self, delay, callback, *args, **kwargs):
         """Schedule a function call at a specific time in the future.
 
@@ -216,10 +216,6 @@ def queue_call(self, delay, callback, *args, **kwargs):
             *args: Positional arguments to be passed to callback.
             **kwargs: Keyword arguments to be passed to callback.
         """
-        if delay is None:
-            self.current.append((callback, args, kwargs))
-            return
-
         when = time.time() + delay if delay < 1e9 else delay
         event = _Event(when, callback, args, kwargs)
         self.insort_event_right(event)
@@ -379,6 +375,12 @@ def add_idle(callback, *args, **kwargs):
     loop.add_idle(callback, *args, **kwargs)
 
 
+def call_soon(callback, *args, **kwargs):
+    """Calls :method:`EventLoop.call_soon` on current event loop. """
+    loop = get_event_loop()
+    loop.call_soon(callback, *args, **kwargs)
+
+
 def queue_call(delay, callback, *args, **kwargs):
     """Calls :method:`EventLoop.queue_call` on current event loop. """
     loop = get_event_loop()

packages/google-cloud-ndb/src/google/cloud/ndb/tasklets.py

@@ -234,8 +234,8 @@ class _TaskletFuture(Future):
     completed and either returned a value or raised an exception.
 
     Args:
-        typing.Generator[Union[tasklets.Future, _remote.RemoteCall], Any, Any]: The
-            generator.
+        typing.Generator[Union[tasklets.Future, _remote.RemoteCall], Any, Any]:
+            The generator.
     """
 
     def __init__(self, generator, context, info="Unknown"):
@@ -276,17 +276,15 @@ def done_callback(yielded):
             # To be called when a future dependency has completed.  Advance the
             # tasklet with the yielded value or error.
             #
-            # It might be worth noting that legacy NDB added a callback to the
-            # event loop which, in turn, called _help_tasklet_along. I don't
-            # see a compelling reason not to go ahead and call _advance_tasklet
-            # immediately here, rather than queue it up to be called soon by
-            # the event loop. This is subject to change if the reason for the
-            # indirection in the original implementation becomes apparent.
+            # It was tempting to call `_advance_tasklet` (`_help_tasklet_along`
+            # in Legacy) directly. Doing so, it has been found, can lead to
+            # exceeding the maximum recursion depth. Queing it up to run on the
+            # event loop avoids this issue by keeping the call stack shallow.
             error = yielded.exception()
             if error:
-                self._advance_tasklet(error=error)
+                _eventloop.call_soon(self._advance_tasklet, error=error)
             else:
-                self._advance_tasklet(yielded.result())
+                _eventloop.call_soon(self._advance_tasklet, yielded.result())
 
         if isinstance(yielded, Future):
             yielded.add_done_callback(done_callback)

packages/google-cloud-ndb/tests/conftest.py

@@ -24,11 +24,18 @@
 
 from google.cloud import environment_vars
 from google.cloud.ndb import context as context_module
+from google.cloud.ndb import _eventloop
 from google.cloud.ndb import model
 
 import pytest
 
 
+class TestingEventLoop(_eventloop.EventLoop):
+    def call_soon(self, callback, *args, **kwargs):
+        """For testing, call the callback immediately."""
+        callback(*args, **kwargs)
+
+
 @pytest.fixture(autouse=True)
 def reset_state(environ):
     """Reset module and class level runtime state.
@@ -76,7 +83,9 @@ def context():
     client = mock.Mock(
         project="testing", namespace=None, spec=("project", "namespace")
     )
-    context = context_module.Context(client, stub=mock.Mock(spec=()))
+    context = context_module.Context(
+        client, stub=mock.Mock(spec=()), eventloop=TestingEventLoop()
+    )
     return context
 
 

packages/google-cloud-ndb/tests/unit/test__datastore_api.py

@@ -764,7 +764,8 @@ def test_commit(datastore_commit, process_commit, in_context):
         rpc = tasklets.Future("_datastore_commit")
         datastore_commit.return_value = rpc
 
-        eventloop = mock.Mock(spec=("queue_rpc", "run"))
+        eventloop = mock.Mock(spec=("queue_rpc", "run", "call_soon"))
+        eventloop.call_soon = lambda f, *args, **kwargs: f(*args, **kwargs)
         with in_context.new(eventloop=eventloop).use():
             future = batch.commit()
 
@@ -788,7 +789,8 @@ def test_commit_error(datastore_commit, process_commit, in_context):
         rpc = tasklets.Future("_datastore_commit")
         datastore_commit.return_value = rpc
 
-        eventloop = mock.Mock(spec=("queue_rpc", "run"))
+        eventloop = mock.Mock(spec=("queue_rpc", "run", "call_soon"))
+        eventloop.call_soon = lambda f, *args, **kwargs: f(*args, **kwargs)
         with in_context.new(eventloop=eventloop).use():
             future = batch.commit()
 
@@ -824,7 +826,8 @@ def test_commit_allocating_ids(
         rpc = tasklets.Future("_datastore_commit")
         datastore_commit.return_value = rpc
 
-        eventloop = mock.Mock(spec=("queue_rpc", "run"))
+        eventloop = mock.Mock(spec=("queue_rpc", "run", "call_soon"))
+        eventloop.call_soon = lambda f, *args, **kwargs: f(*args, **kwargs)
         with in_context.new(eventloop=eventloop).use():
             future = batch.commit()
 

packages/google-cloud-ndb/tests/unit/test__eventloop.py

@@ -122,14 +122,14 @@ def test_insert_event_right_collision(self):
             _Event(2, "baz"),
         ]
 
-    def test_queue_call_now(self):
+    def test_call_soon(self):
         loop = self._make_one()
-        loop.queue_call(None, "foo", "bar", baz="qux")
+        loop.call_soon("foo", "bar", baz="qux")
         assert list(loop.current) == [("foo", ("bar",), {"baz": "qux"})]
         assert not loop.queue
 
     @unittest.mock.patch("google.cloud.ndb._eventloop.time")
-    def test_queue_call_soon(self, time):
+    def test_queue_call_delay(self, time):
         loop = self._make_one()
         time.time.return_value = 5
         loop.queue_call(5, "foo", "bar", baz="qux")
@@ -214,7 +214,7 @@ def test_run0_nothing_to_do(self):
     def test_run0_current(self):
         callback = unittest.mock.Mock(__name__="callback")
         loop = self._make_one()
-        loop.queue_call(None, callback, "foo", bar="baz")
+        loop.call_soon(callback, "foo", bar="baz")
         loop.inactive = 88
         assert loop.run0() == 0
         callback.assert_called_once_with("foo", bar="baz")
@@ -275,7 +275,7 @@ def test_run1_nothing_to_do(self):
     def test_run1_has_work_now(self, time):
         callback = unittest.mock.Mock(__name__="callback")
         loop = self._make_one()
-        loop.queue_call(None, callback)
+        loop.call_soon(callback)
         assert loop.run1() is True
         time.sleep.assert_not_called()
         callback.assert_called_once_with()
@@ -304,7 +304,7 @@ def mock_sleep(seconds):
         runlater = unittest.mock.Mock(__name__="runlater")
         loop = self._make_one()
         loop.add_idle(idler)
-        loop.queue_call(None, runnow)
+        loop.call_soon(runnow)
         loop.queue_call(5, runlater)
         loop.run()
         idler.assert_called_once_with()
@@ -328,6 +328,13 @@ def test_add_idle(context):
         loop.add_idle.assert_called_once_with("foo", "bar", baz="qux")
 
 
+def test_call_soon(context):
+    loop = unittest.mock.Mock(spec=("run", "call_soon"))
+    with context.new(eventloop=loop).use():
+        _eventloop.call_soon("foo", "bar", baz="qux")
+        loop.call_soon.assert_called_once_with("foo", "bar", baz="qux")
+
+
 def test_queue_call(context):
     loop = unittest.mock.Mock(spec=("run", "queue_call"))
     with context.new(eventloop=loop).use():