Fix race condition (#236)

* Fix race condition.

* Black

* Quotes

* Fix task ordering in Python 3.6.

aiorwlock/__init__.py

@@ -85,11 +85,11 @@ async def acquire_read(self) -> bool:
         self._read_waiters.append(fut)
         try:
             await fut
-            self._r_state += 1
             self._owning.append((me, self._RL))
             return True
 
         except asyncio.CancelledError:
+            self._r_state -= 1
             self._wake_up()
             raise
 
@@ -120,11 +120,11 @@ async def acquire_write(self) -> bool:
         self._write_waiters.append(fut)
         try:
             await fut
-            self._w_state += 1
             self._owning.append((me, self._WL))
             return True
 
         except asyncio.CancelledError:
+            self._w_state -= 1
             self._wake_up()
             raise
 
@@ -157,22 +157,18 @@ def _wake_up(self) -> None:
         # waiters.
         if self._r_state == 0 and self._w_state == 0:
             if self._write_waiters:
-                self._wake_up_first(self._write_waiters)
-            elif self._read_waiters:
-                self._wake_up_all(self._read_waiters)
-
-    def _wake_up_first(self, waiters: 'Deque[Future[None]]') -> None:
-        # Wake up the first waiter who isn't cancelled.
-        for fut in waiters:
-            if not fut.done():
-                fut.set_result(None)
-                break
-
-    def _wake_up_all(self, waiters: 'Deque[Future[None]]') -> None:
-        # Wake up all not cancelled waiters.
-        for fut in waiters:
-            if not fut.done():
-                fut.set_result(None)
+                # Wake up the first waiter which isn't cancelled.
+                for fut in self._write_waiters:
+                    if not fut.done():
+                        fut.set_result(None)
+                        self._w_state += 1
+                        return
+
+            # Wake up all not cancelled waiters.
+            for fut in self._read_waiters:
+                if not fut.done():
+                    fut.set_result(None)
+                    self._r_state += 1
 
 
 class _ContextManagerMixin:

tests/test_corner_cases.py

@@ -150,3 +150,28 @@ async def coro(lock):
         pass
 
     assert not rl.locked
+
+
+@pytest.mark.asyncio
+async def test_race_multiple_writers(loop):
+    seq = []
+
+    async def write_wait(lock):
+        async with lock.reader:
+            await asyncio.sleep(0.1)
+            seq.append('READ')
+        async with lock.writer:
+            seq.append('START1')
+            await asyncio.sleep(0.1)
+            seq.append('FIN1')
+
+    async def write(lock):
+        await asyncio.sleep(0)  # PY36 seems to run tasks in the wrong order.
+        async with lock.writer:
+            seq.append('START2')
+            await asyncio.sleep(0.1)
+            seq.append('FIN2')
+
+    lock = RWLock(fast=True)
+    await asyncio.gather(write_wait(lock), write(lock))
+    assert seq == ['READ', 'START2', 'FIN2', 'START1', 'FIN1']