Simplify release of write lock by delaying cancellation until the loc…

…k is acquired

synapse/util/async_helpers.py

@@ -562,31 +562,22 @@ async def _ctx_manager() -> AsyncIterator[None]:
             to_wait_on_defer = defer.gatherResults(to_wait_on)
             try:
                 # Wait for all current readers and the latest writer to finish.
-                # May raise a `CancelledError` if the `Deferred` wrapping us is
-                # cancelled. The `Deferred`s we are waiting on must not be cancelled,
-                # since we do not own them.
-                await make_deferred_yieldable(stop_cancellation(to_wait_on_defer))
+                # May raise a `CancelledError` immediately after the wait if the
+                # `Deferred` wrapping us is cancelled. We must only release the lock
+                # once we have acquired it, hence the delay.
+                await make_deferred_yieldable(
+                    delay_cancellation(to_wait_on_defer, all=True)
+                )
                 yield
             finally:
-
-                def release() -> None:
-                    with PreserveLoggingContext():
-                        new_defer.callback(None)
-                    # `self.key_to_current_writer[key]` may be missing if there was another
-                    # writer waiting for us and it completed entirely within the
-                    # `new_defer.callback()` call above.
-                    if self.key_to_current_writer.get(key) == new_defer:
-                        self.key_to_current_writer.pop(key)
-
-                if to_wait_on_defer.called:
-                    release()
-                else:
-                    # We don't have the lock yet, probably because we were cancelled
-                    # while waiting for it. We can't call `release()` yet, since
-                    # `new_defer` must only resolve once all previous readers and
-                    # writers have finished.
-                    # NB: `release()` won't have a logcontext in this path.
-                    to_wait_on_defer.addBoth(lambda _: release())
+                # Release the lock.
+                with PreserveLoggingContext():
+                    new_defer.callback(None)
+                # `self.key_to_current_writer[key]` may be missing if there was another
+                # writer waiting for us and it completed entirely within the
+                # `new_defer.callback()` call above.
+                if self.key_to_current_writer.get(key) == new_defer:
+                    self.key_to_current_writer.pop(key)
 
         return _ctx_manager()
 

tests/util/test_rwlock.py

@@ -353,13 +353,12 @@ def test_cancellation_while_waiting_for_write_lock(self):
         writer3_d, _ = self._start_nonblocking_writer(rwlock, key, "write 3 completed")
         self.assertFalse(writer3_d.called)
 
-        # 5. The second writer is cancelled.
+        # 5. The second writer is cancelled, but continues waiting for the lock.
         #    The reader, first writer and third writer should not be cancelled.
         #    The first writer should still be waiting on the reader.
-        #    The third writer should still be waiting, even though the second writer has
-        #    been cancelled.
+        #    The third writer should still be waiting on the second writer.
         writer2_d.cancel()
-        self.failureResultOf(writer2_d, CancelledError)
+        self.assertNoResult(writer2_d)
         self.assertFalse(reader_d.called, "Reader was unexpectedly cancelled")
         self.assertFalse(writer1_d.called, "First writer was unexpectedly cancelled")
         self.assertFalse(
@@ -370,9 +369,10 @@ def test_cancellation_while_waiting_for_write_lock(self):
 
         # 6. Unblock the reader, which should complete.
         #    The first writer should be given the lock and block.
-        #    The third writer should still be waiting.
+        #    The third writer should still be waiting on the second writer.
         unblock_reader.callback(None)
         self.assertEqual("read completed", self.successResultOf(reader_d))
+        self.assertNoResult(writer2_d)
         self.assertFalse(
             writer3_d.called,
             "Third writer was unexpectedly given the lock before the first writer "
@@ -383,7 +383,11 @@ def test_cancellation_while_waiting_for_write_lock(self):
         unblock_writer1.callback(None)
         self.assertEqual("write 1 completed", self.successResultOf(writer1_d))
 
-        # 8. The third writer should take the lock and complete.
+        # 8. The second writer should take the lock and release it immediately, since it
+        #    has been cancelled.
+        self.failureResultOf(writer2_d, CancelledError)
+
+        # 9. The third writer should take the lock and complete.
         self.assertTrue(
             writer3_d.called, "Third writer is stuck waiting for a cancelled writer"
         )