add options to client argument and e2e test

Signed-off-by: Filinto Duran <1373693+filintod@users.noreply.github.com>

Author: filintod

durabletask/aio/client.py

@@ -20,15 +20,17 @@
 
 
 class AsyncTaskHubGrpcClient:
-
-    def __init__(self, *,
-                 host_address: Optional[str] = None,
-                 metadata: Optional[list[tuple[str, str]]] = None,
-                 log_handler: Optional[logging.Handler] = None,
-                 log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False,
-                 interceptors: Optional[Sequence[ClientInterceptor]] = None):
-
+    def __init__(
+        self,
+        *,
+        host_address: Optional[str] = None,
+        metadata: Optional[list[tuple[str, str]]] = None,
+        log_handler: Optional[logging.Handler] = None,
+        log_formatter: Optional[logging.Formatter] = None,
+        secure_channel: bool = False,
+        interceptors: Optional[Sequence[ClientInterceptor]] = None,
+        channel_options: Optional[Sequence[tuple[str, Any]]] = None,
+    ):
         if interceptors is not None:
             interceptors = list(interceptors)
             if metadata is not None:
@@ -41,7 +43,8 @@ def __init__(self, *,
         channel = get_grpc_aio_channel(
             host_address=host_address,
             secure_channel=secure_channel,
-            interceptors=interceptors
+            interceptors=interceptors,
+            options=channel_options,
         )
         self._channel = channel
         self._stub = stubs.TaskHubSidecarServiceStub(channel)

durabletask/aio/internal/shared.py

@@ -1,49 +1,72 @@
 # Copyright (c) The Dapr Authors.
 # Licensed under the MIT License.
 
-from typing import Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Union
 
 import grpc
 from grpc import aio as grpc_aio
 
 from durabletask.internal.shared import (
-    get_default_host_address,
-    SECURE_PROTOCOLS,
     INSECURE_PROTOCOLS,
+    SECURE_PROTOCOLS,
+    get_default_host_address,
 )
 
-
 ClientInterceptor = Union[
     grpc_aio.UnaryUnaryClientInterceptor,
     grpc_aio.UnaryStreamClientInterceptor,
     grpc_aio.StreamUnaryClientInterceptor,
-    grpc_aio.StreamStreamClientInterceptor
+    grpc_aio.StreamStreamClientInterceptor,
 ]
 
 
 def get_grpc_aio_channel(
-        host_address: Optional[str],
-        secure_channel: bool = False,
-        interceptors: Optional[Sequence[ClientInterceptor]] = None) -> grpc_aio.Channel:
+    host_address: Optional[str],
+    secure_channel: bool = False,
+    interceptors: Optional[Sequence[ClientInterceptor]] = None,
+    options: Optional[Sequence[tuple[str, Any]]] = None,
+) -> grpc_aio.Channel:
+    """create a grpc asyncio channel
 
+    Args:
+        host_address: The host address of the gRPC server. If None, uses the default address.
+        secure_channel: Whether to use a secure channel (TLS/SSL). Defaults to False.
+        interceptors: Optional sequence of client interceptors to apply to the channel.
+        options: Optional sequence of gRPC channel options as (key, value) tuples. Keys defined in https://grpc.github.io/grpc/core/group__grpc__arg__keys.html
+    """
     if host_address is None:
         host_address = get_default_host_address()
 
     for protocol in SECURE_PROTOCOLS:
         if host_address.lower().startswith(protocol):
             secure_channel = True
-            host_address = host_address[len(protocol):]
+            host_address = host_address[len(protocol) :]
             break
 
     for protocol in INSECURE_PROTOCOLS:
         if host_address.lower().startswith(protocol):
             secure_channel = False
-            host_address = host_address[len(protocol):]
+            host_address = host_address[len(protocol) :]
             break
 
+    # Create the base channel
     if secure_channel:
-        channel = grpc_aio.secure_channel(host_address, grpc.ssl_channel_credentials(), interceptors=interceptors)
-    else:
-        channel = grpc_aio.insecure_channel(host_address, interceptors=interceptors)
+        if options is not None:
+            return grpc_aio.secure_channel(
+                host_address,
+                grpc.ssl_channel_credentials(),
+                interceptors=interceptors,
+                options=options,
+            )
+        return grpc_aio.secure_channel(
+            host_address, grpc.ssl_channel_credentials(), interceptors=interceptors
+        )
 
