Client keep alive support

Fixes #395

Author: cretz

temporalio/bridge/client.py

@@ -38,6 +38,14 @@ class ClientRetryConfig:
     max_retries: int
 
 
+@dataclass
+class ClientKeepAliveConfig:
+    """Python representation of the Rust struct for configuring keep alive."""
+
+    interval_millis: int
+    timeout_millis: int
+
+
 @dataclass
 class ClientConfig:
     """Python representation of the Rust struct for configuring the client."""
@@ -47,6 +55,7 @@ class ClientConfig:
     identity: str
     tls_config: Optional[ClientTlsConfig]
     retry_config: Optional[ClientRetryConfig]
+    keep_alive_config: Optional[ClientKeepAliveConfig]
     client_name: str
     client_version: str
 

temporalio/bridge/src/client.rs

@@ -6,9 +6,9 @@ use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
 use temporal_client::{
-    ClientOptions, ClientOptionsBuilder, ConfiguredClient, HealthService, OperatorService,
-    RetryClient, RetryConfig, TemporalServiceClientWithMetrics, TestService, TlsConfig,
-    WorkflowService,
+    ClientKeepAliveConfig as CoreClientKeepAliveConfig, ClientOptions, ClientOptionsBuilder,
+    ConfiguredClient, HealthService, OperatorService, RetryClient, RetryConfig,
+    TemporalServiceClientWithMetrics, TestService, TlsConfig, WorkflowService,
 };
 use tonic::metadata::MetadataKey;
 use url::Url;
@@ -34,6 +34,7 @@ pub struct ClientConfig {
     identity: String,
     tls_config: Option<ClientTlsConfig>,
     retry_config: Option<ClientRetryConfig>,
+    keep_alive_config: Option<ClientKeepAliveConfig>,
 }
 
 #[derive(FromPyObject)]
@@ -54,6 +55,12 @@ struct ClientRetryConfig {
     pub max_retries: usize,
 }
 
+#[derive(FromPyObject)]
+struct ClientKeepAliveConfig {
+    pub interval_millis: u64,
+    pub timeout_millis: u64,
+}
+
 #[derive(FromPyObject)]
 struct RpcCall {
     rpc: String,
@@ -388,7 +395,8 @@ impl TryFrom<ClientConfig> for ClientOptions {
             .retry_config(
                 opts.retry_config
                     .map_or(RetryConfig::default(), |c| c.into()),
-            );
+            )
+            .keep_alive(opts.keep_alive_config.map(Into::into));
         // Builder does not allow us to set option here, so we have to make
         // a conditional to even call it
         if let Some(tls_config) = opts.tls_config {
@@ -437,3 +445,12 @@ impl From<ClientRetryConfig> for RetryConfig {
         }
     }
 }
+
+impl From<ClientKeepAliveConfig> for CoreClientKeepAliveConfig {
+    fn from(conf: ClientKeepAliveConfig) -> Self {
+        CoreClientKeepAliveConfig {
+            interval: Duration::from_millis(conf.interval_millis),
+            timeout: Duration::from_millis(conf.timeout_millis),
+        }
+    }
+}

temporalio/client.py

@@ -51,7 +51,13 @@
 import temporalio.runtime
 import temporalio.service
 import temporalio.workflow
-from temporalio.service import RetryConfig, RPCError, RPCStatusCode, TLSConfig
+from temporalio.service import (
+    KeepAliveConfig,
+    RetryConfig,
+    RPCError,
+    RPCStatusCode,
+    TLSConfig,
+)
 
 from .types import (
     AnyType,
@@ -96,6 +102,7 @@ async def connect(
         ] = None,
         tls: Union[bool, TLSConfig] = False,
         retry_config: Optional[RetryConfig] = None,
+        keep_alive_config: Optional[KeepAliveConfig] = KeepAliveConfig.default,
         rpc_metadata: Mapping[str, str] = {},
         identity: Optional[str] = None,
         lazy: bool = False,
@@ -128,6 +135,10 @@ async def connect(
                 opted in) or all high-level calls made by this client (which all
                 opt-in to retries by default). If unset, a default retry
                 configuration is used.
+            keep_alive_config: Keep-alive configuration for the client
+                connection. Default is to check every 30s and kill the
+                connection if a response doesn't come back in 15s. Can be set to
+                ``None`` to disable.
             rpc_metadata: Headers to use for all calls to the server. Keys here
                 can be overriden by per-call RPC metadata keys.
             identity: Identity for this client. If unset, a default is created
@@ -141,6 +152,7 @@ async def connect(
             target_host=target_host,
             tls=tls,
             retry_config=retry_config,
+            keep_alive_config=keep_alive_config,
             rpc_metadata=rpc_metadata,
             identity=identity or "",
             lazy=lazy,

temporalio/service.py

@@ -11,7 +11,7 @@
 from dataclasses import dataclass, field
 from datetime import timedelta
 from enum import IntEnum
-from typing import Generic, Mapping, Optional, Type, TypeVar, Union
+from typing import ClassVar, Generic, Mapping, Optional, Type, TypeVar, Union
 
 import google.protobuf.empty_pb2
 import google.protobuf.message
@@ -93,13 +93,36 @@ def _to_bridge_config(self) -> temporalio.bridge.client.ClientRetryConfig:
         )
 
 
+@dataclass(frozen=True)
+class KeepAliveConfig:
+    """Keep-alive configuration for client connections."""
+
+    interval_millis: int = 30000
+    """Interval to send HTTP2 keep alive pings."""
+    timeout_millis: int = 15000
+    """Timeout that the keep alive must be responded to within or the connection
+    will be closed."""
+    default: ClassVar[KeepAliveConfig]
+    """Default keep alive config."""
+
+    def _to_bridge_config(self) -> temporalio.bridge.client.ClientKeepAliveConfig:
+        return temporalio.bridge.client.ClientKeepAliveConfig(
+            interval_millis=self.interval_millis,
+            timeout_millis=self.timeout_millis,
+        )
+
+
+KeepAliveConfig.default = KeepAliveConfig()
+
+
 @dataclass
 class ConnectConfig:
     """Config for connecting to the server."""
 
     target_host: str
     tls: Union[bool, TLSConfig] = False
     retry_config: Optional[RetryConfig] = None
+    keep_alive_config: Optional[KeepAliveConfig] = KeepAliveConfig.default
     rpc_metadata: Mapping[str, str] = field(default_factory=dict)
     identity: str = ""
     lazy: bool = False
@@ -142,6 +165,9 @@ def _to_bridge_config(self) -> temporalio.bridge.client.ClientConfig:
             retry_config=self.retry_config._to_bridge_config()
             if self.retry_config
             else None,
+            keep_alive_config=self.keep_alive_config._to_bridge_config()
+            if self.keep_alive_config
+            else None,
             metadata=self.rpc_metadata,
             identity=self.identity,
             client_name="temporal-python",