folly/python/request_context rework Capsule Storage for Memory Leaks

Summary:
Ok I wrote a test to insure the use counts were not leaking.  But they were `test_multiple_ctx_copies_memory_leak` was ending up with a shared_ptr ref count of the number of tasks we spawned using each context.

This was all using shared_ptr::use_count so it wouldn't matter how many python objects were holding the capsule, we were leaking somehow, even with all the tasks/contextvar  deleted.

I wasn't able to track it down, but I replaced the PyCapsule with the existing Context object.    Context is really a PyCapsule with a easier to read API.

Reviewed By: itamaro, prakashgayasen

Differential Revision: D84285368

fbshipit-source-id: d4ba1f29ce1538d8e256144a649d4fc431bf4831

Author: fried

third-party/folly/src/folly/python/request_context.pxd

@@ -26,13 +26,11 @@ cdef extern from "folly/io/async/Request.h" namespace "folly":
         )
 
 
-cdef extern from "folly/python/request_context_capsule.h" namespace "folly::python":
-    cdef object RequestContextToPyCapsule(shared_ptr[RequestContext] ptr)
-    cdef shared_ptr[RequestContext] PyCapsuleToRequestContext(object ptr)
+cdef object RequestContextToPyCapsule(shared_ptr[RequestContext] ptr) noexcept
+cdef shared_ptr[RequestContext] PyCapsuleToRequestContext(object capsule) noexcept
+cdef bint isRequestContextPyCapsule(object capsule) noexcept
 
-cdef class Context:
-    cdef shared_ptr[RequestContext] _ptr
 
-cpdef Context save() noexcept
-cpdef Context get_from_contextvar() noexcept
-cdef object get_PyContext(object context) except *
+cpdef object save() noexcept
+cpdef object get_from_contextvar() noexcept
+cdef object get_PyContext(object context = ?) except *

third-party/folly/src/folly/python/request_context.pyi

@@ -16,11 +16,11 @@ from contextlib import contextmanager
 from typing import Generator
 
 class Context:
-    pass
+    def use_count(self) -> int: ...
 
 @contextmanager
 def active() -> Generator[Context, None, None]: ...
 def create() -> Context: ...
 def set(ctx: Context) -> Context: ...
 def save() -> Context: ...
-def get_from_contextvar() -> Context | None: ...
+def get_from_contextvar() -> Context: ...

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

@@ -18,7 +18,6 @@ from contextlib import contextmanager
 from libcpp.memory cimport make_shared
 from cpython.object cimport PyObject
 from cpython.contextvars cimport get_value_no_default as get_value, PyContextVar_Set, PyContextVar_Reset, PyContextVar_New, PyContext_CheckExact
-from cpython.pycapsule cimport PyCapsule_CheckExact
 from cpython.pystate cimport PyThreadState
 from libcpp.utility cimport move
 
@@ -29,7 +28,7 @@ cdef object _RequestContext = PyContextVar_New("_RequestContext", NULL)
 cdef object set_PyContext(shared_ptr[RequestContext] ptr) except *:
     return PyContextVar_Set(_RequestContext, RequestContextToPyCapsule(move(ptr)))
 
-cdef object get_PyContext(object context) except *:
+cdef object get_PyContext(object context = None) except *:
     """Return the PyCapsule from the ContextVar"""
     if context is None:
         return get_value(_RequestContext)
@@ -42,33 +41,53 @@ cdef object get_PyContext(object context) except *:
 
 @cython.auto_pickle(False)
 cdef class Context:
+    """ Python Wrapper around shared_ptr[folly::RequestContext] """
+    # This is opague to other cython files
+    cdef shared_ptr[RequestContext] _ptr
+
     def __eq__(self, Context other):
         return self._ptr.get() == other._ptr.get()
 
     def __hash__(self):
         return <unsigned long>(self._ptr.get())
 
     def __repr__(self):
-        return f"<{self.__class__!r}: {<unsigned long>self._ptr.get()}>"
+        return f"<Python Capsule shared_ptr<folly::RequestContext>: {<unsigned long>self._ptr.get()}>"
 
     def __bool__(self):
         return self._ptr.get() != NULL
 
+    def use_count(self):
+        return self._ptr.use_count()
+
+
+cdef object RequestContextToPyCapsule(shared_ptr[RequestContext] ptr) noexcept:
+    cdef Context holder = Context.__new__(Context)
+    holder._ptr = move(ptr)
+    return holder
+
+cdef shared_ptr[RequestContext] PyCapsuleToRequestContext(object capsule) noexcept:
+    cdef Context holder = <Context>capsule
+    return holder._ptr
 
-cpdef Context save() noexcept:
+cdef bint isRequestContextPyCapsule(object capsule) noexcept:
+    return isinstance(capsule, Context)
+
+
+cpdef object save() noexcept:
     """ Return the current setcontext """
     cdef Context ctx = Context.__new__(Context)
-    ctx._ptr = RequestContext.saveContext()
+    ctx._ptr = move(RequestContext.saveContext())
     return ctx
 
 
-cpdef Context get_from_contextvar() noexcept:
+cpdef object get_from_contextvar() noexcept:
     """ Return the current context from the contextvar """
     cdef Context ctx = Context.__new__(Context)
 
     ctx_var = get_value(_RequestContext)
-    if ctx_var is not None and PyCapsule_CheckExact(ctx_var):
-        ctx._ptr = PyCapsuleToRequestContext(ctx_var)
+    if isinstance(ctx_var, Context):
+        return ctx_var
     return ctx
 
 
@@ -77,12 +96,13 @@ def active():
     """ Create a Context and shove it into the python contextvar """
     cdef shared_ptr[RequestContext] rctx = make_shared[RequestContext]()
     prev_rctx = RequestContext.setContext(rctx)
