One Runtime to Rule them All (#157)

Author: njooma

src/viam/robot/client.py

@@ -90,6 +90,7 @@ async def at_address(cls, address: str, options: Options) -> Self:
         Returns:
             Self: the RobotClient
         """
+        logging.setLevel(options.log_level)
         channel = await dial(address, options.dial_options)
         robot = await RobotClient.with_channel(channel, options)
         robot._should_close_channel = True

src/viam/rpc/dial.py

@@ -4,9 +4,10 @@
 import socket
 import ssl
 import sys
+from typing_extensions import Self
 import warnings
 from dataclasses import dataclass
-from typing import Callable, Literal, Optional, Tuple, Type
+from typing import Callable, ClassVar, Literal, Optional, Tuple, Type
 
 from grpclib.client import Channel, Stream
 from grpclib.const import Cardinality
@@ -17,6 +18,7 @@
 from viam.errors import InsecureConnectionError, ViamError
 from viam.proto.rpc.auth import AuthenticateRequest, AuthServiceStub
 from viam.proto.rpc.auth import Credentials as PBCredentials
+from viam.utils import PointerCounter
 
 LOGGER = logging.getLogger(__name__)
 
@@ -151,51 +153,83 @@ async def __aexit__(self, exc_type, exc_value, traceback):
         self.close()
 
 
-async def dial(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
-    opts = options if options else DialOptions()
-    if opts.disable_webrtc:
-        channel = await _dial_direct(address, options)
-        return ViamChannel(channel, lambda: None)
+class _Runtime:
 
-    creds = opts.credentials.payload if opts.credentials else ""
-    insecure = opts.insecure or opts.allow_insecure_with_creds_downgrade or (not creds and opts.allow_insecure_downgrade)
+    _shared: ClassVar[Self]
 
-    libname = pathlib.Path(__file__).parent.absolute() / f"libviam.{'dylib' if sys.platform == 'darwin' else 'so'}"
-    c_lib = ctypes.CDLL(libname.__str__())
-    c_lib.init_rust_runtime.argtypes = ()
-    c_lib.init_rust_runtime.restype = ctypes.c_void_p
+    _lib: ctypes.CDLL
+    _ptr: ctypes.c_void_p
+    _semaphore: PointerCounter = PointerCounter()
 
-    c_lib.dial.argtypes = (ctypes.c_char_p, ctypes.c_char_p, ctypes.c_bool, ctypes.c_void_p)
-    c_lib.dial.restype = ctypes.c_void_p
+    def __new__(cls):
+        if not hasattr(cls, "_shared"):
+            cls._shared = super(_Runtime, cls).__new__(cls)
 
-    c_lib.free_rust_runtime.argtypes = (ctypes.c_void_p,)
-    c_lib.free_rust_runtime.restype = None
+            libname = pathlib.Path(__file__).parent.absolute() / f"libviam.{'dylib' if sys.platform == 'darwin' else 'so'}"
+            cls._shared._lib = ctypes.CDLL(libname.__str__())
+            cls._shared._lib.init_rust_runtime.argtypes = ()
+            cls._shared._lib.init_rust_runtime.restype = ctypes.c_void_p
 
-    c_lib.free_string.argtypes = (ctypes.c_void_p,)
-    c_lib.free_string.restype = None
+            cls._shared._lib.dial.argtypes = (ctypes.c_char_p, ctypes.c_char_p, ctypes.c_bool, ctypes.c_void_p)
+            cls._shared._lib.dial.restype = ctypes.c_void_p
 
-    ptr = c_lib.init_rust_runtime()
+            cls._shared._lib.free_rust_runtime.argtypes = (ctypes.c_void_p,)
+            cls._shared._lib.free_rust_runtime.restype = None
 
-    path_ptr = c_lib.dial(
-        address.encode("utf-8"),
-        creds.encode("utf-8") if creds else None,
-        insecure,
-        ptr,
-    )
-    path = ctypes.cast(path_ptr, ctypes.c_char_p).value
+            cls._shared._lib.free_string.argtypes = (ctypes.c_void_p,)
+            cls._shared._lib.free_string.restype = None
+
+            cls._shared._ptr = cls._shared._lib.init_rust_runtime()
+
+        return cls._shared
+
+    def dial(self, address: str, options: DialOptions) -> Tuple[Optional[str], ctypes.c_void_p]:
+        creds = options.credentials.payload if options.credentials else ""
+        insecure = options.insecure or options.allow_insecure_with_creds_downgrade or (not creds and options.allow_insecure_downgrade)
+
+        path_ptr = self._lib.dial(
+            address.encode("utf-8"),
+            creds.encode("utf-8") if creds else None,
+            insecure,
+            self._ptr,
+        )
+        path = ctypes.cast(path_ptr, ctypes.c_char_p).value
+        path = path.decode("utf-8") if path else ""
+
+        self._semaphore.increment()
+
+        return (path, path_ptr)
+
+    def release(self):
+        self._semaphore.decrement()
+        if self._semaphore.count == 0:
+            self._lib.free_rust_runtime(self._ptr)
+
+    def free_str(self, ptr: ctypes.c_void_p):
+        self._lib.free_string(ptr)
+
+
+async def dial(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
+    opts = options if options else DialOptions()
+    if opts.disable_webrtc:
+        channel = await _dial_direct(address, options)
+        return ViamChannel(channel, lambda: None)
+
+    runtime = _Runtime()
 
+    path, path_ptr = runtime.dial(address, opts)
     if path:
         LOGGER.info(f"Connecting to socket: {path}")
-        chan = Channel(path=path.decode("utf-8"), ssl=None)
+        chan = Channel(path=path, ssl=None)
 
         def release():
-            c_lib.free_string(path_ptr)
-            c_lib.free_rust_runtime(ptr)
+            runtime.free_str(path_ptr)
+            runtime.release()
 
         channel = ViamChannel(chan, release)
         return channel
 
-    c_lib.free_rust_runtime(ptr)
+    runtime.release()
     raise ViamError(f"Unable to establish a connection to {address}")
 
 

src/viam/utils.py

@@ -1,3 +1,4 @@
+from asyncio import Event
 from typing import Any, Dict, List, Mapping, SupportsFloat, Type, TypeVar
 
 from google.protobuf.json_format import MessageToDict, ParseDict
@@ -157,3 +158,29 @@ def sensor_readings_value_to_native(readings: Mapping[str, Value]) -> Mapping[st
             elif kind == "orientation_vector_degrees":
                 prim_readings[key] = Orientation(o_x=reading["ox"], o_y=reading["oy"], o_z=reading["oz"], theta=reading["theta"])
     return prim_readings
+
+
+class PointerCounter:
+    def __init__(self) -> None:
+        self._event = Event()
+        self._count = 0
+        self._event.set()
+
+    def increment(self) -> int:
+        self._count += 1
+        self._event.clear()
+        return self._count
+
+    def decrement(self) -> int:
+        assert self._count > 0, "Pointer count cannot go below zero"
+        self._count -= 1
+        if self._count == 0:
+            self._event.set()
+        return self._count
+
+    async def wait(self) -> None:
+        await self._event.wait()
+
+    @property
+    def count(self) -> int:
+        return self._count

tests/test_utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import pytest
 from google.protobuf.json_format import ParseError
 from google.protobuf.struct_pb2 import ListValue, Struct, Value
@@ -10,6 +11,7 @@
     Vector3,
 )
 from viam.utils import (
+    PointerCounter,
     dict_to_struct,
     message_to_struct,
     primitive_to_value,
@@ -221,3 +223,26 @@ def test_sensor_readings():
     response = sensor_readings_value_to_native(test)
 
     assert response == expected
+
+
+@pytest.mark.asyncio
+async def test_pointer_counter():
+    counter = PointerCounter()
+
+    assert counter.count == 0
+
+    counter.increment()
+    assert counter.count == 1
+
+    async def final_test():
+        assert counter.count > 0
+        await counter.wait()
+        assert counter.count == 0
+
+    task = asyncio.get_running_loop().create_task(final_test())
+
+    await asyncio.sleep(0)  # Needed to start the final_test task
+
+    counter.decrement()
+
+    await task