Cleanup thread state in python-compat runtime. (runtime:python27, vm:true)

bryanmau · pcostell · commit 45dbb9e0b4a0 · 2016-03-25T13:13:05.000-07:00
Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=118080230
diff --git a/ndb/google_imports.py b/ndb/google_imports.py
@@ -18,6 +18,7 @@
 All other NDB code should import its Google App Engine modules from
 this module.  If necessary, add new imports here (in both places).
 """
+# pylint: disable=unused-import, g-import-not-at-top
 
 import os
 import sys
@@ -86,6 +87,10 @@ def set_appengine_imports():
   from google.appengine.ext.blobstore import blobstore as ext_blobstore
   from google.appengine.ext import db
   from google.appengine.ext import gql
+  try:
+    from google.appengine.ext.vmruntime import callback
+  except ImportError:
+    callback = None
   from google.appengine.runtime import apiproxy_errors
   from google.net.proto import ProtocolBuffer
 else:
@@ -107,6 +112,7 @@ def set_appengine_imports():
   from google3.apphosting.ext.blobstore import blobstore as ext_blobstore
   from google3.apphosting.ext import db
   from google3.apphosting.ext import gql
+  from google3.apphosting.ext.vmruntime import callback
   from google3.apphosting.runtime import apiproxy_errors
   from google3.net.proto import ProtocolBuffer
   # Prospective search is optional.
diff --git a/ndb/tasklets.py b/ndb/tasklets.py
@@ -80,6 +80,7 @@ def foo():
 import os
 import sys
 import types
+import weakref
 
 from .google_imports import apiproxy_stub_map
 from .google_imports import apiproxy_rpc
@@ -88,6 +89,7 @@ def foo():
 from .google_imports import datastore_pbs
 from .google_imports import datastore_rpc
 from .google_imports import namespace_manager
+from .google_imports import callback as _request_callback
 
 from . import eventloop
 from . import utils
@@ -102,6 +104,8 @@ def foo():
 
 _logging_debug = utils.logging_debug
 
+_CALLBACK_KEY = '__CALLBACK__'
+
 
 def _is_generator(obj):
   """Helper to test for a generator object.
@@ -115,12 +119,22 @@ def _is_generator(obj):
 class _State(utils.threading_local):
   """Hold thread-local state."""
 
-  current_context = None
-
   def __init__(self):
     super(_State, self).__init__()
+    self.current_context = None
+    self.all_generators = weakref.WeakSet()
     self.all_pending = set()
 
+  def set_context(self, ctx):
+    if _request_callback and _CALLBACK_KEY not in os.environ:
+      _request_callback.SetRequestEndCallback(self.reset)
+      os.environ[_CALLBACK_KEY] = '1'
+    self.current_context = ctx
+
+  def add_generator(self, gen):
+    _logging_debug('all_generators: add %s', gen)
+    self.all_generators.add(gen)
+
   def add_pending(self, fut):
     _logging_debug('all_pending: add %s', fut)
     self.all_pending.add(fut)
@@ -132,9 +146,18 @@ def remove_pending(self, fut, status='success'):
     else:
       _logging_debug('all_pending: %s: not found %s', status, fut)
 
+  def clear_all_generators(self):
+    if self.all_generators:
+      _logging_debug('all_generators: clear %s', self.all_generators)
+      for gen in self.all_generators:
+        gen.close()
+      self.all_generators.clear()
+    else:
+      _logging_debug('all_generators: clear no-op')
+
   def clear_all_pending(self):
     if self.all_pending:
-      logging.info('all_pending: clear %s', self.all_pending)
+      _logging_debug('all_pending: clear %s', self.all_pending)
       self.all_pending.clear()
     else:
       _logging_debug('all_pending: clear no-op')
@@ -149,6 +172,13 @@ def dump_all_pending(self, verbose=False):
       pending.append(line)
     return '\n'.join(pending)
 
+  def reset(self):
+    self.current_context = None
+    ev = eventloop.get_event_loop()
+    ev.clear()
+    self.clear_all_pending()
+    self.clear_all_generators()
+
 
 _state = _State()
 
@@ -1021,6 +1051,7 @@ def tasklet_wrapper(*args, **kwds):
     if _is_generator(result):
       ns = namespace_manager.get_namespace()
       ds_conn = datastore._GetConnection()
+      _state.add_generator(result)
       eventloop.queue_call(None, fut._help_tasklet_along, ns, ds_conn, result)
     else:
       fut.set_result(result)
@@ -1210,7 +1241,7 @@ def _make_cloud_datastore_context(app_id, external_app_ids=()):
 def set_context(new_context):
   # XXX Docstring
   os.environ[_CONTEXT_KEY] = '1'
-  _state.current_context = new_context
+  _state.set_context(new_context)
 
 class _ThrowingStub(object):
   """A Stub implementation which always throws a NotImplementedError."""
