make readwrite more safe

Author: Jineshbansal

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -35,55 +35,69 @@
 
 
 class ReadWriteLock:
+    """
+    A read-write lock that allows multiple concurrent readers
+    or one exclusive writer, implemented using Trio primitives.
+    """
+
     def __init__(self) -> None:
         self._readers = 0
-        self._readers_lock = trio.Lock()  # Protects readers count
-        self._writer_lock = trio.Semaphore(1)  # Ensures mutual exclusion for writers
+        self._readers_lock = trio.Lock()  # Protects access to _readers count
+        self._writer_lock = trio.Semaphore(1)  # Allows only one writer at a time
 
     async def acquire_read(self) -> None:
+        """Acquire a read lock. Multiple readers can hold it simultaneously."""
         try:
             async with self._readers_lock:
-                self._readers += 1
-                if self._readers == 1:
+                if self._readers == 0:
                     await self._writer_lock.acquire()
+                self._readers += 1
         except trio.Cancelled:
-            async with self._readers_lock:
-                if self._readers > 0:
-                    self._readers -= 1
-                    if self._readers == 0:
-                        self._writer_lock.release()
             raise
 
     async def release_read(self) -> None:
+        """Release a read lock."""
         async with self._readers_lock:
-            self._readers -= 1
-            if self._readers == 0:
+            if self._readers == 1:
                 self._writer_lock.release()
+            self._readers -= 1
 
     async def acquire_write(self) -> None:
+        """Acquire an exclusive write lock."""
         try:
             await self._writer_lock.acquire()
         except trio.Cancelled:
             raise
 
     def release_write(self) -> None:
+        """Release the exclusive write lock."""
         self._writer_lock.release()
 
     @asynccontextmanager
     async def read_lock(self) -> AsyncGenerator[None, None]:
-        await self.acquire_read()
+        """Context manager for acquiring and releasing a read lock safely."""
+        acquire = False
         try:
+            await self.acquire_read()
+            acquire = True
             yield
         finally:
-            await self.release_read()
+            if acquire:
+                with trio.CancelScope() as scope:
+                    scope.shield = True
+                    await self.release_read()
 
     @asynccontextmanager
     async def write_lock(self) -> AsyncGenerator[None, None]:
-        await self.acquire_write()
+        """Context manager for acquiring and releasing a write lock safely."""
+        acquire = False
         try:
+            await self.acquire_write()
+            acquire = True
             yield
         finally:
-            self.release_write()
+            if acquire:
+                self.release_write()
 
 
 class MplexStream(IMuxedStream):