crud lock object + acquire and release.

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -117,6 +117,37 @@ async def test_coordination_lock_full_lifecycle(self, aio_connection):
             ),
         )
 
+        lock = client.lock("test_lock", node_path)
+
+        create_resp = await lock.create(init_limit=1, init_data=b"init-data")
+        assert create_resp.status == StatusCode.SUCCESS
+
+        describe_resp = await lock.describe()
+        assert describe_resp.status == StatusCode.SUCCESS
+
+        sem = describe_resp.semaphore_description
+        assert sem.name == "test_lock"
+        assert sem.data == b"init-data"
+        assert sem.count == 0
+        assert sem.ephemeral is False
+        assert list(sem.owners) == []
+        assert list(sem.waiters) == []
+
+        update_resp = await lock.update(new_data=b"updated-data")
+        assert update_resp.status == StatusCode.SUCCESS
+
+        describe_resp2 = await lock.describe()
+        assert describe_resp2.status == StatusCode.SUCCESS
+
+        sem2 = describe_resp2.semaphore_description
+        assert sem2.name == "test_lock"
+        assert sem2.data == b"updated-data"
+        assert sem2.count == 0
+        assert sem2.ephemeral is False
+        assert list(sem2.owners) == []
+        assert list(sem2.waiters) == []
+
+
         lock2_started = asyncio.Event()
         lock2_acquired = asyncio.Event()
 
@@ -127,16 +158,23 @@ async def second_lock_task():
                 await asyncio.sleep(0.5)
 
         async with client.lock("test_lock", node_path) as lock1:
-
             assert lock1._stream is not None
             assert lock1._stream.session_id is not None
+
             resp = await lock1.describe()
             assert resp.status == StatusCode.SUCCESS
 
+            sem_under_lock = resp.semaphore_description
+            assert sem_under_lock.name == "test_lock"
+            assert sem_under_lock.data == b"updated-data"
+            assert sem_under_lock.count == 1
+            assert sem_under_lock.ephemeral is False
+            assert len(list(sem_under_lock.owners)) == 1
+            assert list(sem_under_lock.waiters) == []
+
             t2 = asyncio.create_task(second_lock_task())
             await lock2_started.wait()
 
-
             await asyncio.sleep(0.5)
 
             assert lock1._stream is not None
@@ -148,3 +186,15 @@ async def second_lock_task():
             assert lock3._stream is not None
             assert lock3._stream.session_id is not None
 
+            resp3 = await lock3.describe()
+            assert resp3.status == StatusCode.SUCCESS
+            sem3 = resp3.semaphore_description
+            assert sem3.count == 1
+
+
+        delete_resp = await lock.delete()
+        assert delete_resp.status == StatusCode.SUCCESS
+
+        describe_after_delete = await lock.describe()
+        assert describe_after_delete.status == StatusCode.NOT_FOUND
+

ydb/aio/coordination/lock.py

@@ -7,6 +7,8 @@
     ReleaseSemaphore,
     UpdateSemaphore,
     DescribeSemaphore,
+    CreateSemaphore,
+    DeleteSemaphore,
     FromServer,
 )
 from ydb.aio.coordination.stream import CoordinationStream