-    return channel
+    if options is not None:
+        # validate all options keys prefix starts with `grpc.`
+        if not all(key.startswith('grpc.') for key, _ in options):
+            raise ValueError(
+                f'All options keys must start with `grpc.`. Invalid options: {options}'
+            )
+        return grpc_aio.insecure_channel(host_address, interceptors=interceptors, options=options)
+    return grpc_aio.insecure_channel(host_address, interceptors=interceptors)

durabletask/client.py

@@ -91,15 +91,17 @@ def new_orchestration_state(instance_id: str, res: pb.GetInstanceResponse) -> Op
 
 
 class TaskHubGrpcClient:
-
-    def __init__(self, *,
-                 host_address: Optional[str] = None,
-                 metadata: Optional[list[tuple[str, str]]] = None,
-                 log_handler: Optional[logging.Handler] = None,
-                 log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False,
-                 interceptors: Optional[Sequence[shared.ClientInterceptor]] = None):
-
+    def __init__(
+        self,
+        *,
+        host_address: Optional[str] = None,
+        metadata: Optional[list[tuple[str, str]]] = None,
+        log_handler: Optional[logging.Handler] = None,
+        log_formatter: Optional[logging.Formatter] = None,
+        secure_channel: bool = False,
+        interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+        channel_options: Optional[Sequence[tuple[str, Any]]] = None,
+    ):
         # If the caller provided metadata, we need to create a new interceptor for it and
         # add it to the list of interceptors.
         if interceptors is not None:
@@ -114,7 +116,8 @@ def __init__(self, *,
         channel = shared.get_grpc_channel(
             host_address=host_address,
             secure_channel=secure_channel,
-            interceptors=interceptors
+            interceptors=interceptors,
+            options=channel_options,
         )
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)

durabletask/internal/shared.py

@@ -51,31 +51,54 @@ def get_default_host_address() -> str:
 
 
 def get_grpc_channel(
-        host_address: Optional[str],
-        secure_channel: bool = False,
-        interceptors: Optional[Sequence[ClientInterceptor]] = None) -> grpc.Channel:
+    host_address: Optional[str],
+    secure_channel: bool = False,
+    interceptors: Optional[Sequence[ClientInterceptor]] = None,
+    options: Optional[Sequence[tuple[str, Any]]] = None,
+) -> grpc.Channel:
+    """create a grpc channel
+
+    Args:
+        host_address: The host address of the gRPC server. If None, uses the default address.
+        secure_channel: Whether to use a secure channel (TLS/SSL). Defaults to False.
+        interceptors: Optional sequence of client interceptors to apply to the channel.
+        options: Optional sequence of gRPC channel options as (key, value) tuples. Keys defined in https://grpc.github.io/grpc/core/group__grpc__arg__keys.html
+    """
     if host_address is None:
         host_address = get_default_host_address()
 
     for protocol in SECURE_PROTOCOLS:
         if host_address.lower().startswith(protocol):
             secure_channel = True
             # remove the protocol from the host name
-            host_address = host_address[len(protocol):]
+            host_address = host_address[len(protocol) :]
             break
 
     for protocol in INSECURE_PROTOCOLS:
         if host_address.lower().startswith(protocol):
             secure_channel = False
             # remove the protocol from the host name
-            host_address = host_address[len(protocol):]
+            host_address = host_address[len(protocol) :]
             break
 
     # Create the base channel
-    if secure_channel:
-        channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
+    if options is not None:
+        # validate all options keys prefix starts with `grpc.`
+        if not all(key.startswith('grpc.') for key, _ in options):
+            raise ValueError(
+                f'All options keys must start with `grpc.`. Invalid options: {options}'
+            )
+        if secure_channel:
+            channel = grpc.secure_channel(
+                host_address, grpc.ssl_channel_credentials(), options=options
+            )
+        else:
+            channel = grpc.insecure_channel(host_address, options=options)
     else:
-        channel = grpc.insecure_channel(host_address)
+        if secure_channel:
+            channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
+        else:
+            channel = grpc.insecure_channel(host_address)
 
     # Apply interceptors ONLY if they exist
     if interceptors:

