locks: Restore permits when async context entry is cancelled

tornado/locks.py

@@ -12,6 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import asyncio
 import collections
 import datetime
 import types
@@ -459,7 +460,15 @@ def __exit__(
         self.__enter__()
 
     async def __aenter__(self) -> None:
-        await self.acquire()
+        waiter = asyncio.ensure_future(self.acquire())
+        try:
+            await waiter
+        except asyncio.CancelledError:
+            # A release may have handed us the permit before cancellation
+            # prevented entry into the context (and thus __aexit__).
+            if not waiter.cancelled():
+                self.release()
+            raise
 
     async def __aexit__(
         self,
@@ -571,7 +580,15 @@ def __exit__(
         self.__enter__()
 
     async def __aenter__(self) -> None:
-        await self.acquire()
+        waiter = asyncio.ensure_future(self.acquire())
+        try:
+            await waiter
+        except asyncio.CancelledError:
+            # A release may have handed us the permit before cancellation
+            # prevented entry into the context (and thus __aexit__).
+            if not waiter.cancelled():
+                self.release()
+            raise
 
     async def __aexit__(
         self,

tornado/test/locks_test.py

@@ -530,5 +530,58 @@ def test_context_manager_misuse(self):
                 pass
 
 
+class AsyncContextManagerCancellationTest(AsyncTestCase):
+    @gen_test
+    async def test_cancel_after_handoff(self):
+        for cls in (locks.Semaphore, locks.BoundedSemaphore, locks.Lock):
+            with self.subTest(primitive=cls.__name__):
+                primitive = cls()
+                await primitive.acquire()
+
+                async def worker():
+                    async with primitive:
+                        self.fail("cancelled waiter entered the context")
+
+                task = asyncio.create_task(worker())
+                await asyncio.sleep(0)
+                successor = asyncio.ensure_future(primitive.acquire())
+                primitive.release()
+                task.cancel()
+                with self.assertRaises(asyncio.CancelledError):
+                    await task
+                try:
+                    self.assertTrue(successor.done())
+                    await successor
+                    primitive.release()
+                finally:
+                    successor.cancel()
+
+    @gen_test
+    async def test_cancel_before_handoff(self):
+        for cls in (locks.Semaphore, locks.BoundedSemaphore, locks.Lock):
+            with self.subTest(primitive=cls.__name__):
+                primitive = cls()
+                await primitive.acquire()
+
+                async def worker():
+                    async with primitive:
+                        self.fail("cancelled waiter entered the context")
+
+                task = asyncio.create_task(worker())
+                await asyncio.sleep(0)
+                task.cancel()
+                with self.assertRaises(asyncio.CancelledError):
+                    await task
+                successor = asyncio.ensure_future(primitive.acquire())
+                try:
+                    self.assertFalse(successor.done())
+                    primitive.release()
+                    self.assertTrue(successor.done())
+                    await successor
+                    primitive.release()
+                finally:
+                    successor.cancel()
+
+
 if __name__ == "__main__":
     unittest.main()