Add MRU executor cache to folly/python get_running_executor_drive

Reviewed By: MatzeB

Differential Revision: D116524117

fbshipit-source-id: 608749f4a9fa96bc08ad0aee1a6cea80f5dc0e7e

Author: generatedunixname1395027625275998

third-party/folly/src/folly/python/executor.pyx

@@ -20,11 +20,25 @@ from libcpp.memory cimport make_unique, unique_ptr
 from cython.operator cimport dereference as deref
 from weakref import WeakKeyDictionary
 from cpython.ref cimport PyObject
+from cpython.weakref cimport PyWeakref_NewRef, PyWeakref_GetObject
 
 # asyncio Loops to AsyncioExecutor
 loop_to_q = WeakKeyDictionary()
 _RaiseKeyError = object()
 
+# weak reference to the last seen event loop
+_last_loop = None
+# AsyncioExecutor for the last seen event loop
+_last_executor = None
+
+
+# cleanup callback that clears _last_loop and _last_executor when the loop
+# referenced by _last_loop is collected
+def _clean_last_executor(_wr):
+    global _last_loop, _last_executor
+    _last_loop = None
+    _last_executor = None
+
 
 cdef class AsyncioExecutor:
     def __cinit__(self):
@@ -112,28 +126,41 @@ cdef cAsyncioExecutor* get_running_executor(bint running) noexcept:
 
 cdef cAsyncioExecutor* get_running_executor_drive(
     bint running, bint driveBeforeDealloc) noexcept:
+    global _last_loop, _last_executor
     try:
         if running:
             loop = asyncio.get_running_loop()
         else:
             loop = asyncio.get_event_loop()
     except RuntimeError:
         return NULL
-    try:
-        executor = <AsyncioExecutor>(loop_to_q[loop])
-    except KeyError:
-        if sys.platform == "win32":
-            executor = ProactorExecutor(loop._proactor._iocp)
-            queue = IocpQueue(executor)
-            queue.swap(loop)
-        else:
-            executor = NotificationQueueAsyncioExecutor(driveBeforeDealloc)
-            loop.add_reader(executor.fileno(), executor.drive)
-        loop_to_q[loop] = executor
+    cdef AsyncioExecutor executor = None
+    if _last_loop is not None and <PyObject*>loop is PyWeakref_GetObject(_last_loop):
+        executor = _last_executor
+    if executor is None:
+        try:
+            executor = <AsyncioExecutor>(loop_to_q[loop])
+        except KeyError:
+            if sys.platform == "win32":
+                proactor = ProactorExecutor(loop._proactor._iocp)
+                queue = IocpQueue(proactor)
+                queue.swap(loop)
+                executor = proactor
+            else:
+                nq = NotificationQueueAsyncioExecutor(driveBeforeDealloc)
+                loop.add_reader(nq.fileno(), nq.drive)
+                executor = nq
+            loop_to_q[loop] = executor
+        _last_loop = PyWeakref_NewRef(loop, _clean_last_executor)
+        _last_executor = executor
     return executor._executor
 
 
 cdef int set_executor_for_loop(object loop, cAsyncioExecutor* c_executor) noexcept:
+    global _last_loop, _last_executor
+    # drop the fast-path cache to avoid returning a stale executor
+    _last_loop = None
+    _last_executor = None
     if c_executor == NULL:
         del loop_to_q[loop]
         return 0