Merge branch 'coordination-service-impementation' into coordination-lock

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -1,7 +1,9 @@
+import asyncio
+
 import pytest
 
 import ydb
-from ydb import aio
+from ydb import aio, StatusCode
 
 from ydb.coordination import (
     NodeConfig,
@@ -93,3 +95,106 @@ async def test_coordination_node_lifecycle_async(self, aio_connection):
 
         with pytest.raises(ydb.SchemeError):
             await client.describe_node(node_path)
+
+    async def test_coordination_lock_full_lifecycle(self, aio_connection):
+        client = aio.CoordinationClient(aio_connection)
+
+        node_path = "/local/test_lock_full_lifecycle"
+
+        try:
+            await client.delete_node(node_path)
+        except ydb.SchemeError:
+            pass
+
+        await 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_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()
+
+        async def second_lock_task():
+            lock2_started.set()
+            async with client.lock("test_lock", node_path):
+                lock2_acquired.set()
+                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
+
+        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 = 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/_apis.py

@@ -143,9 +143,9 @@ class QueryService(object):
 
 class CoordinationService(object):
     Stub = ydb_coordination_v1_pb2_grpc.CoordinationServiceStub
-
-    Session = "Session"
     CreateNode = "CreateNode"
     AlterNode = "AlterNode"
     DropNode = "DropNode"
     DescribeNode = "DescribeNode"
+    SessionRequest = "SessionRequest"
+    Session = "Session"

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -55,3 +55,201 @@ def to_proto(self) -> ydb_coordination_pb2.DropNodeRequest:
         return ydb_coordination_pb2.DropNodeRequest(
             path=self.path,
         )
+
+
+@dataclass
+class SessionStart(IToProto):
+    path: str
+    timeout_millis: int
+    description: str = ""
+    session_id: int = 0
+    seq_no: int = 0
+    protection_key: bytes = b""
+
+    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+        return ydb_coordination_pb2.SessionRequest(session_start=ydb_coordination_pb2.SessionRequest.SessionStart(
+            path=self.path,
+            session_id=self.session_id,
+            timeout_millis=self.timeout_millis,
+            description=self.description,
+            seq_no=self.seq_no,
+            protection_key=self.protection_key,
+        ))
+
+
+@dataclass
+class SessionStop(IToProto):
+    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:
+        return ydb_coordination_pb2.SessionRequest(
+            ping=ydb_coordination_pb2.SessionRequest.PingPong(opaque=self.opaque))
+
+
+@dataclass
+class CreateSemaphore(IToProto):
+    name: str
+    req_id: int
+    limit: int
+    data: bytes = b""
+
+    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
+    req_id: int
+    count: int = 1
+    timeout_millis: int = 0
+    data: bytes = b""
+    ephemeral: bool = False
+
+    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,
+                name=self.name,
+                timeout_millis=self.timeout_millis,
+                count=self.count,
+                data=self.data,
+                ephemeral=self.ephemeral,
+            ))
+
+
+@dataclass
+class ReleaseSemaphore(IToProto):
+    name: str
+    req_id: int
+
+    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
+            ))
+
+@dataclass
+class DescribeSemaphore(IToProto):
+    include_owners: bool
+    include_waiters: bool
+    name: str
+    req_id: int
+    watch_data: bool
+    watch_owners: bool
+
+    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,
+                include_waiters=self.include_waiters,
+                name=self.name,
+                req_id=self.req_id,
+                watch_data=self.watch_data,
+                watch_owners=self.watch_owners
+            )
+        )
+
+
+@dataclass
+class FromServer:
+    raw: ydb_coordination_pb2.SessionResponse
+
+    @staticmethod
+    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 status(self) -> typing.Optional[int]:
+        return getattr(self.raw, "status", None)
+
+    @property
+    def session_started(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.SessionStarted]:
+        x = getattr(self.raw, "session_started", None)
+        return x if getattr(x, "session_id", 0) else None
+
+    @property
+    def session_stopped(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.SessionStopped]:
+        return getattr(self.raw, "session_stopped", None) or None
+
+    @property
+    def failure(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.Failure]:
+        return getattr(self.raw, "failure", None) or None
+
+    @property
+    def pong(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.PingPong]:
+        return getattr(self.raw, "pong", None) or None
+
+    @property
+    def ping(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.PingPong]:
+        return getattr(self.raw, "ping", None) or None
+
+    @property
+    def opaque(self) -> typing.Optional[int]:
+        if getattr(self.raw, "ping") is not None:
+            return getattr(getattr(self.raw, "ping"), "opaque", None)
+
+    @property
+    def acquire_semaphore_result(
+        self,
+    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.AcquireSemaphoreResult]:
+        return getattr(self.raw, "acquire_semaphore_result", None) or None
+
+    @property
+    def create_semaphore_result(
+        self,
+    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult]:
+        return getattr(self.raw, "create_semaphore_result", None) or None
+
+    @property
+    def delete_semaphore_result(
+        self,
+    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DeleteSemaphoreResult]:
+        return getattr(self.raw, "delete_semaphore_result", None) or None
+
+    @property
+    def update_semaphore_result(
+        self,
+    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.UpdateSemaphoreResult]:
+        return getattr(self.raw, "update_semaphore_result", None) or None
+
+    @property
+    def describe_semaphore_result(
+        self,
+    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DescribeSemaphoreResult]:
+        return getattr(self.raw, "describe_semaphore_result", None) or None

ydb/_grpc/grpcwrapper/ydb_coordination_public_types.py

@@ -55,3 +55,31 @@ def from_proto(msg: ydb_coordination_pb2.DescribeNodeResponse) -> "NodeConfig":
         result = ydb_coordination_pb2.DescribeNodeResult()
         msg.operation.result.Unpack(result)
         return NodeConfig.from_proto(result.config)
+
+
+@dataclass
+class AcquireSemaphoreResult:
+    req_id: int
+    acquired: bool
+    status: int
+
+    @staticmethod
+    def from_proto(msg: ydb_coordination_pb2.SessionResponse.AcquireSemaphoreResult) -> "AcquireSemaphoreResult":
+        return AcquireSemaphoreResult(
+            req_id=msg.req_id,
+            acquired=msg.acquired,
+            status=msg.status,
+        )
+
+
+@dataclass
+class CreateSemaphoreResult:
+    req_id: int
+    status: int
+
+    @staticmethod
+    def from_proto(msg: ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult) -> "CreateSemaphoreResult":
+        return CreateSemaphoreResult(
+            req_id=msg.req_id,
+            status=msg.status,
+        )

ydb/aio/__init__.py

@@ -1,4 +1,4 @@
 from .driver import Driver  # noqa
 from .table import SessionPool, retry_operation  # noqa
 from .query import QuerySessionPool, QuerySession, QueryTxContext  # noqa
-from .coordination_client import CoordinationClient  # noqa
+from .coordination import CoordinationClient  # noqa

ydb/aio/coordination/__init__.py

@@ -0,0 +1,5 @@
+__all__ = [
+    "CoordinationClient",
+]
+
+from .client import CoordinationClient