Add gRPC keepalive max_connection_age settings on OTel sources (#6745)

data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/server/CreateServer.java

@@ -179,6 +179,12 @@ public Server createGRPCServer(
         if(serverConfiguration.getMaxRequestLength() != null) {
             sb.maxRequestLength(serverConfiguration.getMaxRequestLength().getBytes());
         }
+        if (serverConfiguration.getMaxConnectionAge() != null) {
+            sb.maxConnectionAge(serverConfiguration.getMaxConnectionAge());
+        }
+        if (serverConfiguration.getConnectionDrainDuration() != null) {
+            sb.connectionDrainDuration(serverConfiguration.getConnectionDrainDuration());
+        }
 
         // ACM Cert for SSL takes preference
         if (serverConfiguration.isSsl() || serverConfiguration.useAcmCertForSSL()) {

data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/server/ServerConfiguration.java

@@ -11,6 +11,8 @@
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.codec.CompressionOption;
 
+import java.time.Duration;
+
 public class ServerConfiguration {
     static final int DEFAULT_REQUEST_TIMEOUT_MS = 10000;
     static final boolean DEFAULT_ENABLED_UNFRAMED_REQUESTS = false;
@@ -88,6 +90,14 @@ public class ServerConfiguration {
     @Setter
     private int bufferTimeoutInMillis;
 
+    @Getter
+    @Setter
+    private Duration maxConnectionAge;
+
+    @Getter
+    @Setter
+    private Duration connectionDrainDuration;
+
     public boolean hasHealthCheck() {
         return healthCheck;
     }

data-prepper-plugins/http-common/src/test/java/org/opensearch/dataprepper/plugins/server/CreateServerTest.java

@@ -38,6 +38,7 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -46,6 +47,8 @@
 import java.util.Optional;
 import java.util.function.Function;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.mockito.Mockito.mock;
@@ -86,6 +89,44 @@ public class CreateServerTest {
     @Mock
     private Certificate certificate;
 
+    @Test
+    void createGrpcServer_appliesKeepaliveSettingsWhenConfigured() throws JsonProcessingException {
+        when(authenticationProvider.getAuthenticationInterceptor()).thenReturn(authenticationInterceptor);
+        final Map<String, Object> metadata = createGrpcMetadata(21895, false, 10000, 10, 5, CompressionOption.NONE,
+                null);
+        final ServerConfiguration serverConfiguration = createServerConfig(metadata);
+        serverConfiguration.setMaxConnectionAge(Duration.ofMinutes(30));
+        serverConfiguration.setConnectionDrainDuration(Duration.ofSeconds(15));
+        final CreateServer createServer = new CreateServer(serverConfiguration, LOG, pluginMetrics, TEST_SOURCE_NAME,
+                TEST_PIPELINE_NAME);
+        Buffer<Record<? extends Metric>> buffer = new BlockingBuffer<Record<? extends Metric>>(TEST_PIPELINE_NAME);
+        TestService testService = getTestService(buffer);
+
+        Server server = createServer.createGRPCServer(authenticationProvider, testService, certificateProvider, null);
+
+        assertNotNull(server);
+        assertThat(server.config().maxConnectionAgeMillis(), equalTo(Duration.ofMinutes(30).toMillis()));
+        assertThat(server.config().connectionDrainDurationMicros(),
+                equalTo(Duration.ofSeconds(15).toNanos() / 1_000L));
+    }
+
+    @Test
+    void createGrpcServer_leavesKeepaliveDefaultsWhenNotConfigured() throws JsonProcessingException {
+        when(authenticationProvider.getAuthenticationInterceptor()).thenReturn(authenticationInterceptor);
+        final Map<String, Object> metadata = createGrpcMetadata(21896, false, 10000, 10, 5, CompressionOption.NONE,
+                null);
+        final ServerConfiguration serverConfiguration = createServerConfig(metadata);
+        final CreateServer createServer = new CreateServer(serverConfiguration, LOG, pluginMetrics, TEST_SOURCE_NAME,
+                TEST_PIPELINE_NAME);
+        Buffer<Record<? extends Metric>> buffer = new BlockingBuffer<Record<? extends Metric>>(TEST_PIPELINE_NAME);
+        TestService testService = getTestService(buffer);
+
+        Server server = createServer.createGRPCServer(authenticationProvider, testService, certificateProvider, null);
+
+        assertNotNull(server);
+        assertThat(server.config().maxConnectionAgeMillis(), equalTo(0L));
+    }
+
     @Test
     void createGrpcServerTest() throws JsonProcessingException {
         when(authenticationProvider.getAuthenticationInterceptor()).thenReturn(authenticationInterceptor);

data-prepper-plugins/otel-logs-source/README.md

@@ -22,6 +22,8 @@ source:
 * unframed_requests(Optional) => A boolean to enable requests not framed using the gRPC wire protocol. 
 * thread_count(Optional) => the number of threads to keep in the ScheduledThreadPool. Default is `200`.
 * max_connection_count(Optional) => the maximum allowed number of open connections. Default is `500`.
+* max_connection_age(Optional) => An ISO-8601 duration string (for example `PT30M`) that sets the maximum age of a gRPC server connection before it is gracefully closed. Useful when horizontally scaling, because long-lived gRPC channels do not re-resolve DNS on their own. When unset (default), Armeria places no upper bound on connection age. Allowed range: 1s to 24h.
+* connection_drain_duration(Optional) => An ISO-8601 duration string (for example `PT15S`) that controls how long Armeria waits after `max_connection_age` expires before forcefully terminating outstanding RPCs. When unset (default), Armeria uses its built-in drain duration. Allowed range: 0ms to 1h.
 * compression (Optional) : The compression type applied on the client request payload. Defaults to `none`. Supported values are:
     * `none`: no compression
     * `gzip`: apply GZip de-compression on the incoming request.

data-prepper-plugins/otel-logs-source/build.gradle

@@ -31,6 +31,7 @@ dependencies {
     implementation libs.grpc.inprocess
     implementation 'com.fasterxml.jackson.core:jackson-databind'
     implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
+    implementation libs.hibernate.validator
     implementation libs.commons.lang3
     implementation libs.bouncycastle.bcprov
     implementation libs.bouncycastle.bcpkix

data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsSource.java

@@ -163,6 +163,12 @@ private Server createServer(ServerBuilder serverBuilder, Buffer<Record<Object>>
         if (oTelLogsSourceConfig.getMaxRequestLength() != null) {
             serverBuilder.maxRequestLength(oTelLogsSourceConfig.getMaxRequestLength().getBytes());
         }
+        if (oTelLogsSourceConfig.getMaxConnectionAge() != null) {
+            serverBuilder.maxConnectionAge(oTelLogsSourceConfig.getMaxConnectionAge());
+        }
+        if (oTelLogsSourceConfig.getConnectionDrainDuration() != null) {
+            serverBuilder.connectionDrainDuration(oTelLogsSourceConfig.getConnectionDrainDuration());
+        }
         final int threadCount = oTelLogsSourceConfig.getThreadCount();
         ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(threadCount);
         serverBuilder.blockingTaskExecutor(executor, true);

data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsSourceConfig.java

@@ -15,12 +15,16 @@
 import lombok.NoArgsConstructor;
 
 import org.apache.commons.lang3.StringUtils;
+import org.hibernate.validator.constraints.time.DurationMax;
+import org.hibernate.validator.constraints.time.DurationMin;
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.codec.CompressionOption;
 import org.opensearch.dataprepper.model.configuration.PluginModel;
 import org.opensearch.dataprepper.plugins.otel.codec.OTelOutputFormat;
 import org.opensearch.dataprepper.plugins.server.RetryInfoConfig;
 
+import java.time.Duration;
+
 @Builder
 @NoArgsConstructor
 @AllArgsConstructor
@@ -126,6 +130,16 @@ public class OTelLogsSourceConfig {
     @JsonProperty(RETRY_INFO)
     private RetryInfoConfig retryInfo;
 
+    @JsonProperty("max_connection_age")
+    @DurationMin(seconds = 1)
+    @DurationMax(hours = 24)
+    private Duration maxConnectionAge;
+
+    @JsonProperty("connection_drain_duration")
+    @DurationMin(millis = 0)
+    @DurationMax(hours = 1)
+    private Duration connectionDrainDuration;
+
     @AssertTrue(message = "path should start with /")
     boolean isPathValid() {
         return path == null || path.startsWith("/");
@@ -251,5 +265,13 @@ public RetryInfoConfig getRetryInfo() {
     public void setRetryInfo(RetryInfoConfig retryInfo) {
         this.retryInfo = retryInfo;
     }
+
+    public Duration getMaxConnectionAge() {
+        return maxConnectionAge;
+    }
+
+    public Duration getConnectionDrainDuration() {
+        return connectionDrainDuration;
+    }
 }
 

data-prepper-plugins/otel-logs-source/src/test/java/org/opensearch/dataprepper/plugins/source/otellogs/OtelLogsSourceConfigTests.java

@@ -304,6 +304,24 @@ void testRetryInfoConfig() {
         assertThat(retryInfo.getMinDelay(), equalTo(Duration.ofMillis(50)));
     }
 
+    @Test
+    void maxConnectionAge_defaultsToNull() {
+        final OTelLogsSourceConfig config = new OTelLogsSourceConfig();
+        assertNull(config.getMaxConnectionAge());
+        assertNull(config.getConnectionDrainDuration());
+    }
+
+    @Test
+    void maxConnectionAge_deserializesFromYaml() {
+        final ObjectMapper mapper = new ObjectMapper().registerModule(new JavaTimeModule());
+        final Map<String, Object> settings = new HashMap<>();
+        settings.put("max_connection_age", "PT45M");
+        settings.put("connection_drain_duration", "PT10S");
+        final OTelLogsSourceConfig config = mapper.convertValue(settings, OTelLogsSourceConfig.class);
+        assertThat(config.getMaxConnectionAge(), equalTo(Duration.ofMinutes(45)));
+        assertThat(config.getConnectionDrainDuration(), equalTo(Duration.ofSeconds(10)));
+    }
+
     private PluginSetting completePluginSettingForOtelLogsSource(final int requestTimeoutInMillis,
                                                                  final int port,
                                                                  final String path,

data-prepper-plugins/otel-metrics-source/README.md

@@ -22,6 +22,8 @@ source:
 * unframed_requests(Optional) => A boolean to enable requests not framed using the gRPC wire protocol. When ```health_check_service``` is true and ```unframed_requests``` is true, enables HTTP health check service under ```/health```.
 * thread_count(Optional) => the number of threads to keep in the ScheduledThreadPool. Default is `200`.
 * max_connection_count(Optional) => the maximum allowed number of open connections. Default is `500`.
+* max_connection_age(Optional) => An ISO-8601 duration string (for example `PT30M`) that sets the maximum age of a gRPC server connection before it is gracefully closed. Useful when horizontally scaling, because long-lived gRPC channels do not re-resolve DNS on their own. When unset (default), Armeria places no upper bound on connection age. Allowed range: 1s to 24h.
+* connection_drain_duration(Optional) => An ISO-8601 duration string (for example `PT15S`) that controls how long Armeria waits after `max_connection_age` expires before forcefully terminating outstanding RPCs. When unset (default), Armeria uses its built-in drain duration. Allowed range: 0ms to 1h.
 * compression (Optional) : The compression type applied on the client request payload. Defaults to `none`. Supported values are:
     * `none`: no compression
     * `gzip`: apply GZip de-compression on the incoming request.

data-prepper-plugins/otel-metrics-source/build.gradle

@@ -28,6 +28,7 @@ dependencies {
     implementation libs.grpc.inprocess
     implementation 'com.fasterxml.jackson.core:jackson-databind'
     implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
+    implementation libs.hibernate.validator
     implementation libs.commons.lang3
     implementation libs.bouncycastle.bcprov
     implementation libs.bouncycastle.bcpkix

data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/ConvertConfiguration.java

@@ -21,6 +21,8 @@ public static ServerConfiguration convertConfiguration(final OTelMetricsSourceCo
             serverConfiguration.setRetryInfo(oTelMetricsSourceConfig.getRetryInfo());
             serverConfiguration.setThreadCount(oTelMetricsSourceConfig.getThreadCount());
             serverConfiguration.setMaxConnectionCount(oTelMetricsSourceConfig.getMaxConnectionCount());
+            serverConfiguration.setMaxConnectionAge(oTelMetricsSourceConfig.getMaxConnectionAge());
+            serverConfiguration.setConnectionDrainDuration(oTelMetricsSourceConfig.getConnectionDrainDuration());
 
         return serverConfiguration;
     }

data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSourceConfig.java

@@ -9,12 +9,15 @@
 import jakarta.validation.constraints.AssertTrue;
 import jakarta.validation.constraints.Size;
 import org.apache.commons.lang3.StringUtils;
+import org.hibernate.validator.constraints.time.DurationMax;
+import org.hibernate.validator.constraints.time.DurationMin;
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.codec.CompressionOption;
 import org.opensearch.dataprepper.plugins.server.RetryInfoConfig;
 import org.opensearch.dataprepper.model.configuration.PluginModel;
 import org.opensearch.dataprepper.plugins.otel.codec.OTelOutputFormat;
 
+import java.time.Duration;
 import java.util.Set;
 
 public class OTelMetricsSourceConfig {
@@ -126,6 +129,16 @@ public class OTelMetricsSourceConfig {
     @JsonProperty(RETRY_INFO)
     private RetryInfoConfig retryInfo;
 
+    @JsonProperty("max_connection_age")
+    @DurationMin(seconds = 1)
+    @DurationMax(hours = 24)
+    private Duration maxConnectionAge;
+
+    @JsonProperty("connection_drain_duration")
+    @DurationMin(millis = 0)
+    @DurationMax(hours = 1)
+    private Duration connectionDrainDuration;
+
     @AssertTrue(message = "buffer_partition_keys only supports 'name' and 'service_name'. 'name' is mandatory")
     boolean isBufferKeysValid() {
         if (bufferPartitionKeys == null) {
@@ -274,5 +287,13 @@ public RetryInfoConfig getRetryInfo() {
     public void setRetryInfo(RetryInfoConfig retryInfo) {
         this.retryInfo = retryInfo;
     }
+
+    public Duration getMaxConnectionAge() {
+        return maxConnectionAge;
+    }
+
+    public Duration getConnectionDrainDuration() {
+        return connectionDrainDuration;
+    }
 }
 

data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OtelMetricsSourceConfigTests.java

@@ -329,6 +329,35 @@ void testRetryInfoConfig() {
         assertThat(retryInfo.getMinDelay(), equalTo(Duration.ofMillis(50)));
     }
 
+    @Test
+    void maxConnectionAge_defaultsToNull() {
+        final OTelMetricsSourceConfig config = new OTelMetricsSourceConfig();
+        assertNull(config.getMaxConnectionAge());
+        assertNull(config.getConnectionDrainDuration());
+    }
+
+    @Test
+    void maxConnectionAge_deserializesFromYaml() {
+        final Map<String, Object> settings = new HashMap<>();
+        settings.put("max_connection_age", "PT60M");
+        settings.put("connection_drain_duration", "PT20S");
+        final OTelMetricsSourceConfig config = OBJECT_MAPPER.convertValue(settings, OTelMetricsSourceConfig.class);
+        assertThat(config.getMaxConnectionAge(), equalTo(Duration.ofMinutes(60)));
+        assertThat(config.getConnectionDrainDuration(), equalTo(Duration.ofSeconds(20)));
+    }
+
+    @Test
+    void convertConfiguration_propagatesKeepaliveSettings() {
+        final Map<String, Object> settings = new HashMap<>();
+        settings.put("max_connection_age", "PT12M");
+        settings.put("connection_drain_duration", "PT7S");
+        final OTelMetricsSourceConfig config = OBJECT_MAPPER.convertValue(settings, OTelMetricsSourceConfig.class);
+        final org.opensearch.dataprepper.plugins.server.ServerConfiguration server =
+                ConvertConfiguration.convertConfiguration(config);
+        assertThat(server.getMaxConnectionAge(), equalTo(Duration.ofMinutes(12)));
+        assertThat(server.getConnectionDrainDuration(), equalTo(Duration.ofSeconds(7)));
+    }
+
     private PluginSetting completePluginSettingForOtelMetricsSource(final int requestTimeoutInMillis,
                                                                     final int port,
                                                                     final String path,

data-prepper-plugins/otel-trace-source/README.md

@@ -36,6 +36,8 @@ For more information on migrating from Data Prepper 1.x to Data Prepper 2.x, see
 * unframed_requests(Optional) => A boolean to enable requests not framed using the gRPC wire protocol. When ```health_check_service``` is true and ```unframed_requests``` is true, enables HTTP health check service under ```/health```.
 * thread_count(Optional) => the number of threads to keep in the ScheduledThreadPool. Default is `200`.
 * max_connection_count(Optional) => the maximum allowed number of open connections. Default is `500`. 
+* max_connection_age(Optional) => An ISO-8601 duration string (for example `PT30M`) that sets the maximum age of a gRPC server connection before it is gracefully closed. Useful when horizontally scaling, because long-lived gRPC channels do not re-resolve DNS on their own. When unset (default), Armeria places no upper bound on connection age. Allowed range: 1s to 24h.
+* connection_drain_duration(Optional) => An ISO-8601 duration string (for example `PT15S`) that controls how long Armeria waits after `max_connection_age` expires before forcefully terminating outstanding RPCs. When unset (default), Armeria uses its built-in drain duration. Allowed range: 0ms to 1h.
 * authentication(Optional) => An authentication configuration. By default, this runs an unauthenticated server. See below for more information.
 * record_type(Optional) => A string represents the supported record data type that is written into the buffer plugin. Value options are `otlp` or `event`. Default is `otlp`.
 * output_format(Optional) => A string that sets the output format for decoded spans. Default is `opensearch`. Supported values are:

data-prepper-plugins/otel-trace-source/build.gradle

@@ -26,6 +26,7 @@ dependencies {
     implementation libs.grpc.inprocess
     implementation 'com.fasterxml.jackson.core:jackson-databind'
     implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
+    implementation libs.hibernate.validator
     implementation libs.commons.lang3
     implementation libs.bouncycastle.bcprov
     implementation libs.bouncycastle.bcpkix

data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceSource.java

@@ -155,6 +155,12 @@ private void configureHeadersAndHealthCheck(ServerBuilder serverBuilder) {
         if(oTelTraceSourceConfig.getMaxRequestLength() != null) {
             serverBuilder.maxRequestLength(oTelTraceSourceConfig.getMaxRequestLength().getBytes());
         }
+        if (oTelTraceSourceConfig.getMaxConnectionAge() != null) {
+            serverBuilder.maxConnectionAge(oTelTraceSourceConfig.getMaxConnectionAge());
+        }
+        if (oTelTraceSourceConfig.getConnectionDrainDuration() != null) {
+            serverBuilder.connectionDrainDuration(oTelTraceSourceConfig.getConnectionDrainDuration());
+        }
         serverBuilder.maxNumConnections(oTelTraceSourceConfig.getMaxConnectionCount());
     }
 

data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceSourceConfig.java

@@ -7,6 +7,8 @@
 
 import jakarta.validation.constraints.AssertTrue;
 import jakarta.validation.constraints.Size;
+import org.hibernate.validator.constraints.time.DurationMax;
+import org.hibernate.validator.constraints.time.DurationMin;
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.codec.CompressionOption;
 import org.opensearch.dataprepper.model.configuration.PluginModel;
@@ -15,6 +17,8 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.commons.lang3.StringUtils;
 
+import java.time.Duration;
+
 public class OTelTraceSourceConfig {
     static final String REQUEST_TIMEOUT = "request_timeout";
     static final String PORT = "port";
@@ -116,6 +120,16 @@ public class OTelTraceSourceConfig {
     @JsonProperty(RETRY_INFO)
     private RetryInfoConfig retryInfo;
 
+    @JsonProperty("max_connection_age")
+    @DurationMin(seconds = 1)
+    @DurationMax(hours = 24)
+    private Duration maxConnectionAge;
+
+    @JsonProperty("connection_drain_duration")
+    @DurationMin(millis = 0)
+    @DurationMax(hours = 1)
+    private Duration connectionDrainDuration;
+
     @AssertTrue(message = "path should start with /")
     boolean isPathValid() {
         return path == null || path.startsWith("/");
@@ -245,4 +259,12 @@ public ByteCount getMaxRequestLength() {
     public RetryInfoConfig getRetryInfo() {
         return retryInfo;
     }
+
+    public Duration getMaxConnectionAge() {
+        return maxConnectionAge;
+    }
+
+    public Duration getConnectionDrainDuration() {
+        return connectionDrainDuration;
+    }
 }