-    cdef Context ctx = Context.__new__(Context)
-    ctx._ptr = rctx
+    cdef Context ctx = get_from_contextvar()
+    assert ctx._ptr == rctx
+    rctx.reset() # So we have correct use_count for tests
     try:
         yield ctx
     finally:
-        assert RequestContext.setContext(prev_rctx) == rctx
+        assert RequestContext.setContext(move(prev_rctx)) == ctx._ptr
 
 
 cdef extern from "folly/python/request_context.h":
@@ -107,11 +127,11 @@ cdef int _watcher(PyContextEvent event, PyObject* pycontext):
         return 0
 
     py_ctx = get_value(_RequestContext)
-    if py_ctx is not None and PyCapsule_CheckExact(py_ctx):
+    if isinstance(py_ctx, Context):
         ctx = PyCapsuleToRequestContext(py_ctx)
 
     # This is always called so we don't leak RC between PyContext switches.
-    RequestContext.setContext(ctx)
+    RequestContext.setContext(move(ctx))
     return 0
 
 
@@ -140,7 +160,7 @@ cdef void _setContextWatcher(const shared_ptr[RequestContext]& prev_ctx, const s
         # We don't have a contextvar set and we don't have a context, so we don't need to do anything
         return
 
-    if py_ctx is not None and PyCapsule_CheckExact(py_ctx):
+    if isinstance(py_ctx, Context):
         if PyCapsuleToRequestContext(py_ctx).get() == curr_ctx.get():
             # We triggered this change in our context watcher, so we don't need to do anything
             return

third-party/folly/src/folly/python/request_context_capsule.h

@@ -14,6 +14,8 @@
 
 import asyncio
 import itertools
+import random
+
 from contextvars import copy_context
 from functools import partial
 from unittest import IsolatedAsyncioTestCase
@@ -33,6 +35,79 @@ async def get_Context(pass_ctx) -> request_context.Context | None:
 
 
 class RequestContextTest(IsolatedAsyncioTestCase):
+    def test_use_count(self) -> None:
+        ctx = request_context.Context()
+        # ctx is a nullptr capsule, its use count is 0
+        self.assertEqual(ctx.use_count(), 0)
+        self.assertFalse(ctx)
+        del ctx
+        # Set the frc to something.
+        frc_helper.setContext()
+        # Fetch that ctx, from the PyContext
+        ctx = request_context.get_from_contextvar()
+        # ctx, and folly threadlocal
+        self.assertEqual(ctx.use_count(), 2)
+        new_ctx = request_context.save()
+        self.assertEqual(ctx.use_count(), 3)
+        del new_ctx
+        self.assertEqual(ctx.use_count(), 2)
+        # Set the frc to something else
+        frc_helper.setContext()
+        # Just ctx still exists.
+        self.assertEqual(ctx.use_count(), 1)
+
+    async def test_memory_leaks(self) -> None:
+        with request_context.active() as ctx:
+            # Owners
+            # 1. ctx ( also the same as the PyContext capsule )
+            # 2. folly::RequestContext:: threadlocal
+            self.assertEqual(ctx.use_count(), 2)
+            self.assertEqual(await get_Context(ctx), ctx)
+            self.assertEqual(ctx.use_count(), 2)
+            # Copy_context shouldn't leak
+            context = copy_context()
+            # Still only one capsule.
+            self.assertEqual(ctx.use_count(), 2)
+        # Owner 2 goes away because the FRC is not set
+        self.assertEqual(ctx.use_count(), 1)
+        del context
+        self.assertEqual(ctx.use_count(), 1)
+
+    async def test_multiple_ctx_copies_memory_leak(self) -> None:
+        tasks = []
+        with request_context.active() as ctx1:
+            for _ in range(20):
+                tasks.append(asyncio.create_task(get_Context(ctx1)))
+
+        with request_context.active() as ctx2:
+            for _ in range(28):
+                tasks.append(asyncio.create_task(get_Context(ctx2)))
+        # Single PyCapsule and ctx object
+        self.assertEqual(ctx1.use_count(), 1)
+        self.assertEqual(ctx2.use_count(), 1)
+
+        # Now there are 20 PyContext objects
+
+        async def test() -> None:
+            await asyncio.sleep(0)
+            self.assertEqual(ctx1.use_count(), 1)
+            await asyncio.sleep(0)
+            self.assertEqual(ctx2.use_count(), 1)
+
+        tasks.append(asyncio.create_task(test()))
+
+        random.shuffle(tasks)
+
+        await asyncio.gather(*tasks)
+
+        self.assertEqual(ctx1.use_count(), 1)
+        self.assertEqual(ctx2.use_count(), 1)
+
+        del tasks
+
+        self.assertEqual(ctx1.use_count(), 1)
+        self.assertEqual(ctx2.use_count(), 1)
+
     async def test_request_context(self) -> None:
         with request_context.active() as ctx1:
             # Create task used copy_context() during its creation
@@ -44,6 +119,7 @@ async def test_request_context(self) -> None:
                 request_context.get_from_contextvar(), request_context.save()
             )
             task1 = asyncio.create_task(get_Context(ctx1))
+            self.assertEqual(ctx1.use_count(), 2)  # New Context but same PyCapsule
 
         with request_context.active() as ctx2:
             task2 = asyncio.create_task(get_Context(ctx2))

third-party/folly/src/folly/python/test/request_context.py

@@ -14,7 +14,8 @@
 
 from folly.request_context cimport RequestContext
 from libcpp.memory cimport make_shared, shared_ptr
+from libcpp.utility cimport move
 
 def setContext():
     cdef shared_ptr[RequestContext] ptr = make_shared[RequestContext]()
-    RequestContext.setContext(ptr)
+    RequestContext.setContext(move(ptr))