refactor logic lock + reconnecor + stream - > lock should be resosible only for lock logic (release and so on) , stream about session work -> making stream_stream canal and so on, reconnector -> make this stream stable

Author: slampy97

ydb/aio/coordination/lock.py

@@ -6,7 +6,6 @@
     AcquireSemaphore,
     ReleaseSemaphore,
     FromServer,
-    SessionStart,
 )
 from ydb.aio.coordination.stream import CoordinationStream
 from ydb.aio.coordination.reconnector import CoordinationReconnector
@@ -31,13 +30,15 @@ def __init__(
         self._timeout_millis: int = timeout_millis
         self._next_req_id: int = 1
 
-        self._closed: asyncio.Event = asyncio.Event()
         self._request_queue: asyncio.Queue = asyncio.Queue()
         self._stream: Optional[CoordinationStream] = None
-        self._reader_task: Optional[asyncio.Task] = None
-        self.session_id: Optional[int] = None
-        self._session_ready: asyncio.Event = asyncio.Event()
-        self._reconnector = CoordinationReconnector(self)
+
+        self._reconnector = CoordinationReconnector(
+            driver=self._driver,
+            request_queue=self._request_queue,
+            node_path=self._node_path,
+            timeout_millis=self._timeout_millis,
+        )
 
         self._wait_timeout: float = self._timeout_millis / 1000.0
 
@@ -51,33 +52,17 @@ async def send(self, req):
             raise issues.Error("Stream is not started yet")
         await self._stream.send(req)
 
-    async def _start_session(self):
-        if self.session_id is not None:
+    async def _ensure_session(self):
+        if self._stream is not None and self._stream.session_id is not None:
             return
 
         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=self._timeout_millis,
-            ).to_proto()
-        )
-
         self._reconnector.start()
-        await self._session_ready.wait()
+        await self._reconnector.wait_ready()
 
-    async def _stop_session(self):
-        self._closed.set()
-        if self._stream:
-            await self._stream.close()
-            self._stream = None
-
-        await self._reconnector.stop()
-        self.session_id = None
-        self._node_path = None
+        self._stream = self._reconnector.get_stream()
 
     async def _wait_for_acquire_response(self):
         try:
@@ -95,7 +80,7 @@ async def _wait_for_acquire_response(self):
             )
 
     async def __aenter__(self):
-        await self._start_session()
+        await self._ensure_session()
 
         self._req_id = self.next_req_id()
 
@@ -123,7 +108,9 @@ async def __aexit__(self, exc_type, exc, tb):
             except issues.Error:
                 pass
 
-        await self._stop_session()
+        await self._reconnector.stop()
+        self._stream = None
+        self._node_path = None
 
     async def acquire(self):
         return await self.__aenter__()

ydb/aio/coordination/reconnector.py

@@ -6,51 +6,99 @@
 
 
 class CoordinationReconnector:
-    def __init__(self, client):
-        self._client = client
+    def __init__(
+        self,
+        driver,
+        request_queue: asyncio.Queue,
+        node_path: str,
+        timeout_millis: int,
+    ):
+        self._driver = driver
+        self._request_queue = request_queue
+        self._node_path = node_path
+        self._timeout_millis = timeout_millis
+
         self._task: Optional[asyncio.Task] = None
+        self._stream: Optional[CoordinationStream] = None
+
+        self._ready = asyncio.Event()
+        self._stopped = False
+
         self._first_error: asyncio.Future = asyncio.get_running_loop().create_future()
         self._state_changed = asyncio.Event()
 
     def start(self):
+        if self._stopped:
+            return
         if self._task is None or self._task.done():
             self._task = asyncio.create_task(self._connection_loop())
 
     async def stop(self):
+        self._stopped = True
+
         if self._task:
             self._task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await self._task
             self._task = None
 
