fix: refactor transactions to use their own event loops (#443)

Referring to issue #426, the problem ultimately turned out to be that
we could fall out of the transaction scope and trigger a commit while
there is still work left queued on the event loop, including, in this
case, the tasklet that would eventually schedule the call to delete,
causing the delete to never actually happen.

The fix is to go ahead and consume the event loop queues before
scheduling the call to COMMIT. However, if there are other tasks
happening in parallel, this can really mess with the natural sequence of
events in ways that can cause things to blow up. (All of the
`parallel_transaction` tests in `tests/system/test_misc.py` for
instance, will fail.) The fix for that is to give each transaction its
own event loop, so that when it calls `_eventloop.run` prior to commit,
it is only flushing tasks that pertain to it.

Fixes #426

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/_batch.py

@@ -14,8 +14,6 @@
 
 """Support for batching operations."""
 
-from google.cloud.ndb import _eventloop
-
 
 def get_batch(batch_cls, options=None):
     """Gets a data structure for storing batched calls to Datastore Lookup.
@@ -68,5 +66,5 @@ def idle():
         return idle
 
     batches[options_key] = batch = batch_cls(options)
-    _eventloop.add_idle(idler(batch))
+    context.eventloop.add_idle(idler(batch))
     return batch

packages/google-cloud-ndb/google/cloud/ndb/_datastore_api.py

@@ -519,6 +519,19 @@ def commit_callback(rpc):
         rpc.add_done_callback(commit_callback)
 
 
+def prepare_to_commit(transaction):
+    """Signal that we're ready to commit a transaction.
+
+    Currently just used to signal to the commit batch that we're not going to
+    need to call `AllocateIds`, because we're ready to commit now.
+
+    Args:
+        transaction (bytes): The transaction id about to be committed.
+    """
+    batch = _get_commit_batch(transaction, _options.Options())
+    batch.preparing_to_commit = True
+
+
 def commit(transaction, retries=None, timeout=None):
     """Commit a transaction.
 
@@ -605,6 +618,7 @@ def __init__(self, transaction, options):
         self.allocating_ids = []
         self.incomplete_mutations = []
         self.incomplete_futures = []
+        self.preparing_to_commit = False
 
     def put(self, entity_pb):
         """Add an entity to batch to be stored.
@@ -657,8 +671,9 @@ def delete(self, key):
 
     def idle_callback(self):
         """Call AllocateIds on any incomplete keys in the batch."""
-        if not self.incomplete_mutations:
-            # This will happen if `commit` is called first.
+        # If there are no incomplete mutations, or if we're already preparing
+        # to commit, there's no need to allocate ids.
+        if self.preparing_to_commit or not self.incomplete_mutations:
             return
 
         # Signal to a future commit that there is an id allocation in
@@ -728,11 +743,6 @@ def commit(self, retries=None, timeout=None):
             if not future.done():
                 yield future
 
-        # Head off making any more AllocateId calls. Any remaining incomplete
-        # keys will get ids as part of the Commit call.
-        self.incomplete_mutations = []
-        self.incomplete_futures = []
-
         future = tasklets.Future("Commit")
         futures = self.futures
 

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

@@ -34,7 +34,6 @@
     "queue_call",
     "queue_rpc",
     "run",
-    "run0",
     "run1",
 ]
 
@@ -396,13 +395,7 @@ def run():
     loop.run()
 
 
-def run0():
-    """Calls :method:`EventLoop.run0` on current event loop."""
-    loop = get_event_loop()
-    loop.run0()
-
-
 def run1():
     """Calls :method:`EventLoop.run1` on current event loop."""
     loop = get_event_loop()
-    loop.run1()
+    return loop.run1()

packages/google-cloud-ndb/google/cloud/ndb/_transaction.py

@@ -130,15 +130,36 @@ def _transaction_async(context, callback, read_only=False):
     tx_context = context.new(
         transaction=transaction_id,
         on_commit_callbacks=on_commit_callbacks,
-        cache=None,  # Use new, empty cache for transaction
+        batches=None,
+        commit_batches=None,
+        cache=None,
+        # We could just pass `None` here and let the `Context` constructor
+        # instantiate a new event loop, but our unit tests inject a subclass of
+        # `EventLoop` that makes testing a little easier. This makes sure the
+        # new event loop is of the same type as the current one, to propagate
+        # the event loop class used for testing.
+        eventloop=type(context.eventloop)(),
     )
+
+    # The outer loop is dependent on the inner loop
+    def run_inner_loop(inner_context):
+        with inner_context.use():
+            if inner_context.eventloop.run1():
+                return True  # schedule again
+
+    context.eventloop.add_idle(run_inner_loop, tx_context)
+
     with tx_context.use():
         try:
             # Run the callback
             result = callback()
             if isinstance(result, tasklets.Future):
                 result = yield result
 
+            # Make sure we've run everything we can run before calling commit
+            _datastore_api.prepare_to_commit(transaction_id)
+            tx_context.eventloop.run()
+
             # Commit the transaction
             yield _datastore_api.commit(transaction_id, retries=0)
 

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

@@ -123,7 +123,10 @@ def wait(self):
         after a call to this method.
         """
         while not self._done:
-            _eventloop.run1()
+            if not _eventloop.run1():
+                raise RuntimeError(
+                    "Eventloop is exhausted with unfinished futures."
+                )
 
     def check_success(self):
         """Check whether a future has completed without raising an exception.
@@ -348,16 +351,20 @@ def done_callback(yielded):
 
             error = yielded.exception()
             if error:
-                _eventloop.call_soon(self._advance_tasklet, error=error)
+                self.context.eventloop.call_soon(
+                    self._advance_tasklet, error=error
+                )
             else:
-                _eventloop.call_soon(self._advance_tasklet, yielded.result())
+                self.context.eventloop.call_soon(
+                    self._advance_tasklet, yielded.result()
+                )
 
         if isinstance(yielded, Future):
             yielded.add_done_callback(done_callback)
             self.waiting_on = yielded
 
         elif isinstance(yielded, _remote.RemoteCall):
-            _eventloop.queue_rpc(yielded, done_callback)
+            self.context.eventloop.queue_rpc(yielded, done_callback)
             self.waiting_on = yielded
 
         elif isinstance(yielded, (list, tuple)):
@@ -515,7 +522,10 @@ def wait_any(futures):
             if future.done():
                 return future
 
-        _eventloop.run1()
+        if not _eventloop.run1():
+            raise RuntimeError(
+                "Eventloop is exhausted with unfinished futures."
+            )
 
 
 def wait_all(futures):

packages/google-cloud-ndb/tests/system/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+import operator
 import time
 
 KIND = "SomeKind"
@@ -61,3 +63,23 @@ def eventually(f, predicate, timeout=120, interval=2):
         time.sleep(interval)
 
     assert predicate(value)
+
+
+def length_equals(n):
+    """Returns predicate that returns True if passed a sequence of length `n`.
+
+    For use with `eventually`.
+    """
+
+    def predicate(sequence):
+        return len(sequence) == n
+
+    return predicate
+
+
+def equals(n):
+    """Returns predicate that returns True if passed `n`.
+
+    For use with `eventually`.
+    """
+    return functools.partial(operator.eq, n)

packages/google-cloud-ndb/tests/system/test_crud.py

@@ -16,8 +16,6 @@
 System tests for Create, Update, Delete. (CRUD)
 """
 import datetime
-import functools
-import operator
 import os
 import pickle
 import random
@@ -37,15 +35,11 @@
 from google.cloud.ndb import _cache
 from google.cloud.ndb import global_cache as global_cache_module
 
-from tests.system import KIND, eventually
+from tests.system import KIND, eventually, equals
 
 USE_REDIS_CACHE = bool(os.environ.get("REDIS_CACHE_URL"))
 
 
-def _equals(n):
-    return functools.partial(operator.eq, n)
-
-
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity(ds_entity):
     entity_id = test_utils.system.unique_resource_id()
@@ -526,7 +520,7 @@ class SomeKind(ndb.Model):
         # Sneaky. Delete entity out from under cache so we know we're getting
         # cached copy.
         key.delete()
-        eventually(key.get, _equals(None))
+        eventually(key.get, equals(None))
 
     retrieved = key.get()
     assert retrieved.foo == 42
@@ -772,6 +766,29 @@ def delete_entity():
     assert key.get() is None
 
 
+def test_delete_entity_in_transaction_with_global_cache(
+    client_context, ds_entity
+):
+    """Regression test for #426
+
+    https://github.com/googleapis/python-ndb/issues/426
+    """
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    entity_id = test_utils.system.unique_resource_id()
+    ds_entity(KIND, entity_id, foo=42)
+
+    global_cache = global_cache_module._InProcessGlobalCache()
+    with client_context.new(global_cache=global_cache).use():
+        key = ndb.Key(KIND, entity_id)
+        assert key.get().foo == 42
+
+        ndb.transaction(key.delete)
+        assert key.get() is None
+
+
 @pytest.mark.usefixtures("client_context")
 def test_delete_entity_in_transaction_then_rollback(ds_entity):
     entity_id = test_utils.system.unique_resource_id()

packages/google-cloud-ndb/tests/system/test_misc.py

@@ -22,6 +22,8 @@
 
 from google.cloud import ndb
 
+from tests.system import eventually, length_equals
+
 USE_REDIS_CACHE = bool(os.environ.get("REDIS_CACHE_URL"))
 
 
@@ -271,3 +273,25 @@ def update(id, add, fail=False):
 
         entity = SomeKind.get_by_id(id)
         assert entity.foo == 142
+
+
+@pytest.mark.usefixtures("client_context")
+def test_insert_entity_in_transaction_without_preallocating_id(dispose_of):
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+        bar = ndb.StringProperty()
+
+    def save_entity():
+        # By not waiting on the Future, we don't force a call to AllocateIds
+        # before the transaction is committed.
+        SomeKind(foo=42, bar="none").put_async()
+
+    ndb.transaction(save_entity)
+
+    query = SomeKind.query()
+    eventually(query.fetch, length_equals(1))
+    retrieved = query.fetch()[0]
+    dispose_of(retrieved._key._key)
+
+    assert retrieved.foo == 42
+    assert retrieved.bar == "none"