Support http2 keep-alive

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/client/base/channel.py

@@ -5,6 +5,7 @@
 
 import dataclasses
 import pathlib
+from datetime import timedelta
 from typing import assert_never
 from urllib.parse import parse_qs, urlparse
 
@@ -41,6 +42,24 @@ class SslOptions:
     """
 
 
+@dataclasses.dataclass(frozen=True)
+class KeepAliveOptions:
+    """Options for HTTP2 keep-alive pings."""
+
+    enabled: bool = True
+    """Whether keep-alive should be enabled."""
+
+    interval: timedelta | None = None
+    """The interval between pings.
+
+    If None, keep-alive is disabled."""
+
+    timeout: timedelta = timedelta(seconds=20)
+    """The time in milliseconds to wait for a keep-alive response.
+
+    This is only used if `interval` is not None and `enabled` is `True`."""
+
+
 @dataclasses.dataclass(frozen=True)
 class ChannelOptions:
     """Options for a gRPC channel."""
@@ -51,6 +70,9 @@ class ChannelOptions:
     ssl: SslOptions = SslOptions()
     """SSL options for the channel."""
 
+    keep_alive: KeepAliveOptions = KeepAliveOptions()
+    """HTTP2 keep-alive options for the channel."""
+
 
 def parse_grpc_uri(
     uri: str,
@@ -120,6 +142,23 @@ def parse_grpc_uri(
         parsed_uri.netloc if parsed_uri.port else f"{parsed_uri.netloc}:{defaults.port}"
     )
 
+    channel_options = (
+        [
+            ("grpc.http2.max_pings_without_data", 0),
+            ("grpc.keepalive_permit_without_calls", 1),
+            (
+                "grpc.keepalive_time_ms",
+                (defaults.keep_alive.interval.total_seconds() * 1000),
+            ),
+            (
+                "grpc.keepalive_timeout_ms",
+                defaults.keep_alive.timeout.total_seconds() * 1000,
+            ),
+        ]
+        if defaults.keep_alive.enabled and defaults.keep_alive.interval
+        else None
+    )
+
     ssl = defaults.ssl.enabled if options.ssl is None else options.ssl
     if ssl:
         return secure_channel(
@@ -141,8 +180,9 @@ def parse_grpc_uri(
                     defaults.ssl.certificate_chain,
                 ),
             ),
+            channel_options,
         )
-    return insecure_channel(target)
+    return insecure_channel(target, channel_options)
 
 
 def _to_bool(value: str) -> bool:

tests/test_channel.py

@@ -5,6 +5,7 @@
 
 import dataclasses
 import pathlib
+from datetime import timedelta
 from unittest import mock
 
 import pytest
@@ -13,6 +14,7 @@
 
 from frequenz.client.base.channel import (
     ChannelOptions,
+    KeepAliveOptions,
     SslOptions,
     _to_bool,
     parse_grpc_uri,
@@ -136,6 +138,70 @@ class _ValidUrlTestCase:
                 ),
             ),
         ),
+        _ValidUrlTestCase(
+            title="Keep-alive no defaults",
+            uri="grpc://localhost:1234",
+            defaults=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=True,
+                    interval=timedelta(minutes=5),
+                    timeout=timedelta(minutes=1),
+                ),
+            ),
+            expected_host="localhost",
+            expected_port=1234,
+            expected_options=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=True,
+                    interval=timedelta(minutes=5),
+                    timeout=timedelta(minutes=1),
+                ),
+            ),
+        ),
+        _ValidUrlTestCase(
+            title="Keep-alive default timeout",
+            uri="grpc://localhost:1234",
+            defaults=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=False, interval=timedelta(minutes=5)
+                ),
+            ),
+            expected_host="localhost",
+            expected_port=1234,
+            expected_options=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=False, interval=timedelta(minutes=5)
+                ),
+            ),
+        ),
+        _ValidUrlTestCase(
+            title="Keep-alive default interval",
+            uri="grpc://localhost:1234",
+            defaults=ChannelOptions(
+                keep_alive=KeepAliveOptions(enabled=True, timeout=timedelta(minutes=1)),
+            ),
+            expected_host="localhost",
+            expected_port=1234,
+            expected_options=ChannelOptions(
+                keep_alive=KeepAliveOptions(enabled=True, timeout=timedelta(minutes=1)),
+            ),
+        ),
+        _ValidUrlTestCase(
+            title="keep-alive no defaults",
+            uri="grpc://localhost:1234",
+            defaults=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=False, interval=None, timeout=timedelta(minutes=1)
+                ),
+            ),
+            expected_host="localhost",
+            expected_port=1234,
+            expected_options=ChannelOptions(
+                keep_alive=KeepAliveOptions(
+                    enabled=False, interval=None, timeout=timedelta(minutes=1)
+                ),
+            ),
+        ),
     ],
     ids=lambda case: case.title,
 )
@@ -198,6 +264,22 @@ def test_parse_uri_ok(  # pylint: disable=too-many-locals
 
     assert channel == expected_channel
     expected_target = f"{expected_host}:{expected_port}"
+    expected_channel_options = (
+        [
+            ("grpc.http2.max_pings_without_data", 0),
+            ("grpc.keepalive_permit_without_calls", 1),
+            (
+                "grpc.keepalive_time_ms",
+                (defaults.keep_alive.interval.total_seconds() * 1000),
+            ),
+            (
+                "grpc.keepalive_timeout_ms",
+                defaults.keep_alive.timeout.total_seconds() * 1000,
+            ),
+        ]
+        if defaults.keep_alive.enabled and defaults.keep_alive.interval
+        else None
+    )
     if expected_ssl:
         if isinstance(expected_root_certificates, pathlib.Path):
             get_contents_mock.assert_any_call(
@@ -223,10 +305,12 @@ def test_parse_uri_ok(  # pylint: disable=too-many-locals
             certificate_chain=expected_certificate_chain,
         )
         secure_channel_mock.assert_called_once_with(
-            expected_target, expected_credentials
+            expected_target, expected_credentials, expected_channel_options
         )
     else:
-        insecure_channel_mock.assert_called_once_with(expected_target)
+        insecure_channel_mock.assert_called_once_with(
+            expected_target, expected_channel_options
+        )
 
 
 @pytest.mark.parametrize("value", ["true", "on", "1", "TrUe", "On", "ON", "TRUE"])