rewrite stream with wrapper plus sync client

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -1,9 +1,11 @@
 import asyncio
+import threading
+import time
 
 import pytest
 
 import ydb
-from ydb import aio, StatusCode
+from ydb import aio, StatusCode, logger
 
 from ydb.coordination import (
     NodeConfig,
@@ -155,8 +157,6 @@ 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: DescribeLockResult = await lock1.describe()
             assert resp.status == StatusCode.SUCCESS
@@ -172,14 +172,10 @@ async def second_lock_task():
 
             await asyncio.sleep(0.5)
 
-            assert lock1._stream is not None
-
         await asyncio.wait_for(lock2_acquired.wait(), timeout=5)
         await asyncio.wait_for(t2, timeout=5)
 
         async with client.lock("test_lock", node_path) as lock3:
-            assert lock3._stream is not None
-            assert lock3._stream.session_id is not None
 
             resp3: DescribeLockResult = await lock3.describe()
             assert resp3.status == StatusCode.SUCCESS
@@ -190,3 +186,97 @@ async def second_lock_task():
 
         describe_after_delete: DescribeLockResult = await lock.describe()
         assert describe_after_delete.status == StatusCode.NOT_FOUND
+
+    def test_coordination_lock_full_lifecycle_sync(self, driver_sync):
+        client = CoordinationClient(driver_sync)
+        node_path = "/local/test_lock_full_lifecycle"
+
+        # --- cleanup ---
+        try:
+            client.delete_node(node_path)
+        except ydb.SchemeError:
+            pass
+
+        client.create_node(
+            node_path,
+            NodeConfig(
+                session_grace_period_millis=1000,
+                attach_consistency_mode=ConsistencyMode.STRICT,
+                read_consistency_mode=ConsistencyMode.STRICT,
+                rate_limiter_counters_mode=RateLimiterCountersMode.UNSET,
+                self_check_period_millis=0,
+            ),
+        )
+
+        lock = client.lock("test_lock", node_path)
+
+        # --- create/update/describe ---
+        create_resp: CreateSemaphoreResult = lock.create(init_limit=1, init_data=b"init-data")
+        assert create_resp.status == StatusCode.SUCCESS
+
+        describe_resp: DescribeLockResult = lock.describe()
+        assert describe_resp.status == StatusCode.SUCCESS
+        assert describe_resp.data == b"init-data"
+
+        update_resp = lock.update(new_data=b"updated-data")
+        assert update_resp.status == StatusCode.SUCCESS
+        assert lock.describe().data == b"updated-data"
+
+        # --- threading coordination ---
+        lock2_ready = threading.Event()
+        lock2_acquired = threading.Event()
+        thread_exc = {"err": None}
+
+        def second_lock_task():
+            try:
+                lock2_ready.set()  # сигнал, что поток готов
+                with client.lock("test_lock", node_path):
+                    lock2_acquired.set()  # сигнал, что захватил lock
+                    logger.info("Second thread acquired lock")
+            except Exception as e:
+                logger.exception("second_lock_task failed")
+                thread_exc["err"] = e
+
+        t2 = threading.Thread(target=second_lock_task)
+
+        # --- main thread acquires first lock ---
+        with client.lock("test_lock", node_path) as lock1:
+            resp = lock1.describe()
+            assert resp.status == StatusCode.SUCCESS
+            assert resp.count == 1
+
+            # запускаем второй поток
+            t2.start()
+            started = lock2_ready.wait(timeout=2.0)
+            assert started, "Second thread did not signal readiness to acquire lock"
+
+        # --- main thread released lock, второй поток должен захватить ---
+        acquired = lock2_acquired.wait(timeout=10.0)  # увеличенный таймаут
+        t2.join(timeout=5.0)
+
+        if not acquired:
+            if thread_exc["err"]:
+                raise AssertionError(f"Second thread raised exception: {thread_exc['err']!r}") from thread_exc["err"]
+            else:
+                raise AssertionError(
+                    "Second thread did not acquire the lock in time. Check logs for details."
+                )
+
+        assert not t2.is_alive(), "Second thread did not finish after acquiring lock"
+
+        # --- проверяем, что lock можно снова взять в главном потоке ---
+        with client.lock("test_lock", node_path) as lock3:
+            resp3: DescribeLockResult = lock3.describe()
+            assert resp3.status == StatusCode.SUCCESS
+            assert resp3.count == 1
+
+        # --- cleanup ---
+        delete_resp = lock.delete()
+        assert delete_resp.status == StatusCode.SUCCESS
+        # небольшая пауза для удаления на сервере
+        time.sleep(0.1)
+        describe_after_delete: DescribeLockResult = lock.describe()
+        assert describe_after_delete.status == StatusCode.NOT_FOUND
+
+
+

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -1,21 +1,21 @@
 import typing
 from dataclasses import dataclass
 
-
 if typing.TYPE_CHECKING:
     from ..v4.protos import ydb_coordination_pb2
 else:
     from ..common.protos import ydb_coordination_pb2
 
 from .common_utils import IToProto
 
+# ---------- CRUD для узлов ----------
 
 @dataclass
 class CreateNodeRequest(IToProto):
     path: str
     config: typing.Any
 
-    def to_proto(self) -> ydb_coordination_pb2.CreateNodeRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.CreateNodeRequest":
         cfg_proto = self.config.to_proto() if self.config else None
         return ydb_coordination_pb2.CreateNodeRequest(
             path=self.path,
@@ -28,7 +28,7 @@ class AlterNodeRequest(IToProto):
     path: str
     config: typing.Any
 
-    def to_proto(self) -> ydb_coordination_pb2.AlterNodeRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.AlterNodeRequest":
         cfg_proto = self.config.to_proto() if self.config else None
         return ydb_coordination_pb2.AlterNodeRequest(
             path=self.path,
@@ -40,7 +40,7 @@ def to_proto(self) -> ydb_coordination_pb2.AlterNodeRequest:
 class DescribeNodeRequest(IToProto):
     path: str
 
-    def to_proto(self) -> ydb_coordination_pb2.DescribeNodeRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.DescribeNodeRequest":
         return ydb_coordination_pb2.DescribeNodeRequest(
             path=self.path,
         )
@@ -50,11 +50,12 @@ def to_proto(self) -> ydb_coordination_pb2.DescribeNodeRequest:
 class DropNodeRequest(IToProto):
     path: str
 
-    def to_proto(self) -> ydb_coordination_pb2.DropNodeRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.DropNodeRequest":
         return ydb_coordination_pb2.DropNodeRequest(
             path=self.path,
         )
 
+# ---------- Сессии и семафоры ----------
 
 @dataclass
 class SessionStart(IToProto):
@@ -65,7 +66,7 @@ class SessionStart(IToProto):
     seq_no: int = 0
     protection_key: bytes = b""
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
             session_start=ydb_coordination_pb2.SessionRequest.SessionStart(
                 path=self.path,
@@ -80,15 +81,17 @@ def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
 
 @dataclass
 class SessionStop(IToProto):
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
-        return ydb_coordination_pb2.SessionRequest(session_stop=ydb_coordination_pb2.SessionRequest.SessionStop())
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
+        return ydb_coordination_pb2.SessionRequest(
+            session_stop=ydb_coordination_pb2.SessionRequest.SessionStop()
+        )
 
 
 @dataclass
 class Ping(IToProto):
     opaque: int = 0
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
             ping=ydb_coordination_pb2.SessionRequest.PingPong(opaque=self.opaque)
         )
@@ -101,42 +104,14 @@ class CreateSemaphore(IToProto):
     limit: int
     data: bytes = b""
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
             create_semaphore=ydb_coordination_pb2.SessionRequest.CreateSemaphore(
                 req_id=self.req_id, name=self.name, limit=self.limit, data=self.data
             )
         )
 
 
-@dataclass
-class UpdateSemaphore(IToProto):
-    name: str
-    req_id: int
-    data: bytes
-
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
-        return ydb_coordination_pb2.SessionRequest(
-            update_semaphore=ydb_coordination_pb2.SessionRequest.UpdateSemaphore(
-                req_id=self.req_id, name=self.name, data=self.data
-            )
-        )
-
-
-@dataclass
-class DeleteSemaphore(IToProto):
-    name: str
-    req_id: int
-    force: bool = False
-
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
-        return ydb_coordination_pb2.SessionRequest(
-            delete_semaphore=ydb_coordination_pb2.SessionRequest.DeleteSemaphore(
-                req_id=self.req_id, name=self.name, force=self.force
-            )
-        )
-
-
 @dataclass
 class AcquireSemaphore(IToProto):
     name: str
@@ -146,7 +121,7 @@ class AcquireSemaphore(IToProto):
     data: bytes = b""
     ephemeral: bool = False
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
             acquire_semaphore=ydb_coordination_pb2.SessionRequest.AcquireSemaphore(
                 req_id=self.req_id,
@@ -164,9 +139,11 @@ class ReleaseSemaphore(IToProto):
     name: str
     req_id: int
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
-            release_semaphore=ydb_coordination_pb2.SessionRequest.ReleaseSemaphore(req_id=self.req_id, name=self.name)
+            release_semaphore=ydb_coordination_pb2.SessionRequest.ReleaseSemaphore(
+                req_id=self.req_id, name=self.name
+            )
         )
 
 
@@ -179,7 +156,7 @@ class DescribeSemaphore(IToProto):
     watch_data: bool
     watch_owners: bool
 
-    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
         return ydb_coordination_pb2.SessionRequest(
             describe_semaphore=ydb_coordination_pb2.SessionRequest.DescribeSemaphore(
                 include_owners=self.include_owners,
@@ -191,20 +168,47 @@ def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
             )
         )
 
+@dataclass
+class UpdateSemaphore(IToProto):
+    name: str
+    req_id: int
+    data: bytes
+
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
+        return ydb_coordination_pb2.SessionRequest(
+            update_semaphore=ydb_coordination_pb2.SessionRequest.UpdateSemaphore(
+                req_id=self.req_id, name=self.name, data=self.data
+            )
+        )
+
+
+@dataclass
+class DeleteSemaphore(IToProto):
+    name: str
+    req_id: int
+    force: bool = False
+
+    def to_proto(self) -> "ydb_coordination_pb2.SessionRequest":
+        return ydb_coordination_pb2.SessionRequest(
+            delete_semaphore=ydb_coordination_pb2.SessionRequest.DeleteSemaphore(
+                req_id=self.req_id, name=self.name, force=self.force
+            )
+        )
+
 
 @dataclass
 class FromServer:
-    raw: ydb_coordination_pb2.SessionResponse
+    raw: "ydb_coordination_pb2.SessionResponse"
 
     @staticmethod
-    def from_proto(resp: ydb_coordination_pb2.SessionResponse) -> "FromServer":
+    def from_proto(resp: "ydb_coordination_pb2.SessionResponse") -> "FromServer":
         return FromServer(raw=resp)
 
     def __getattr__(self, name: str):
         return getattr(self.raw, name)
 
     @property
-    def session_started(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.SessionStarted]:
+    def session_started(self) -> typing.Optional["ydb_coordination_pb2.SessionResponse.SessionStarted"]:
         s = self.raw.session_started
         return s if s.session_id else None
 
@@ -215,23 +219,9 @@ def opaque(self) -> typing.Optional[int]:
         return None
 
     @property
-    def acquire_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.AcquireSemaphoreResult]:
+    def acquire_semaphore_result(self):
         return self.raw.acquire_semaphore_result if self.raw.HasField("acquire_semaphore_result") else None
 
     @property
-    def create_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult]:
+    def create_semaphore_result(self):
         return self.raw.create_semaphore_result if self.raw.HasField("create_semaphore_result") else None
-
-    @property
-    def delete_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DeleteSemaphoreResult]:
-        return self.raw.delete_semaphore_result if self.raw.HasField("delete_semaphore_result") else None
-
-    @property
-    def update_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.UpdateSemaphoreResult]:
-        return self.raw.update_semaphore_result if self.raw.HasField("update_semaphore_result") else None
-
-    @property
-    def describe_semaphore_result(
-        self,
-    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DescribeSemaphoreResult]:
-        return self.raw.describe_semaphore_result if self.raw.HasField("describe_semaphore_result") else None