-        if self._client._stream:
-            await self._client._stream.close()
-            self._client._stream = None
+        if self._stream:
+            await self._stream.close()
+            self._stream = None
+
+        self._ready.clear()
+
+    async def wait_ready(self):
+        await self._ready.wait()
+
+    def get_stream(self) -> CoordinationStream:
+        if self._stream is None or self._stream.session_id is None:
+            raise RuntimeError("Coordination stream is not ready")
+        return self._stream
 
     async def _connection_loop(self):
         attempt = 0
         backoff = 0.1
-        while not self._client._closed.is_set():
+
+        while not self._stopped:
             try:
-                self._client._stream = CoordinationStream(
-                    self._client._driver, self._client._request_queue
+                stream = CoordinationStream(
+                    self._driver,
+                    self._request_queue,
+                )
+
+                await stream.start_session(
+                    self._node_path,
+                    self._timeout_millis,
+                )
+
+                self._stream = stream
+                self._ready.set()
+
+                await asyncio.wait(
+                    stream._background_tasks,
+                    return_when=asyncio.FIRST_EXCEPTION,
                 )
-                await self._client._stream.start()
-                self._client.session_id = self._client._stream.session_id
-                self._client._session_ready.set()
 
-                await asyncio.wait(self._client._stream._background_tasks, return_when=asyncio.FIRST_EXCEPTION)
             except asyncio.CancelledError:
-                raise
+                break
+
             except Exception as exc:
-                self._client.session_id = None
-                self._client._session_ready.clear()
+                self._ready.clear()
+                self._stream = None
+
                 if not self._first_error.done():
                     self._first_error.set_result(exc)
+                    self._state_changed.set()
+
+                if self._stopped:
+                    break
+
                 await asyncio.sleep(backoff)
                 attempt += 1
                 backoff = min(backoff * 2, 3.0)
+
             finally:
-                if self._client._stream:
-                    await self._client._stream.close()
-                    self._client._stream = None
+                if self._stream:
+                    await self._stream.close()
+                    self._stream = None

ydb/aio/coordination/stream.py

@@ -3,7 +3,7 @@
 from typing import Set, Optional
 
 from ydb import issues, _apis
-from ydb._grpc.grpcwrapper.ydb_coordination import FromServer, Ping
+from ydb._grpc.grpcwrapper.ydb_coordination import FromServer, Ping, SessionStart
 
 
 class CoordinationStream:
@@ -19,7 +19,21 @@ def __init__(self, driver: "ydb.aio.Driver", request_queue: asyncio.Queue):
         self.session_id: Optional[int] = None
         self._started: bool = False
 
-    async def start(self):
+    async def start_session(self, path: str, timeout_millis: int):
+        if self._started:
+            raise issues.Error("CoordinationStream already started")
+
+        await self.send(
+            SessionStart(
+                path=path,
+                session_id=0,
+                timeout_millis=timeout_millis,
+            ).to_proto()
+        )
+
+        await self._start_internal()
+
+    async def _start_internal(self):
         if self._started:
             raise issues.Error("CoordinationStream already started")
         self._started = True
@@ -30,13 +44,16 @@ async def request_gen():
                 yield req
 
         self._stream = await self._driver(
-            request_gen(), _apis.CoordinationService.Stub, _apis.CoordinationService.Session
+            request_gen(),
+            _apis.CoordinationService.Stub,
+            _apis.CoordinationService.Session,
         )
 
         try:
             async for resp in self._stream:
-                if FromServer.from_proto(resp).session_started:
-                    self.session_id = FromServer.from_proto(resp).session_started
+                fs = FromServer.from_proto(resp)
+                if fs.session_started:
+                    self.session_id = fs.session_started
                     break
         except Exception as exc:
             self._set_first_error(exc)
@@ -84,6 +101,8 @@ async def close(self):
                 self._stream.close()
             except Exception:
                 pass
+
+        self.session_id = None
         self._state_changed.set()
 
     def _set_first_error(self, exc: Exception):
@@ -98,4 +117,4 @@ def _get_first_error(self):
     def _check_error(self):
         err = self._get_first_error()
         if err:
-            raise err
+            raise err