From aadcd6f47b31f2cd3acaca57ba073343b43a3dfe Mon Sep 17 00:00:00 2001
From: Yury Selivanov <yury@magic.io>
Date: Tue, 7 Aug 2018 12:56:06 -0400
Subject: [PATCH] Fix a memory leak related to contextvars support. (#192)

Initial patch and memleak discovery by Victor K. @hellysmile.

Fixes #191.
---
 tests/test_context.py | 24 ++++++++++++++++++++++++
 uvloop/cbhandles.pyx  | 18 ++++++++++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tests/test_context.py b/tests/test_context.py
index c57e359..c167abe 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -3,6 +3,7 @@ import decimal
 import random
 import sys
 import unittest
+import weakref
 
 from uvloop import _testbase as tb
 
@@ -117,6 +118,29 @@ class _ContextBaseTests:
 
         self.assertEqual(cvar.get(), -1)
 
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_task_context_4(self):
+        import contextvars
+        cvar = contextvars.ContextVar('cvar', default='nope')
+
+        class TrackMe:
+            pass
+        tracked = TrackMe()
+        ref = weakref.ref(tracked)
+
+        async def sub():
+            cvar.set(tracked)
+            self.loop.call_soon(lambda: None)
+
+        async def main():
+            await self.loop.create_task(sub())
+
+        task = self.loop.create_task(main())
+        self.loop.run_until_complete(task)
+
+        del tracked
+        self.assertIsNone(ref())
+
 
 class Test_UV_Context(_ContextBaseTests, tb.UVTestCase):
 
diff --git a/uvloop/cbhandles.pyx b/uvloop/cbhandles.pyx
index c342064..031a7ec 100644
--- a/uvloop/cbhandles.pyx
+++ b/uvloop/cbhandles.pyx
@@ -15,9 +15,11 @@ cdef class Handle:
             self._source_traceback = extract_stack()
 
     cdef inline _set_context(self, object context):
+        cdef PyContext* current_context
+
         if PY37:
             if context is None:
-                context = <object>PyContext_CopyCurrent()
+                context = copy_current_context()
             self.context = context
         else:
             if context is not None:
@@ -179,7 +181,7 @@ cdef class TimerHandle:
 
         if PY37:
             if context is None:
-                context = <object>PyContext_CopyCurrent()
+                context = copy_current_context()
             self.context = context
         else:
             if context is not None:
@@ -400,3 +402,15 @@ cdef extract_stack():
 
     stack.reverse()
     return stack
+
+
+cdef copy_current_context():
+    cdef PyContext* current_context
+
+    if PY37:
+        current_context = PyContext_CopyCurrent()
+        py_context = <object>current_context
+        Py_XDECREF(<PyObject*>current_context)
+        return py_context
+
+    raise NotImplementedError('"contextvars" support requires Python 3.7+')