RSDK-5960 - initial connection options (#853)

Author: stuqdog

src/viam/robot/client.py

@@ -52,7 +52,7 @@
 from viam.resource.registry import Registry
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase, ResourceRPCClientBase
 from viam.resource.types import API, RESOURCE_TYPE_COMPONENT, RESOURCE_TYPE_SERVICE
-from viam.rpc.dial import DialOptions, ViamChannel, dial
+from viam.rpc.dial import DialOptions, ViamChannel, dial, _dial_inner
 from viam.services.service_base import ServiceBase
 from viam.sessions_client import SessionsClient
 from viam.utils import datetime_to_timestamp, dict_to_struct
@@ -403,7 +403,7 @@ async def _check_connection(self, check_every: int, reconnect_every: int):
                 try:
                     self._sessions_client.reset()
 
-                    channel = await dial(self._address, self._options.dial_options)
+                    channel = await _dial_inner(self._address, self._options.dial_options)
 
                     client: RobotServiceStub
                     if isinstance(channel, Channel):

src/viam/rpc/dial.py

@@ -70,6 +70,14 @@ class DialOptions:
     max_reconnect_attempts: int = 3
     """Max number of times the client attempts to reconnect when connection is lost"""
 
+    initial_connection_attempts: int = 3
+    """Max number of times the client will attempt to establish an initial connection
+    If set to a non-positive integer, then there will be no limit to initial connection attempts"""
+
+    initial_connection_attempt_timeout: float
+    """Number of seconds before dial connection times out on initial connection attempts
+    Defaults to whatever value is set in the `timeout` field"""
+
     timeout: float = 20
     """Number of seconds before the dial connection times out
     Set to 20sec to match _defaultOfferDeadline in goutils/rpc/wrtc_call_queue.go"""
@@ -85,6 +93,8 @@ def __init__(
         allow_insecure_with_creds_downgrade: bool = False,
         max_reconnect_attempts: int = 3,
         timeout: float = 20,
+        initial_connection_attempts: int = 3,
+        initial_connection_attempt_timeout: Optional[float] = None,
     ) -> None:
         self.disable_webrtc = disable_webrtc
         self.auth_entity = auth_entity
@@ -94,6 +104,8 @@ def __init__(
         self.allow_insecure_with_creds_downgrade = allow_insecure_with_creds_downgrade
         self.max_reconnect_attempts = max_reconnect_attempts
         self.timeout = timeout
+        self.initial_connection_attempts = initial_connection_attempts
+        self.initial_connection_attempt_timeout = initial_connection_attempt_timeout if initial_connection_attempt_timeout else timeout
 
     @classmethod
     def with_api_key(cls, api_key: str, api_key_id: str) -> Self:
@@ -279,6 +291,27 @@ def free_str(self, ptr: ctypes.c_void_p):
 
 
 async def dial(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
+    options = options if options else DialOptions()
+    timeout = options.timeout
+    options.timeout = options.initial_connection_attempt_timeout
+    if options.initial_connection_attempts == 0:
+        options.initial_connection_attempts = -1
+    attempt_countdown = options.initial_connection_attempts
+    exception: Exception
+    while attempt_countdown != 0:
+        try:
+            chan = await _dial_inner(address, options)
+            options.timeout = timeout
+            return chan
+        except Exception as e:
+            exception = e
+            attempt_countdown -= 1
+    # the only way we could get here is if we failed at least once which means we've set the
+    # exception, so typechecker concerns about a possibly unbounded variable are unfounded
+    raise exception # type: ignore
+
+
+async def _dial_inner(address: str, options: Optional[DialOptions] = None) -> ViamChannel:
     async def send_request(event: SendRequest):
         event.metadata["viam-client"] = f"python;v{SDK_VERSION};v{API_VERSION}"