feat(sources): add gRPC max connection age

Author: fpytloun

src/components/validation/runner/io.rs

@@ -22,7 +22,7 @@ use crate::{
         Client as VectorClient, HealthCheckRequest, HealthCheckResponse, PushEventsRequest,
         PushEventsResponse, Server as VectorServer, Service as VectorService, ServingStatus,
     },
-    sources::util::grpc::run_grpc_server,
+    sources::util::grpc::{GrpcKeepaliveConfig, run_grpc_server},
 };
 
 #[derive(Clone)]
@@ -166,6 +166,7 @@ pub fn spawn_grpc_server<S>(
             listen_addr.as_socket_addr(),
             tls_settings,
             service,
+            GrpcKeepaliveConfig::default(),
             shutdown_signal,
         );
         pin!(server);

src/sources/opentelemetry/config.rs

@@ -14,7 +14,7 @@ use crate::{
             grpc::Service,
             http::{build_warp_filter, run_http_server},
         },
-        util::grpc::run_grpc_server_with_routes,
+        util::grpc::{GrpcKeepaliveConfig, run_grpc_server_with_routes},
     },
 };
 use futures::FutureExt;
@@ -173,12 +173,17 @@ pub struct GrpcConfig {
     #[configurable(derived)]
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub tls: Option<TlsEnableableConfig>,
+
+    #[configurable(derived)]
+    #[serde(default)]
+    pub keepalive: GrpcKeepaliveConfig,
 }
 
 fn example_grpc_config() -> GrpcConfig {
     GrpcConfig {
         address: "0.0.0.0:4317".parse().unwrap(),
         tls: None,
+        keepalive: GrpcKeepaliveConfig::default(),
     }
 }
 
