From f8ca80f67b5d79d121dec22d90598de5e5dc2965 Mon Sep 17 00:00:00 2001
From: Dylan Baker <dylan.c.baker@intel.com>
Date: Wed, 19 Apr 2023 10:25:11 -0700
Subject: [PATCH] bin/pick: Add a Read/Write lock class

---
 bin/pick/core.py | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/bin/pick/core.py b/bin/pick/core.py
index b0db62bdefa..d4944669f48 100644
--- a/bin/pick/core.py
+++ b/bin/pick/core.py
@@ -21,6 +21,7 @@
 """Core data structures and routines for pick."""
 
 import asyncio
+import contextlib
 import enum
 import json
 import pathlib
@@ -44,7 +45,7 @@ if typing.TYPE_CHECKING:
         resolution: typing.Optional[int]
         main_sha: typing.Optional[str]
         because_sha: typing.Optional[str]
-        notes: typing.Optional[str] = attr.ib(None)
+        notes: typing.Optional[str]
 
 IS_FIX = re.compile(r'^\s*fixes:\s*([a-f0-9]{6,40})', flags=re.MULTILINE | re.IGNORECASE)
 # FIXME: I dislike the duplication in this regex, but I couldn't get it to work otherwise
@@ -62,6 +63,39 @@ git_toplevel = subprocess.check_output(['git', 'rev-parse', '--show-toplevel'],
 pick_status_json = pathlib.Path(git_toplevel) / '.pick_status.json'
 
 
+@attr.s(slots=True, cmp=False)
+class AsyncRWLock:
+
+    """An asynchronous Read/Write lock.
+    
+    This is a very simple read/write lock that prioritizes reads.
+
+    As an implementation detail, this relies on python's global locking to drop
+    the need for a lock to protect the `readers` attribute.
+    """
+
+    readers: int = attr.ib(0, init=False)
+    global_lock: asyncio.Lock = attr.ib(factory=asyncio.Lock, init=False)
+    read_lock: asyncio.Lock = attr.ib(factory=asyncio.Lock, init=False)
+
+    @contextlib.asynccontextmanager
+    async def read(self) -> typing.AsyncIterator[None]:
+        async with self.read_lock:
+            self.readers += 1
+            if self.readers == 1:
+                await self.global_lock.acquire()
+        yield
+        async with self.read_lock:
+            self.readers -= 1
+            if self.readers == 0:
+                self.global_lock.release()
+
+    @contextlib.asynccontextmanager
+    async def write(self) -> typing.AsyncIterator[None]:
+        async with self.global_lock:
+            yield
+
+
 class PickUIException(Exception):
     pass