@@ -44,6 +46,7 @@ def __init__(
 
         self._wait_timeout: float = self._timeout_millis / 1000.0
 
+
     def next_req_id(self) -> int:
         r = self._next_req_id
         self._next_req_id += 1
@@ -66,52 +69,62 @@ async def _ensure_session(self):
 
         self._stream = self._reconnector.get_stream()
 
-    async def _wait_for_acquire_response(self):
+    async def _wait_for_response(self, req_id: int, *, kind: str):
         try:
             while True:
                 resp = await asyncio.wait_for(
                     self._stream._incoming_queue.get(),
                     timeout=self._wait_timeout,
                 )
-                acquire_resp = FromServer.from_proto(resp).acquire_semaphore_result
-                if acquire_resp and acquire_resp.req_id == self._req_id:
-                    return acquire_resp
-        except asyncio.TimeoutError:
-            raise issues.Error(
-                f"Timeout waiting for lock {self._name} acquisition"
-            )
+                fs = FromServer.from_proto(resp)
+
+                if kind == "acquire":
+                    r = fs.acquire_semaphore_result
+                elif kind == "describe":
+                    r = fs.describe_semaphore_result
+                elif kind == "create":
+                    r = fs.create_semaphore_result
+                elif kind == "update":
+                    r = fs.update_semaphore_result
+                elif kind == "delete":
+                    r = fs.delete_semaphore_result
+                else:
+                    r = None
+
+                if r and r.req_id == req_id:
+                    return r
 
-    async def _wait_for_describe_response(self, req_id: int):
-        try:
-            while True:
-                resp = await asyncio.wait_for(
-                    self._stream._incoming_queue.get(),
-                    timeout=self._wait_timeout,
-                )
-                describe_resp = FromServer.from_proto(resp).describe_semaphore_result
-                if describe_resp and describe_resp.req_id == req_id:
-                    return describe_resp
         except asyncio.TimeoutError:
+            action = {
+                "acquire": "acquisition",
+                "describe": "describe",
+                "update": "update",
+                "delete": "delete",
+                "create": "create"
+            }.get(kind, "operation")
+
             raise issues.Error(
-                f"Timeout waiting for lock {self._name} describe"
+                f"Timeout waiting for lock {self._name} {action}"
             )
 
+
     async def __aenter__(self):
         await self._ensure_session()
 
-        self._req_id = self.next_req_id()
+        req_id = self.next_req_id()
+        self._req_id = req_id
 
         req = AcquireSemaphore(
-            req_id=self._req_id,
+            req_id=req_id,
             name=self._name,
             count=self._count,
-            ephemeral=True,
+            ephemeral=False,
             timeout_millis=self._timeout_millis,
         ).to_proto()
 
         await self.send(req)
 
-        resp = await self._wait_for_acquire_response()
+        resp = await self._wait_for_response(req_id, kind="acquire")
         if resp.acquired:
             return self
         else:
@@ -120,21 +133,57 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc, tb):
         if self._req_id is not None:
             try:
-                req = ReleaseSemaphore(req_id=self._req_id, name=self._name)
+                req = ReleaseSemaphore(
+                    req_id=self._req_id,
+                    name=self._name,
+                ).to_proto()
                 await self.send(req)
             except issues.Error:
                 pass
 
         await self._reconnector.stop()
         self._stream = None
         self._node_path = None
+        self._req_id = None
 
     async def acquire(self):
         return await self.__aenter__()
 
     async def release(self):
         await self.__aexit__(None, None, None)
 
+    async def create(self, init_limit, init_data):
+        await self._ensure_session()
+
+        req_id = self.next_req_id()
+
+        req = CreateSemaphore(
+            req_id=req_id,
+            name=self._name,
+            limit=init_limit,
+            data=init_data
+        ).to_proto()
+
+        await self.send(req)
+
+        resp = await self._wait_for_response(req_id, kind="create")
+        return resp
+
+    async def delete(self):
+        await self._ensure_session()
+
+        req_id = self.next_req_id()
+
+        req = DeleteSemaphore(
+            req_id=req_id,
+            name=self._name,
+        ).to_proto()
+
+        await self.send(req)
+
+        resp = await self._wait_for_response(req_id, kind="delete")
+        return resp
+
 
     async def describe(self):
         await self._ensure_session()
@@ -152,8 +201,20 @@ async def describe(self):
 
         await self.send(req)
 
-        resp = await self._wait_for_describe_response(req_id)
+        resp = await self._wait_for_response(req_id, kind="describe")
         return resp
 
-    async def update(self, *, limit: Optional[int] = None, settings=None):
-        pass
+    async def update(self, new_data):
+        await self._ensure_session()
+
+        req_id = self.next_req_id()
+        req = UpdateSemaphore(
+            req_id=req_id,
+            name=self._name,
+            data=new_data
+        ).to_proto()
+
+        await self.send(req)
+
+        resp = await self._wait_for_response(req_id, kind="update")
+        return resp