Merge pull request #747 from uriyyo/fix-context-stack-issue

Fix issue with concurrent contextual stack modification

Author: wwwjfy

mysql_tests/test_engine.py

@@ -38,7 +38,8 @@ async def test_issue_79():
         async with e.acquire():
             pass  # pragma: no cover
     # noinspection PyProtectedMember
-    assert len(e._ctx.get([])) == 0
+    ctx = e._ctx.get()
+    assert ctx and len(ctx.stack) == 0
 
 
 async def test_reuse(engine):
@@ -198,15 +199,15 @@ async def test_lazy(mocker):
     init_size = qsize(engine)
     async with engine.acquire(lazy=True):
         assert qsize(engine) == init_size
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
     assert engine._ctx.get() is None
     assert qsize(engine) == init_size
     async with engine.acquire(lazy=True):
         assert qsize(engine) == init_size
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
         assert await engine.scalar("select 1")
         assert qsize(engine) == init_size - 1
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
     assert engine._ctx.get() is None
     assert qsize(engine) == init_size
 

src/gino/engine.py

@@ -558,22 +558,39 @@ async def prepare(self, clause):
         return await self._execute(clause, (_bypass_no_param,), {}).prepare(clause)
 
 
+_StackCtx = collections.namedtuple("_StackCtx", "task,stack")
+
+
+def _current_task(loop=None):
+    if loop is None:
+        loop = asyncio.get_event_loop()
+
+    if sys.version_info >= (3, 7):
+        return asyncio.current_task(loop=loop)
+    else:
+        return asyncio.Task.current_task(loop=loop)
+
+
 class _ContextualStack:
     __slots__ = ("_ctx", "_stack")
 
     def __init__(self, ctx):
         self._ctx = ctx
-        self._stack = ctx.get()
-        if self._stack is None:
+        curr_ctx = self._ctx.get()
+
+        if curr_ctx is None or curr_ctx.task is not _current_task():
             self._stack = collections.deque()
-            ctx.set(self._stack)
+            self._ctx.set(_StackCtx(_current_task(), self._stack))
+        else:
+            self._stack = curr_ctx.stack
 
     def __bool__(self):
         return bool(self._stack)
 
     @property
     def top(self):
-        return self._stack[-1]
+        if self._stack:
+            return self._stack[-1]
 
     def push(self, value):
         self._stack.append(value)
@@ -738,9 +755,9 @@ def current_connection(self):
         :return: :class:`.GinoConnection`
 
         """
-        stack = self._ctx.get()
-        if stack:
-            return stack[-1].gino_conn
+        ctx = self._ctx.get()
+        if ctx and ctx.stack:
+            return ctx.stack[-1].gino_conn
 
     async def close(self):
         """

tests/test_engine.py

@@ -38,7 +38,8 @@ async def test_issue_79():
         async with e.acquire():
             pass  # pragma: no cover
     # noinspection PyProtectedMember
-    assert len(e._ctx.get([])) == 0
+    ctx = e._ctx.get()
+    assert ctx and len(ctx.stack) == 0
 
 
 async def test_reuse(engine):
@@ -94,6 +95,38 @@ async def test_reuse(engine):
     assert qsize(engine) == init_size
 
 
+async def test_reuse_conn_in_task(request, engine):
+    loop = asyncio.get_event_loop()
+
+    sub_task_result1 = loop.create_future()
+    sub_task_result2 = loop.create_future()
+
+    main_task_check = asyncio.Event()
+
+    async with engine.acquire(reuse=False) as conn:
+        async def _task():
+            async with engine.acquire(reuse=True) as _task_conn:
+                sub_task_result1.set_result(_task_conn.raw_connection is not conn.raw_connection)
+                await asyncio.sleep(0)
+
+            async with engine.acquire(reuse=False) as _task_conn:
+                sub_task_result2.set_result(_task_conn.raw_connection is not conn.raw_connection)
+                await asyncio.sleep(0)
+                await asyncio.wait_for(main_task_check.wait(), 5)
+
+        task = loop.create_task(_task())
+        request.addfinalizer(task.cancel)
+
+        assert await asyncio.wait_for(sub_task_result1, 5)
+        assert await asyncio.wait_for(sub_task_result2, 5)
+
+        async with engine.acquire(reuse=True) as sub_coon:
+            assert conn.raw_connection is sub_coon.raw_connection
+
+        main_task_check.set()
+        await task
+
+
 async def test_compile(engine):
     stmt, params = engine.compile(User.query.where(User.id == 3))
     assert params[0] == 3
@@ -236,15 +269,15 @@ async def test_lazy(mocker):
     init_size = qsize(engine)
     async with engine.acquire(lazy=True):
         assert qsize(engine) == init_size
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
     assert engine._ctx.get() is None
     assert qsize(engine) == init_size
     async with engine.acquire(lazy=True):
         assert qsize(engine) == init_size
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
         assert await engine.scalar("select 1")
         assert qsize(engine) == init_size - 1
-        assert len(engine._ctx.get()) == 1
+        assert len(engine._ctx.get().stack) == 1
     assert engine._ctx.get() is None
     assert qsize(engine) == init_size