first working version of sync and asunc client with lock api

Author: slampy97

ydb/aio/coordination/lock.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import Optional
 
-from ydb import issues
+from ydb import issues, StatusCode
 from ydb._grpc.grpcwrapper.ydb_coordination import (
     AcquireSemaphore,
     ReleaseSemaphore,
@@ -130,17 +130,15 @@ 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,
+                )
                 await self.send(req)
             except issues.Error:
                 pass
 
-        if self._reconnector:
-            await self._reconnector.stop(flush=True)
-
-        self._stream = None
         self._req_id = None
-        self._node_path = None
 
     async def acquire(self):
         return await self.__aenter__()
@@ -160,18 +158,31 @@ async def create(self, init_limit, init_data):
         resp = await self._wait_for_response(req_id, kind="create")
         return CreateSemaphoreResult.from_proto(resp)
 
-    async def delete(self):
+    async def delete(self, wait_empty_timeout: float = 5.0, poll_interval: float = 0.05):
         await self._ensure_session()
 
-        req_id = self.next_req_id()
-
-        req = DeleteSemaphore(
-            req_id=req_id,
-            name=self._name,
-        )
+        deadline = asyncio.get_running_loop().time() + wait_empty_timeout
+        while True:
+            try:
+                desc = await self.describe()
+                if desc.count == 0 and not desc.owners:
+                    break
+            except issues.Error as e:
+                if getattr(e, "status", None) == StatusCode.NOT_FOUND:
+                    break
+                raise
+
+            now = asyncio.get_running_loop().time()
+            if now > deadline:
+                raise issues.Error(
+                    f"Timeout waiting for semaphore '{self._name}' to become empty before delete. "
+                    f"count={desc.count}, owners={list(desc.owners)}"
+                )
+            await asyncio.sleep(poll_interval)
 
+        req_id = self.next_req_id()
+        req = DeleteSemaphore(req_id=req_id, name=self._name)
         await self.send(req)
-
         resp = await self._wait_for_response(req_id, kind="delete")
         return resp
 
@@ -206,18 +217,17 @@ async def update(self, new_data):
         return resp
 
     async def close(self, flush: bool = True):
-
         try:
             if self._req_id is not None:
                 req = ReleaseSemaphore(req_id=self._req_id, name=self._name)
-                await self.send(req)
+                if self._stream is not None:
+                    await self.send(req)
         except issues.Error:
             pass
 
-        if self._reconnector:
+        if self._reconnector is not None:
             await self._reconnector.stop(flush)
 
         self._stream = None
         self._req_id = None
         self._node_path = None
-

ydb/aio/coordination/reconnector.py

@@ -6,14 +6,6 @@
 
 
 class CoordinationReconnector:
