feat: record time spent waiting on rpc calls (#472)

Two new attributes have been added to the top level context object,
which can be used in looking at performance of underlying Datastore
calls. `rpc_time` records the total amount of time all rpc calls made
in that context have taken. `wait_time` records the total amount of time
spent waiting (blocking) for rpc calls to complete. Because of
parallelism, `wait_time` should be less than `rpc_time`, although in
practice they're almost always quite close.

Author: Chris Rossi

google/cloud/ndb/__init__.py

@@ -30,6 +30,7 @@
 from google.cloud.ndb.context import Context
 from google.cloud.ndb.context import ContextOptions
 from google.cloud.ndb.context import get_context
+from google.cloud.ndb.context import get_toplevel_context
 from google.cloud.ndb.context import TransactionOptions
 from google.cloud.ndb._datastore_api import EVENTUAL
 from google.cloud.ndb._datastore_api import EVENTUAL_CONSISTENCY
@@ -218,6 +219,7 @@
     "add_flow_exception",
     "Future",
     "get_context",
+    "get_toplevel_context",
     "make_context",
     "make_default_context",
     "QueueFuture",

google/cloud/ndb/_datastore_api.py

@@ -82,6 +82,8 @@ def make_call(rpc_name, request, retries=None, timeout=None):
 
     @tasklets.tasklet
     def rpc_call():
+        context = context_module.get_toplevel_context()
+
         call = method.future(request, timeout=timeout)
         rpc = _remote.RemoteCall(call, "{}({})".format(rpc_name, request))
         log.debug(rpc)
@@ -93,6 +95,8 @@ def rpc_call():
             if isinstance(error, grpc.Call):
                 error = core_exceptions.from_grpc_error(error)
             raise error
+        finally:
+            context.rpc_time += rpc.elapsed_time
 
         raise tasklets.Return(result)
 

google/cloud/ndb/_eventloop.py

@@ -17,6 +17,7 @@
 This should handle both asynchronous ``ndb`` objects and arbitrary callbacks.
 """
 import collections
+import logging
 import uuid
 import time
 
@@ -26,23 +27,8 @@
 except ImportError:  # pragma: NO PY3 COVER
     import Queue as queue
 
-__all__ = [
-    "add_idle",
-    "call_soon",
-    "EventLoop",
-    "get_event_loop",
-    "queue_call",
-    "queue_rpc",
-    "run",
-    "run1",
-]
-
-
-def _logging_debug(*args, **kw):
-    """Placeholder.
-
-    See #6360."""
 
+log = logging.getLogger(__name__)
 
 _Event = collections.namedtuple(
     "_Event", ("when", "callback", "args", "kwargs")
@@ -149,21 +135,21 @@ def clear(self):
             idlers = self.idlers
             queue = self.queue
             rpcs = self.rpcs
-            _logging_debug("Clearing stale EventLoop instance...")
+            log.debug("Clearing stale EventLoop instance...")
             if current:
-                _logging_debug("  current = %s", current)
+                log.debug("  current = %s", current)
             if idlers:
-                _logging_debug("  idlers = %s", idlers)
+                log.debug("  idlers = %s", idlers)
             if queue:
-                _logging_debug("  queue = %s", queue)
+                log.debug("  queue = %s", queue)
             if rpcs:
-                _logging_debug("  rpcs = %s", rpcs)
+                log.debug("  rpcs = %s", rpcs)
             self.__init__()
             current.clear()
             idlers.clear()
             queue[:] = []
             rpcs.clear()
-            _logging_debug("Cleared")
+            log.debug("Cleared")
 
     def insort_event_right(self, event):
         """Insert event in queue with sorting.
@@ -267,12 +253,12 @@ def run_idle(self):
             return False
         idler = self.idlers.popleft()
         callback, args, kwargs = idler
-        _logging_debug("idler: %s", callback.__name__)
+        log.debug("idler: %s", callback.__name__)
         result = callback(*args, **kwargs)
 
         # See add_idle() for meaning of callback return value.
         if result is None:
-            _logging_debug("idler %s removed", callback.__name__)
+            log.debug("idler %s removed", callback.__name__)
         else:
             if result:
                 self.inactive = 0
@@ -292,7 +278,6 @@ def _run_current(self):
 
         self.inactive = 0
         callback, args, kwargs = self.current.popleft()
-        _logging_debug("nowevent: %s", callback.__name__)
         callback(*args, **kwargs)
         return True
 
@@ -312,15 +297,24 @@ def run0(self):
             if delay <= 0:
                 self.inactive = 0
                 _, callback, args, kwargs = self.queue.pop(0)
-                _logging_debug("event: %s", callback.__name__)
+                log.debug("event: %s", callback.__name__)
                 callback(*args, **kwargs)
                 return 0
 
         if self.rpcs:
+            # Avoid circular import
+            from google.cloud.ndb import context as context_module
+
+            context = context_module.get_toplevel_context()
+
             # This potentially blocks, waiting for an rpc to finish and put its
             # result on the queue. Functionally equivalent to the ``wait_any``
             # call that was used here in legacy NDB.
+            start_time = time.time()
             rpc_id, rpc = self.rpc_results.get()
+            elapsed = time.time() - start_time
+            log.debug("Blocked for {}s awaiting RPC results.".format(elapsed))
+            context.wait_time += elapsed
 
             callback = self.rpcs.pop(rpc_id)
             callback(rpc)

google/cloud/ndb/_remote.py

@@ -17,6 +17,7 @@
 # In its own module to avoid circular import between _datastore_api and
 # tasklets modules.
 import grpc
+import time
 
 from google.cloud.ndb import exceptions
 
@@ -39,6 +40,13 @@ class RemoteCall(object):
     def __init__(self, future, info):
         self.future = future
         self.info = info
+        self.start_time = time.time()
+        self.elapsed_time = 0
+
+        def record_time(future):
+            self.elapsed_time = time.time() - self.start_time
+
+        future.add_done_callback(record_time)
 
     def __repr__(self):
         return self.info

google/cloud/ndb/context.py

@@ -26,20 +26,12 @@
 from google.cloud.ndb import tasklets
 
 
-__all__ = [
-    "AutoBatcher",
-    "Context",
-    "ContextOptions",
-    "get_context",
-    "TransactionOptions",
-]
-
-
 class _LocalState(threading.local):
     """Thread local state."""
 
     def __init__(self):
         self.context = None
+        self.toplevel_context = None
 
 
 _state = _LocalState()
@@ -73,6 +65,40 @@ def get_context(raise_context_error=True):
         raise exceptions.ContextError()
 
 
+def get_toplevel_context(raise_context_error=True):
+    """Get the current top level context.
+
+    This function should be called within a context established by
+    :meth:`google.cloud.ndb.client.Client.context`.
+
+    The toplevel context is the context created by the call to
+    :meth:`google.cloud.ndb.client.Client.context`. At times, this context will
+    be superceded by subcontexts, which are used, for example, during
+    transactions. This function will always return the top level context
+    regardless of whether one of these subcontexts is the current one.
+
+    Args:
+        raise_context_error (bool): If set to :data:`True`, will raise an
+            exception if called outside of a context. Set this to :data:`False`
+            in order to have it just return :data:`None` if called outside of a
+            context. Default: :data:`True`
+
+    Returns:
+        Context: The current context.
+
+    Raises:
+        .ContextError: If called outside of a context
+            established by :meth:`google.cloud.ndb.client.Client.context` and
+            ``raise_context_error`` is :data:`True`.
+    """
+    context = _state.toplevel_context
+    if context:
+        return context
+
+    if raise_context_error:
+        raise exceptions.ContextError()
+
+
 def _default_policy(attr_name, value_type):
     """Factory for producing default policies.
 
@@ -192,6 +218,8 @@ def __new__(
         datastore_policy=None,
         on_commit_callbacks=None,
         legacy_data=True,
+        rpc_time=None,
+        wait_time=None,
     ):
         # Prevent circular import in Python 2.7
         from google.cloud.ndb import _cache
@@ -253,11 +281,17 @@ def use(self):
         """
         prev_context = _state.context
         _state.context = self
+        if not prev_context:
+            _state.toplevel_context = self
+            self.rpc_time = 0
+            self.wait_time = 0
         try:
             yield self
         finally:
             if prev_context:
                 prev_context.cache.update(self.cache)
+            else:
+                _state.toplevel_context = None
             _state.context = prev_context
 
     @tasklets.tasklet

tests/unit/test__eventloop.py

@@ -22,16 +22,10 @@
 import grpc
 import pytest
 
-from . import utils
-
 from google.cloud.ndb import exceptions
 from google.cloud.ndb import _eventloop
 
 
-def test___all__():
-    utils.verify___all__(_eventloop)
-
-
 def _Event(when=0, what="foo", args=(), kw={}):
     return _eventloop._Event(when, what, args, kw)
 
@@ -258,6 +252,7 @@ def test_run0_next_now(self, time):
         assert len(loop.queue) == 1
         assert loop.inactive == 0
 
+    @pytest.mark.usefixtures("in_context")
     def test_run0_rpc(self):
         rpc = mock.Mock(spec=grpc.Future)
         callback = mock.Mock(spec=())

tests/unit/test_context.py

@@ -28,11 +28,36 @@
 from google.cloud.ndb import model
 from google.cloud.ndb import _options
 
-from . import utils
 
+class Test_get_context:
+    @staticmethod
+    def test_in_context(in_context):
+        assert context_module.get_context() is in_context
+
+    @staticmethod
+    def test_no_context_raise():
+        with pytest.raises(exceptions.ContextError):
+            context_module.get_context()
+
+    @staticmethod
+    def test_no_context_dont_raise():
+        assert context_module.get_context(False) is None
 
-def test___all__():
-    utils.verify___all__(context_module)
+
+class Test_get_toplevel_context:
+    @staticmethod
+    def test_in_context(in_context):
+        with in_context.new().use():
+            assert context_module.get_toplevel_context() is in_context
+
+    @staticmethod
+    def test_no_context_raise():
+        with pytest.raises(exceptions.ContextError):
+            context_module.get_toplevel_context()
+
+    @staticmethod
+    def test_no_context_dont_raise():
+        assert context_module.get_toplevel_context(False) is None
 
 
 class TestContext: