fix lock logic

Author: slampy97

ydb/aio/coordination/lock.py

@@ -13,15 +13,22 @@
 
 
 class CoordinationLock:
-    def __init__(self, client, name: str, node_path: Optional[str] = None):
+    def __init__(
+        self,
+        client,
+        name: str,
+        node_path: Optional[str] = None,
+        count: int = 1,
+        timeout_millis: int = 30000,
+    ):
         self._client = client
         self._driver = client._driver
         self._name = name
         self._node_path = node_path
 
         self._req_id: Optional[int] = None
-        self._count: int = 1
-        self._timeout_millis: int = 30000
+        self._count: int = count
+        self._timeout_millis: int = timeout_millis
         self._next_req_id: int = 1
 
         self._closed: asyncio.Event = asyncio.Event()
@@ -32,6 +39,7 @@ def __init__(self, client, name: str, node_path: Optional[str] = None):
         self._session_ready: asyncio.Event = asyncio.Event()
         self._reconnector = CoordinationReconnector(self)
 
+        self._wait_timeout: float = self._timeout_millis / 1000.0
 
     def next_req_id(self) -> int:
         r = self._next_req_id
@@ -43,22 +51,51 @@ async def send(self, req):
             raise issues.Error("Stream is not started yet")
         await self._stream.send(req)
 
-    async def __aenter__(self):
-        if self.session_id is None:
-            if not self._node_path:
-                raise issues.Error("node_path is not set for CoordinationLock")
-
-            await self._request_queue.put(
-                SessionStart(
-                    path=self._node_path,
-                    session_id=0,
-                    timeout_millis=30000,
-                ).to_proto()
-            )
+    async def _start_session(self):
+        if self.session_id is not None:
+            return
+
+        if not self._node_path:
+            raise issues.Error("node_path is not set for CoordinationLock")
 
-            self._reconnector.start()
+        await self._request_queue.put(
+            SessionStart(
+                path=self._node_path,
+                session_id=0,
+                timeout_millis=self._timeout_millis,
+            ).to_proto()
+        )
+
+        self._reconnector.start()
+        await self._session_ready.wait()
+
+    async def _stop_session(self):
+        self._closed.set()
+        if self._stream:
+            await self._stream.close()
+            self._stream = None
 
-            await self._session_ready.wait()
+        await self._reconnector.stop()
+        self.session_id = None
+        self._node_path = None
+
+    async def _wait_for_acquire_response(self):
+        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"
+            )
+
+    async def __aenter__(self):
+        await self._start_session()
 
         self._req_id = self.next_req_id()
 
@@ -78,29 +115,18 @@ async def __aenter__(self):
         else:
             raise issues.Error(f"Failed to acquire lock: {resp.issues}")
 
-    async def _wait_for_acquire_response(self):
-        try:
-            while True:
-                resp = await asyncio.wait_for(
-                    self._stream._incoming_queue.get(),
-                    timeout=30.0,
-                )
-                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")
-
     async def __aexit__(self, exc_type, exc, tb):
         if self._req_id is not None:
-            req = ReleaseSemaphore(req_id=self._req_id, name=self._name)
-            await self.send(req)
+            try:
+                req = ReleaseSemaphore(req_id=self._req_id, name=self._name)
+                await self.send(req)
+            except issues.Error:
+                pass
 
-        self._closed.set()
-        if self._stream:
-            await self._stream.close()
-            self._stream = None
+        await self._stop_session()
 
-        await self._reconnector.stop()
-        self.session_id = None
-        self._node_path = None
+    async def acquire(self):
+        return await self.__aenter__()
+
+    async def release(self):
+        await self.__aexit__(None, None, None)