-    """
-    Простейшая затычка реконнектора:
-    - не делает автопереподключение
-    - сразу бросает ошибку, если stream не стартан
-    - интерфейс start/stop/get_stream/wait_ready полностью рабочий
-    - не зависает ни на каких таймаутах
-    """
-
     def __init__(
         self,
         driver,
@@ -29,18 +21,23 @@ def __init__(
         self._task: Optional[asyncio.Task] = None
         self._stream: Optional[CoordinationStream] = None
 
-        self._ready = asyncio.Event()
+        self._ready: Optional[asyncio.Event] = None
         self._stopped = False
         self._first_error: Optional[Exception] = None
 
     def start(self):
-        """Запуск реконнектора (фактически старт одного stream)"""
         if self._stopped:
             return
+
+        if self._ready is None:
+            self._ready = asyncio.Event()
+
+        self._first_error = None
+
         if self._task is None or self._task.done():
             self._task = asyncio.create_task(self._connection_loop())
 
-    async def stop(self, flush: bool = True):
+    async def stop(self, flush: bool):
         self._stopped = True
 
         if self._task:
@@ -50,29 +47,23 @@ async def stop(self, flush: bool = True):
             self._task = None
 
         if self._stream:
-            if flush:
-                try:
-                    await self._stream.close()
-                except Exception:
-                    pass
-            else:
-                self._stream._closed = True
+            with contextlib.suppress(Exception):
+                await self._stream.close()
             self._stream = None
 
-        # 3. Сбрасываем состояние
-        self._ready.clear()
-        self._first_error = None
+        if self._ready:
+            self._ready.clear()
 
     async def wait_ready(self):
-        """Ждём, пока stream будет готов, или сразу бросаем ошибку"""
         if self._first_error:
             raise self._first_error
+        if not self._ready:
+            raise RuntimeError("Reconnector not started")
         await self._ready.wait()
         if self._first_error:
             raise self._first_error
 
     def get_stream(self) -> CoordinationStream:
-        """Получить готовый stream"""
         if self._stream is None or self._stream.session_id is None:
             raise RuntimeError("Coordination stream is not ready")
         return self._stream
@@ -84,19 +75,27 @@ async def _connection_loop(self):
         try:
             self._stream = CoordinationStream(self._driver)
             await self._stream.start_session(self._node_path, self._timeout_millis)
-            self._ready.set()
+            if self._ready:
+                self._ready.set()
 
             if self._stream._background_tasks:
-                await asyncio.wait(
+                done, pending = await asyncio.wait(
                     self._stream._background_tasks,
                     return_when=asyncio.FIRST_EXCEPTION,
                 )
 
+                for d in done:
+                    if d.cancelled():
+                        continue
+                    exc = d.exception()
+                    if exc:
+                        raise exc
+
         except Exception as exc:
             self._first_error = exc
-            self._ready.clear()
+            if self._ready:
+                self._ready.clear()
             if self._stream:
                 with contextlib.suppress(Exception):
                     await self._stream.close()
             self._stream = None
-

ydb/aio/coordination/stream.py

@@ -1,12 +1,15 @@
 import asyncio
 import contextlib
+import logging
 from typing import Optional, Set
 
 import ydb
 from ydb import issues, _apis
 from ydb._grpc.grpcwrapper.common_utils import IToProto, GrpcWrapperAsyncIO
 from ydb._grpc.grpcwrapper.ydb_coordination import FromServer, Ping, SessionStart
 
+logger = logging.getLogger(__name__)
+
 
 class CoordinationStream:
     def __init__(self, driver: "ydb.aio.Driver"):
@@ -32,6 +35,7 @@ async def start_session(self, path: str, timeout_millis: int):
         self._stream.write(start_msg)
 
         try:
+            # Wait for SessionStart response
             async for resp in self._stream.from_server_grpc:
                 fs = FromServer.from_proto(resp)
                 if fs.session_started:
@@ -44,7 +48,9 @@ async def start_session(self, path: str, timeout_millis: int):
         except Exception as e:
             raise issues.Error(f"Failed to start session: {e}")
 
-        self._background_tasks.add(asyncio.create_task(self._reader_loop()))
+        task = asyncio.get_running_loop().create_task(self._reader_loop())
+        self._background_tasks.add(task)
+        logger.debug("CoordinationStream: started reader task %r", task)
 
     async def _reader_loop(self):
         try:
@@ -54,13 +60,18 @@ async def _reader_loop(self):
 
                 fs = FromServer.from_proto(resp)
                 if fs.opaque:
-                    self._stream.write(Ping(fs.opaque))
+                    try:
+                        self._stream.write(Ping(fs.opaque))
+                    except Exception:
+                        self._set_first_error(RuntimeError("Failed to write Ping"))
                 else:
                     await self._incoming_queue.put(resp)
                     self._state_changed.set()
         except asyncio.CancelledError:
+            logger.debug("CoordinationStream: reader loop cancelled")
             pass
         except Exception as exc:
+            logger.exception("CoordinationStream: reader loop error")
             self._set_first_error(exc)
 
     async def send(self, req: IToProto):
@@ -89,26 +100,25 @@ async def close(self):
             return
         self._closed = True
 
-        tasks = list(self._background_tasks)
-        for task in tasks:
+        logger.debug("CoordinationStream: closing, cancelling %d background tasks", len(self._background_tasks))
+        for task in list(self._background_tasks):
             task.cancel()
 
         with contextlib.suppress(asyncio.CancelledError):
-            await asyncio.gather(*tasks, return_exceptions=True)
+            await asyncio.gather(*self._background_tasks, return_exceptions=True)
 
         self._background_tasks.clear()
 
         if self._stream:
-            close_coro = getattr(self._stream, "close", None)
-            if close_coro and asyncio.iscoroutinefunction(close_coro):
-                await self._stream.close()
-            else:
+            try:
                 self._stream.close()
+            except Exception:
+                logger.exception("CoordinationStream: error closing underlying stream")
+            self._stream = None
 
         self.session_id = None
-        self._started = False
         self._state_changed.set()
-        self._incoming_queue = asyncio.Queue()
+        logger.debug("CoordinationStream: closed")
 
     def _set_first_error(self, exc: Exception):
         if not self._first_error.done():
@@ -118,6 +128,7 @@ def _set_first_error(self, exc: Exception):
     def _get_first_error(self):
         if self._first_error.done():
             return self._first_error.result()
+        return None
 
     def _check_error(self):
         err = self._get_first_error()