[RSDK-11569] fix race when creating session (#1001)

Author: npmenard

src/viam/sessions_client.py

@@ -7,6 +7,8 @@
 from enum import IntEnum
 from threading import Lock, Thread
 from typing import MutableMapping, Optional
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
 
 from grpclib import Status
 from grpclib.client import Channel
@@ -46,6 +48,7 @@ class SessionsClient:
     _heartbeat_interval: Optional[timedelta]
     _supported: _SupportedState
     _thread: Optional[Thread]
+    _pool: ThreadPoolExecutor
 
     _HEARTBEAT_MONITORED_METHODS: MutableMapping[str, bool] = {}
 
@@ -71,6 +74,7 @@ def __init__(
         self._heartbeat_interval = None
         self._supported = _SupportedState.UNKNOWN
         self._thread = None
+        self._pool = ThreadPoolExecutor()
 
         listen(self.channel, SendRequest, self._send_request)
         listen(self.channel, RecvTrailingMetadata, self._recv_trailers)
@@ -104,40 +108,45 @@ async def _recv_trailers(self, event: RecvTrailingMetadata):
         if event.status == Status.INVALID_ARGUMENT and event.status_message == "SESSION_EXPIRED":
             LOGGER.debug("Session expired")
             self.reset()
+    @asynccontextmanager
+    async def _acquire_lock_async(self):
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(self._pool, self._lock.acquire)
+        try:
+            yield
+        finally:
+            self._lock.release()
 
     @property
     async def metadata(self) -> _MetadataLike:
-        with self._lock:
+        async with self._acquire_lock_async():
             if self._disabled or self._supported != _SupportedState.UNKNOWN:
                 return self._metadata
 
-        request = StartSessionRequest(resume=self._current_id)
-        try:
-            response: StartSessionResponse = await self.client.StartSession(request)
-        except GRPCError as error:
-            if error.status == Status.UNIMPLEMENTED:
-                with self._lock:
-                    self._reset()
-                    self._supported = _SupportedState.FALSE
-                    return self._metadata
-            else:
-                raise
-
-        if response is None:
-            raise GRPCError(status=Status.INTERNAL, message="Expected response to start session")
-
-        if response.heartbeat_window is None:
-            raise GRPCError(status=Status.INTERNAL, message="Expected heartbeat window in response to start session")
+            request = StartSessionRequest(resume=self._current_id)
+            try:
+                response: StartSessionResponse = await self.client.StartSession(request)
+            except GRPCError as error:
+                if error.status == Status.UNIMPLEMENTED:
+                        self._reset()
+                        self._supported = _SupportedState.FALSE
+                        return self._metadata
+                else:
+                    raise
+
+            if response is None:
+                raise GRPCError(status=Status.INTERNAL, message="Expected response to start session")
+
+            if response.heartbeat_window is None:
+                raise GRPCError(status=Status.INTERNAL, message="Expected heartbeat window in response to start session")
 
-        with self._lock:
             self._supported = _SupportedState.TRUE
             self._heartbeat_interval = response.heartbeat_window.ToTimedelta()
             self._current_id = response.id
 
-        # tick once to ensure heartbeats are supported
-        await self._heartbeat_tick(self.client)
+            # tick once to ensure heartbeats are supported
+            await self._heartbeat_tick(self.client)
 
-        with self._lock:
             if self._thread is not None:
                 self._reset()
             if self._supported == _SupportedState.TRUE:
@@ -156,17 +165,16 @@ async def metadata(self) -> _MetadataLike:
             return self._metadata
 
     async def _heartbeat_tick(self, client: RobotServiceStub):
-        with self._lock:
-            if not self._current_id:
-                LOGGER.debug("Failed to send heartbeat, session client reset")
-                return
-            request = SendSessionHeartbeatRequest(id=self._current_id)
+        if not self._current_id:
+            LOGGER.debug("Failed to send heartbeat, session client reset")
+            return
+        request = SendSessionHeartbeatRequest(id=self._current_id)
 
         try:
             await client.SendSessionHeartbeat(request)
         except (GRPCError, StreamTerminatedError):
             LOGGER.debug("Heartbeat terminated", exc_info=True)
-            self.reset()
+            self._reset()
         else:
             LOGGER.debug("Sent heartbeat successfully")
 
@@ -185,10 +193,10 @@ async def _heartbeat_process(self, wait: float):
         channel = await dial(address=addr, options=self._dial_options)
         client = RobotServiceStub(channel.channel)
         while True:
-            with self._lock:
+            async with self._acquire_lock_async():
                 if self._supported != _SupportedState.TRUE:
                     return
-            await self._heartbeat_tick(client)
+                await self._heartbeat_tick(client)
             await asyncio.sleep(wait)
 
     @property

tests/test_sessions_client.py

@@ -106,7 +106,12 @@ async def test_sessions_heartbeat_thread_blocked():
     channel = await dial(address=addr, options=options)
 
     client = SessionsClient(channel.channel, addr, options)
-    assert await client.metadata == {SESSION_METADATA_KEY: MockRobot.SESSION_ID}
+    t1 = asyncio.create_task(client.metadata)
+    t2 = asyncio.create_task(client.metadata)
+
+    await asyncio.gather(t1,t2)
+    assert t1.result() == {SESSION_METADATA_KEY: MockRobot.SESSION_ID}
+    assert t2.result() == {SESSION_METADATA_KEY: MockRobot.SESSION_ID}
 
     assert client._supported == _SupportedState.TRUE
     assert client._heartbeat_interval and client._heartbeat_interval.total_seconds() == MockRobot.HEARTBEAT_INTERVAL