add grpc options with retry/keepalive possible via env vars

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

Author: filintod

durabletask/aio/internal/shared.py

@@ -1,13 +1,14 @@
 # Copyright (c) The Dapr Authors.
 # Licensed under the MIT License.
 
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence, Union, Any
 
 import grpc
 from grpc import aio as grpc_aio
 
 from durabletask.internal.shared import (
     get_default_host_address,
+    build_grpc_channel_options,
     SECURE_PROTOCOLS,
     INSECURE_PROTOCOLS,
 )
@@ -24,8 +25,9 @@
 def get_grpc_aio_channel(
         host_address: Optional[str],
         secure_channel: bool = False,
-        interceptors: Optional[Sequence[ClientInterceptor]] = None) -> grpc_aio.Channel:
-
+        interceptors: Optional[Sequence[ClientInterceptor]] = None,
+        options: Optional[Sequence[tuple[str, Any]]] = None,
+) -> grpc_aio.Channel:
     if host_address is None:
         host_address = get_default_host_address()
 
@@ -41,9 +43,21 @@ def get_grpc_aio_channel(
             host_address = host_address[len(protocol):]
             break
 
+    # Build channel options (merge provided options with env-driven keepalive/retry)
+    channel_options = build_grpc_channel_options(options)
+
+    # Create the base channel
     if secure_channel:
-        channel = grpc_aio.secure_channel(host_address, grpc.ssl_channel_credentials(), interceptors=interceptors)
+        if channel_options is not None:
+            channel = grpc_aio.secure_channel(
+                host_address, grpc.ssl_channel_credentials(), interceptors=interceptors, options=channel_options,
+            )
+        else:
+            channel = grpc_aio.secure_channel(host_address, grpc.ssl_channel_credentials(), interceptors=interceptors)
     else:
-        channel = grpc_aio.insecure_channel(host_address, interceptors=interceptors)
+        if channel_options is not None:
+            channel = grpc_aio.insecure_channel(host_address, interceptors=interceptors, options=channel_options)
+        else:
+            channel = grpc_aio.insecure_channel(host_address, interceptors=interceptors)
 
     return channel

durabletask/internal/shared.py

@@ -4,8 +4,9 @@
 import dataclasses
 import json
 import logging
+import os
 from types import SimpleNamespace
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Union, Iterable
 
 import grpc
 
@@ -24,6 +25,126 @@
 INSECURE_PROTOCOLS = ["http://", "grpc://"]
 
 
+def _get_env_bool(name: str, default: bool) -> bool:
+    """ helper to convert the environment variable to a bool"""
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    return val.strip().lower() in {"1", "true", "t", "yes", "y"}
+
+
+def _get_env_int(name: str, default: int) -> int:
+    """ helper to convert the env var to an int or if we could not to the default value given """
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    try:
+        return int(val)
+    except Exception:
+        return default
+
+
+def _get_env_float(name: str, default: float) -> float:
+    """ helper to convert the env var to a float or if we could not to the default value given """
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    try:
+        return float(val)
+    except Exception:
+        return default
+
+
+def _get_env_csv(name: str, default_csv: str) -> list[str]:
+    """ helper to convert the env var to a list or if we could not to the default value given """
+    val = os.environ.get(name, default_csv)
+    return [s.strip().upper() for s in val.split(",") if s.strip()]
+
+
+def get_grpc_keepalive_options() -> list[tuple[str, Any]]:
+    """Build gRPC keepalive channel options from environment variables.
+
+    Environment variables (defaults in parentheses):
+    - DAPR_GRPC_KEEPALIVE_ENABLED (false)
+    - DAPR_GRPC_KEEPALIVE_TIME_MS (120000)
+    - DAPR_GRPC_KEEPALIVE_TIMEOUT_MS (20000)
+    - DAPR_GRPC_KEEPALIVE_PERMIT_WITHOUT_CALLS (false)
+    """
+    enabled = _get_env_bool("DAPR_GRPC_KEEPALIVE_ENABLED", False)
+    if not enabled:
+        return []
+    time_ms = _get_env_int("DAPR_GRPC_KEEPALIVE_TIME_MS", 120000)
+    timeout_ms = _get_env_int("DAPR_GRPC_KEEPALIVE_TIMEOUT_MS", 20000)
+    permit_without_calls = (
+        1 if _get_env_bool("DAPR_GRPC_KEEPALIVE_PERMIT_WITHOUT_CALLS", False) else 0
+    )
+    return [
+        ("grpc.keepalive_time_ms", time_ms),
+        ("grpc.keepalive_timeout_ms", timeout_ms),
+        ("grpc.keepalive_permit_without_calls", permit_without_calls),
+    ]
+
+
+def get_grpc_retry_service_config_option() -> Optional[tuple[str, str]]:
+    """Return ("grpc.service_config", json_str) if retry is enabled via env; else None.
+
+    Environment variables (defaults in parentheses):
+    - DAPR_GRPC_RETRY_ENABLED (false)
+    - DAPR_GRPC_RETRY_MAX_ATTEMPTS (4)
+    - DAPR_GRPC_RETRY_INITIAL_BACKOFF_MS (100)
+    - DAPR_GRPC_RETRY_MAX_BACKOFF_MS (1000)
+    - DAPR_GRPC_RETRY_BACKOFF_MULTIPLIER (2.0)
+    - DAPR_GRPC_RETRY_CODES (UNAVAILABLE,DEADLINE_EXCEEDED)
+    """
+    enabled = _get_env_bool("DAPR_GRPC_RETRY_ENABLED", False)
+    if not enabled:
+        return None
+
+    max_attempts = _get_env_int("DAPR_GRPC_RETRY_MAX_ATTEMPTS", 4)
+    initial_backoff_ms = _get_env_int("DAPR_GRPC_RETRY_INITIAL_BACKOFF_MS", 100)
+    max_backoff_ms = _get_env_int("DAPR_GRPC_RETRY_MAX_BACKOFF_MS", 1000)
+    backoff_multiplier = _get_env_float("DAPR_GRPC_RETRY_BACKOFF_MULTIPLIER", 2.0)
+    codes = _get_env_csv("DAPR_GRPC_RETRY_CODES", "UNAVAILABLE,DEADLINE_EXCEEDED")
+
+    # service_config ref => https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto#L44
+    service_config = {
+        "methodConfig": [
+            {
+                "name": [{"service": ""}],
+                "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,
+                },
+            }
+        ]
+    }
+    # we are not applying retry throttling policy (but a user can pass the whole option string via options)
+    return "grpc.service_config", json.dumps(service_config)
+
+
+def build_grpc_channel_options(
+        base_options: Optional[Iterable[tuple[str, Any]]] = None,
+) -> Optional[list[tuple[str, Any]]]:
+    """Combine base options + env-driven keepalive and retry service config.
+
+    The returned list is safe to pass as the `options` argument to grpc.secure_channel/insecure_channel.
+    """
+    combined: list[tuple[str, Any]] = []
+    if base_options:
+        combined.extend(list(base_options))
+
+    keepalive = get_grpc_keepalive_options()
+    if keepalive:
+        combined.extend(keepalive)
+    retry_opt = get_grpc_retry_service_config_option()
+    if retry_opt is not None:
+        combined.append(retry_opt)
+    return combined if combined else None
+
+
 def get_default_host_address() -> str:
     """Resolve the default Durable Task sidecar address.
 
@@ -53,7 +174,9 @@ 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:
+        interceptors: Optional[Sequence[ClientInterceptor]] = None,
+        options: Optional[Sequence[tuple[str, Any]]] = None,
+) -> grpc.Channel:
     if host_address is None:
         host_address = get_default_host_address()
 
@@ -71,11 +194,22 @@ def get_grpc_channel(
             host_address = host_address[len(protocol):]
             break
 
+    # Build channel options (merge provided options with env-driven keepalive/retry)
+    channel_options = build_grpc_channel_options(options)
+
     # Create the base channel
     if secure_channel:
-        channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
+        if channel_options is not None:
+            channel = grpc.secure_channel(
+                host_address, grpc.ssl_channel_credentials(), options=channel_options
+            )
+        else:
+            channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
     else:
-        channel = grpc.insecure_channel(host_address)
+        if channel_options is not None:
+            channel = grpc.insecure_channel(host_address, options=channel_options)
+        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,90 @@
+import json
+from unittest.mock import ANY, 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):
+    # Enable keepalive and retry via env
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_ENABLED", "true")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_TIME_MS", "101000")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_TIMEOUT_MS", "30303")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_PERMIT_WITHOUT_CALLS", "true")
+
+    monkeypatch.setenv("DAPR_GRPC_RETRY_ENABLED", "true")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_MAX_ATTEMPTS", "7")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_INITIAL_BACKOFF_MS", "150")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_MAX_BACKOFF_MS", "3000")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_BACKOFF_MULTIPLIER", "2.25")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_CODES", "UNAVAILABLE,RESOURCE_EXHAUSTED")
+
+    base_options = [("another.custom.opt", 99)]
+
+    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"]
+
+        # Base option present
+        assert ("another.custom.opt", 99) in opts
+
+        # Keepalive options present
+        assert ("grpc.keepalive_time_ms", 101000) in opts
+        assert ("grpc.keepalive_timeout_ms", 30303) in opts
+        assert ("grpc.keepalive_permit_without_calls", 1) in opts  # true -> 1
+
+        # 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"] == 7
+        assert retry_policy["initialBackoff"] == f"{150/1000.0}s"
+        assert retry_policy["maxBackoff"] == f"{3000/1000.0}s"
+        assert retry_policy["backoffMultiplier"] == 2.25
+        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", [])

tests/durabletask/test_grpc_channel_options.py

@@ -0,0 +1,80 @@
+import json
+from unittest.mock import ANY, patch
+
+import pytest
+
+from durabletask.internal.shared import get_grpc_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_sync_channel_passes_base_options_and_max_lengths():
+    base_options = [
+        ("grpc.max_send_message_length", 1234),
+        ("grpc.max_receive_message_length", 5678),
+        ("grpc.primary_user_agent", "durabletask-tests"),
+    ]
+    with patch("grpc.insecure_channel") as mock_channel:
+        get_grpc_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", 1234) in opts
+        assert ("grpc.max_receive_message_length", 5678) in opts
+        assert ("grpc.primary_user_agent", "durabletask-tests") in opts
+
+
+def test_sync_channel_merges_env_keepalive_and_retry(monkeypatch: pytest.MonkeyPatch):
+    # Enable keepalive and retry via env
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_ENABLED", "true")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_TIME_MS", "111000")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_TIMEOUT_MS", "22222")
+    monkeypatch.setenv("DAPR_GRPC_KEEPALIVE_PERMIT_WITHOUT_CALLS", "true")
+
+    monkeypatch.setenv("DAPR_GRPC_RETRY_ENABLED", "true")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_MAX_ATTEMPTS", "5")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_INITIAL_BACKOFF_MS", "250")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_MAX_BACKOFF_MS", "2000")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_BACKOFF_MULTIPLIER", "1.5")
+    monkeypatch.setenv("DAPR_GRPC_RETRY_CODES", "UNAVAILABLE,DEADLINE_EXCEEDED,ABORTED")
+
+    base_options = [("custom.opt", 42)]
+
+    with patch("grpc.insecure_channel") as mock_channel:
+        get_grpc_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"]
+
+        # Base option present
+        assert ("custom.opt", 42) in opts
+
+        # Keepalive options present
+        assert ("grpc.keepalive_time_ms", 111000) in opts
+        assert ("grpc.keepalive_timeout_ms", 22222) in opts
+        assert ("grpc.keepalive_permit_without_calls", 1) in opts  # true -> 1
+
+        # 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"] == 5
+        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 "ABORTED" in retry_policy["retryableStatusCodes"]