@@ -325,6 +330,7 @@ impl SourceConfig for OpentelemetryConfig {
             self.grpc.address,
             grpc_tls_settings,
             builder.routes(),
+            self.grpc.keepalive.clone(),
             cx.shutdown.clone(),
         )
         .map_err(|error| {

src/sources/opentelemetry/integration_tests.rs

@@ -53,6 +53,7 @@ async fn receive_logs_legacy_namespace() {
             grpc: GrpcConfig {
                 address: source_grpc_address().parse().unwrap(),
                 tls: Default::default(),
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: source_http_address().parse().unwrap(),
@@ -152,6 +153,7 @@ async fn receive_trace() {
             grpc: GrpcConfig {
                 address: source_grpc_address().parse().unwrap(),
                 tls: Default::default(),
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: source_http_address().parse().unwrap(),
@@ -257,6 +259,7 @@ async fn receive_metric() {
             grpc: GrpcConfig {
                 address: source_grpc_address().parse().unwrap(),
                 tls: Default::default(),
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: source_http_address().parse().unwrap(),

src/sources/opentelemetry/tests.rs

@@ -206,6 +206,30 @@ fn generate_config() {
     test_util::test_generate_config::<OpentelemetryConfig>();
 }
 
+#[test]
+fn config_grpc_keepalive() {
+    let config: OpentelemetryConfig = toml::from_str(
+        r#"
+            [grpc]
+            address = "0.0.0.0:4317"
+
+            [grpc.keepalive]
+            max_connection_age_secs = 300
+            max_connection_age_grace_secs = 30
+
+            [http]
+            address = "0.0.0.0:4318"
+        "#,
+    )
+    .unwrap();
+
+    assert_eq!(config.grpc.keepalive.max_connection_age_secs, Some(300));
+    assert_eq!(
+        config.grpc.keepalive.max_connection_age_grace_secs,
+        Some(30)
+    );
+}
+
 #[tokio::test]
 async fn receive_grpc_logs_vector_namespace() {
     assert_source_compliance(&SOURCE_TAGS, async {
@@ -1175,6 +1199,7 @@ fn get_source_config_with_headers(
         grpc: GrpcConfig {
             address: grpc_addr,
             tls: Default::default(),
+            keepalive: Default::default(),
         },
         http: HttpConfig {
             address: http_addr,
@@ -1510,6 +1535,7 @@ pub async fn build_otlp_test_env(
         grpc: GrpcConfig {
             address: grpc_addr,
             tls: Default::default(),
+            keepalive: Default::default(),
         },
         http: HttpConfig {
             address: http_addr,
@@ -1589,6 +1615,7 @@ async fn http_logs_use_otlp_decoding_emits_metric() {
         grpc: GrpcConfig {
             address: grpc_addr,
             tls: Default::default(),
+            keepalive: Default::default(),
         },
         http: HttpConfig {
             address: http_addr,
@@ -1823,6 +1850,7 @@ mod otlp_decoding_config_tests {
             grpc: GrpcConfig {
                 address: "0.0.0.0:4317".parse().unwrap(),
                 tls: None,
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: "0.0.0.0:4318".parse().unwrap(),
@@ -1863,6 +1891,7 @@ mod otlp_decoding_config_tests {
             grpc: GrpcConfig {
                 address: "0.0.0.0:4317".parse().unwrap(),
                 tls: None,
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: "0.0.0.0:4318".parse().unwrap(),
@@ -1906,6 +1935,7 @@ mod otlp_decoding_config_tests {
             grpc: GrpcConfig {
                 address: "0.0.0.0:4317".parse().unwrap(),
                 tls: None,
+                keepalive: Default::default(),
             },
             http: HttpConfig {
                 address: "0.0.0.0:4318".parse().unwrap(),

src/sources/util/grpc/mod.rs

@@ -1,12 +1,23 @@
-use std::{convert::Infallible, net::SocketAddr, time::Duration};
+use std::{
+    convert::Infallible,
+    net::SocketAddr,
+    pin::Pin,
+    task::{Context, Poll},
+    time::Duration,
+};
 
-use futures::FutureExt;
+use futures::{FutureExt, StreamExt};
 use http::{Request, Response};
 use hyper::Body;
+use tokio::{
+    io::{AsyncRead, AsyncWrite, ReadBuf},
+    net::TcpStream,
+    time::{Sleep, sleep},
+};
 use tonic::{
     body::BoxBody,
     server::NamedService,
-    transport::server::{Routes, Server},
+    transport::server::{Connected, Routes, Server},
 };
 use tower::Service;
 use tower_http::{
@@ -18,16 +29,127 @@ use tracing::Span;
 use crate::{
     internal_events::{GrpcServerRequestReceived, GrpcServerResponseSent},
     shutdown::{ShutdownSignal, ShutdownSignalToken},
-    tls::MaybeTlsSettings,
+    tls::{MaybeTlsIncomingStream, MaybeTlsSettings},
 };
+use vector_lib::configurable::configurable_component;
 
 mod decompression;
 pub use self::decompression::{DecompressionAndMetrics, DecompressionAndMetricsLayer};
 
+/// Configuration of gRPC server keepalive parameters.
+#[configurable_component]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+#[serde(deny_unknown_fields)]
+pub struct GrpcKeepaliveConfig {
+    /// The maximum amount of time a connection may exist before the server closes it.
+    ///
+    /// When unset, connections are not closed based on age.
+    #[serde(default)]
+    #[configurable(metadata(docs::examples = 300))]
+    #[configurable(metadata(docs::type_unit = "seconds"))]
+    #[configurable(metadata(docs::human_name = "Maximum Connection Age"))]
+    pub max_connection_age_secs: Option<u64>,
+
+    /// The grace period added to `max_connection_age_secs` before the server closes the connection.
+    ///
+    /// This setting only applies when `max_connection_age_secs` is set.
+    #[serde(default)]
+    #[configurable(metadata(docs::examples = 30))]
+    #[configurable(metadata(docs::type_unit = "seconds"))]
+    #[configurable(metadata(docs::human_name = "Maximum Connection Age Grace"))]
+    pub max_connection_age_grace_secs: Option<u64>,
+}
+
+impl GrpcKeepaliveConfig {
+    fn max_connection_lifetime(&self) -> Option<Duration> {
+        self.max_connection_age_secs.map(|max_connection_age_secs| {
+            let age = Duration::from_secs(max_connection_age_secs);
+            let grace = self
+                .max_connection_age_grace_secs
+                .map(Duration::from_secs)
+                .unwrap_or_default();
+
+            age.checked_add(grace).unwrap_or(Duration::MAX)
+        })
+    }
+}
+
+struct MaxConnectionAgeIo {
+    inner: MaybeTlsIncomingStream<TcpStream>,
+    deadline: Option<Pin<Box<Sleep>>>,
+}
+
+impl MaxConnectionAgeIo {
+    fn new(inner: MaybeTlsIncomingStream<TcpStream>, lifetime: Option<Duration>) -> Self {
+        Self {
+            inner,
+            deadline: lifetime.map(|lifetime| Box::pin(sleep(lifetime))),
+        }
+    }
+
+    fn is_expired(&mut self, cx: &mut Context<'_>) -> bool {
+        self.deadline
+            .as_mut()
+            .is_some_and(|deadline| deadline.as_mut().poll(cx).is_ready())
+    }
+}
+
+impl AsyncRead for MaxConnectionAgeIo {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        let this = self.get_mut();
+        if this.is_expired(cx) {
+            Poll::Ready(Ok(()))
+        } else {
+            Pin::new(&mut this.inner).poll_read(cx, buf)
+        }
+    }
+}
+
+impl AsyncWrite for MaxConnectionAgeIo {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        let this = self.get_mut();
+        if this.is_expired(cx) {
+            Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()))
+        } else {
+            Pin::new(&mut this.inner).poll_write(cx, buf)
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        let this = self.get_mut();
+        if this.is_expired(cx) {
+            Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()))
+        } else {
+            Pin::new(&mut this.inner).poll_flush(cx)
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
+    }
+}
+
+impl Connected for MaxConnectionAgeIo {
+    type ConnectInfo = <MaybeTlsIncomingStream<TcpStream> as Connected>::ConnectInfo;
+
+    fn connect_info(&self) -> Self::ConnectInfo {
+        self.inner.connect_info()
+    }
+}
+
 pub async fn run_grpc_server<S>(
     address: SocketAddr,
     tls_settings: MaybeTlsSettings,
     service: S,
+    keepalive: GrpcKeepaliveConfig,
     shutdown: ShutdownSignal,
 ) -> crate::Result<()>
 where
@@ -41,7 +163,10 @@ where
     let span = Span::current();
     let (tx, rx) = tokio::sync::oneshot::channel::<ShutdownSignalToken>();
     let listener = tls_settings.bind(&address).await?;
-    let stream = listener.accept_stream();
+    let max_connection_lifetime = keepalive.max_connection_lifetime();
+    let stream = listener
+        .accept_stream()
+        .map(move |stream| stream.map(|io| MaxConnectionAgeIo::new(io, max_connection_lifetime)));
 
     info!(%address, "Building gRPC server.");
 
@@ -72,12 +197,16 @@ pub async fn run_grpc_server_with_routes(
     address: SocketAddr,
     tls_settings: MaybeTlsSettings,
     routes: Routes,
+    keepalive: GrpcKeepaliveConfig,
     shutdown: ShutdownSignal,
 ) -> crate::Result<()> {
     let span = Span::current();
     let (tx, rx) = tokio::sync::oneshot::channel::<ShutdownSignalToken>();
     let listener = tls_settings.bind(&address).await?;
-    let stream = listener.accept_stream();
+    let max_connection_lifetime = keepalive.max_connection_lifetime();
+    let stream = listener
+        .accept_stream()
+        .map(move |stream| stream.map(|io| MaxConnectionAgeIo::new(io, max_connection_lifetime)));
 
     info!(%address, "Building gRPC server.");