tests/durabletask/test_grpc_aio_channel_options.py

@@ -0,0 +1,94 @@
+import json
+from unittest.mock import patch
+
+import pytest
+
+from durabletask.aio.internal.shared import get_grpc_aio_channel
+
+HOST_ADDRESS = 'localhost:50051'
+
+
+def _find_option(options, key):
+    for k, v in options:
+        if k == key:
+            return v
+    raise AssertionError(f'Option with key {key} not found in options: {options}')
+
+
+def test_aio_channel_passes_base_options_and_max_lengths():
+    base_options = [
+        ('grpc.max_send_message_length', 4321),
+        ('grpc.max_receive_message_length', 8765),
+        ('grpc.primary_user_agent', 'durabletask-aio-tests'),
+    ]
+    with patch('durabletask.aio.internal.shared.grpc_aio.insecure_channel') as mock_channel:
+        get_grpc_aio_channel(HOST_ADDRESS, False, options=base_options)
+        # Ensure called with options kwarg
+        assert mock_channel.call_count == 1
+        args, kwargs = mock_channel.call_args
+        assert args[0] == HOST_ADDRESS
+        assert 'options' in kwargs
+        opts = kwargs['options']
+        # Check our base options made it through
+        assert ('grpc.max_send_message_length', 4321) in opts
+        assert ('grpc.max_receive_message_length', 8765) in opts
+        assert ('grpc.primary_user_agent', 'durabletask-aio-tests') in opts
+
+
+def test_aio_channel_merges_env_keepalive_and_retry(monkeypatch: pytest.MonkeyPatch):
+    # retry grpc option
+    # service_config ref => https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto#L44
+    max_attempts = 4
+    initial_backoff_ms = 250
+    max_backoff_ms = 2000
+    backoff_multiplier = 1.5
+    codes = ['RESOURCE_EXHAUSTED']
+    service_config = {
+        'methodConfig': [
+            {
+                'name': [{'service': ''}],  # match all services/methods
+                'retryPolicy': {
+                    'maxAttempts': max_attempts,
+                    'initialBackoff': f'{initial_backoff_ms / 1000.0}s',
+                    'maxBackoff': f'{max_backoff_ms / 1000.0}s',
+                    'backoffMultiplier': backoff_multiplier,
+                    'retryableStatusCodes': codes,
+                },
+            }
+        ]
+    }
+
+    base_options = [('grpc.service_config', json.dumps(service_config))]
+
+    with patch('durabletask.aio.internal.shared.grpc_aio.insecure_channel') as mock_channel:
+        get_grpc_aio_channel(HOST_ADDRESS, False, options=base_options)
+
+        args, kwargs = mock_channel.call_args
+        assert args[0] == HOST_ADDRESS
+        assert 'options' in kwargs
+        opts = kwargs['options']
+
+        # Retry service config present and parses correctly
+        svc_cfg_str = _find_option(opts, 'grpc.service_config')
+        svc_cfg = json.loads(svc_cfg_str)
+        assert 'methodConfig' in svc_cfg and isinstance(svc_cfg['methodConfig'], list)
+        retry_policy = svc_cfg['methodConfig'][0]['retryPolicy']
+        assert retry_policy['maxAttempts'] == 4
+        assert retry_policy['initialBackoff'] == f'{250 / 1000.0}s'
+        assert retry_policy['maxBackoff'] == f'{2000 / 1000.0}s'
+        assert retry_policy['backoffMultiplier'] == 1.5
+        # Codes are upper-cased list
+        assert 'RESOURCE_EXHAUSTED' in retry_policy['retryableStatusCodes']
+
+
+def test_aio_secure_channel_receives_options_when_secure_true():
+    base_options = [('grpc.max_receive_message_length', 999999)]
+    with (
+        patch('durabletask.aio.internal.shared.grpc_aio.secure_channel') as mock_channel,
+        patch('grpc.ssl_channel_credentials') as mock_credentials,
+    ):
+        get_grpc_aio_channel(HOST_ADDRESS, True, options=base_options)
+        args, kwargs = mock_channel.call_args
+        assert args[0] == HOST_ADDRESS
+        assert args[1] == mock_credentials.return_value
+        assert ('grpc.max_receive_message_length', 999999) in kwargs.get('options', [])