Support otel metrics source with partition keys when persistent buffer is used (opensearch-project#6373)

* Support otel metrics source with partition keys when persistent buffer is used

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Addressed review comments

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Optimized  splitExportMetricsServiceRequestByKeys

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

---------

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

Author: kkondaka

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

@@ -27,6 +27,8 @@
 
 import java.time.Instant;
 import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class OTelMetricsGrpcService extends MetricsServiceGrpc.MetricsServiceImplBase {
@@ -49,15 +51,18 @@ public class OTelMetricsGrpcService extends MetricsServiceGrpc.MetricsServiceImp
     private final Counter recordsDroppedCounter;
     private final DistributionSummary payloadSizeSummary;
     private final Timer requestProcessDuration;
+    private final Set<String> bufferPartitionKeys;
 
     public OTelMetricsGrpcService(int bufferWriteTimeoutInMillis,
                                   final OTelProtoCodec.OTelProtoDecoder oTelProtoDecoder,
                                   Buffer<Record<? extends Metric>> buffer,
+                                  final Set<String> bufferPartitionKeys,
                                   final PluginMetrics pluginMetrics,
                                   final String metricsPrefix) {
         this.bufferWriteTimeoutInMillis = bufferWriteTimeoutInMillis;
         this.buffer = buffer;
 
+        this.bufferPartitionKeys = bufferPartitionKeys;
         if (metricsPrefix != null) {
             requestsReceivedCounter = pluginMetrics.counter(REQUESTS_RECEIVED, metricsPrefix);
             successRequestsCounter = pluginMetrics.counter(SUCCESS_REQUESTS, metricsPrefix);
@@ -96,7 +101,15 @@ public void export(final ExportMetricsServiceRequest request, final StreamObserv
     private void processRequest(final ExportMetricsServiceRequest request, final StreamObserver<ExportMetricsServiceResponse> responseObserver) {
         try {
             if (buffer.isByteBuffer()) {
-                buffer.writeBytes(request.toByteArray(), null, bufferWriteTimeoutInMillis);
+                if (bufferPartitionKeys != null && !bufferPartitionKeys.isEmpty()) {
+                    Map<String, ExportMetricsServiceRequest> requestsMap =
+                        oTelProtoDecoder.splitExportMetricsServiceRequestByKeys(request, bufferPartitionKeys);
+                    for (Map.Entry<String, ExportMetricsServiceRequest> entry: requestsMap.entrySet()) {
+                        buffer.writeBytes(entry.getValue().toByteArray(), entry.getKey(), bufferWriteTimeoutInMillis);
+                    }
+                } else {
+                    buffer.writeBytes(request.toByteArray(), null, bufferWriteTimeoutInMillis);
+                }
             } else {
                 Collection<Record<? extends Metric>> metrics;
 

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

@@ -85,6 +85,7 @@ public void start(Buffer<Record<? extends Metric>> buffer) {
                     (int) (oTelMetricsSourceConfig.getRequestTimeoutInMillis() * 0.8),
                     oTelMetricsSourceConfig.getOutputFormat() == OTelOutputFormat.OPENSEARCH ? new OTelProtoOpensearchCodec.OTelProtoDecoder() : new OTelProtoStandardCodec.OTelProtoDecoder(),
                     buffer,
+                    oTelMetricsSourceConfig.getBufferPartitionKeys(),
                     pluginMetrics,
                     null
             );

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

@@ -15,6 +15,8 @@
 import org.opensearch.dataprepper.model.configuration.PluginModel;
 import org.opensearch.dataprepper.plugins.otel.codec.OTelOutputFormat;
 
+import java.util.Set;
+
 public class OTelMetricsSourceConfig {
     static final String REQUEST_TIMEOUT = "request_timeout";
     static final String PORT = "port";
@@ -23,6 +25,8 @@ public class OTelMetricsSourceConfig {
     static final String USE_ACM_CERT_FOR_SSL = "useAcmCertForSSL";
     static final String ACM_CERT_ISSUE_TIME_OUT_MILLIS = "acmCertIssueTimeOutMillis";
     static final String HEALTH_CHECK_SERVICE = "health_check_service";
+    static final String OUTPUT_FORMAT = "output_format";
+    static final String BUFFER_PARTITION_KEYS = "partition_keys";
     static final String PROTO_REFLECTION_SERVICE = "proto_reflection_service";
     static final String SSL_KEY_CERT_FILE = "sslKeyCertChainFile";
     static final String SSL_KEY_FILE = "sslKeyFile";
@@ -46,6 +50,8 @@ public class OTelMetricsSourceConfig {
     static final int DEFAULT_ACM_CERT_ISSUE_TIME_OUT_MILLIS = 120000;
     private static final String S3_PREFIX = "s3://";
     static final String UNAUTHENTICATED_HEALTH_CHECK = "unauthenticated_health_check";
+    private static final String NAME_KEY = "name";
+    private static final String SERVICE_NAME_KEY = "service_name";
 
     @JsonProperty(REQUEST_TIMEOUT)
     private int requestTimeoutInMillis = DEFAULT_REQUEST_TIMEOUT_MS;
@@ -69,9 +75,13 @@ public class OTelMetricsSourceConfig {
     @JsonProperty(SSL)
     private boolean ssl = DEFAULT_SSL;
 
-    @JsonProperty("output_format")
+    @JsonProperty(OUTPUT_FORMAT)
     private OTelOutputFormat outputFormat = OTelOutputFormat.OPENSEARCH;
 
+    @JsonProperty(BUFFER_PARTITION_KEYS)
+    @Size(min = 1, max = 2, message = "partition_keys must contain 1 or 2 elements")
+    private Set<String> bufferPartitionKeys = null;
+
     @JsonProperty(USE_ACM_CERT_FOR_SSL)
     private boolean useAcmCertForSSL = DEFAULT_USE_ACM_CERT_FOR_SSL;
 
@@ -116,6 +126,17 @@ public class OTelMetricsSourceConfig {
     @JsonProperty(RETRY_INFO)
     private RetryInfoConfig retryInfo;
 
+    @AssertTrue(message = "buffer_partition_keys only supports 'name' and 'service_name'. 'name' is mandatory")
+    boolean isBufferKeysValid() {
+        if (bufferPartitionKeys == null) {
+            return true;
+        }
+
+        return bufferPartitionKeys.contains(NAME_KEY) &&
+           (bufferPartitionKeys.size() == 1 || (bufferPartitionKeys.size() == 2
+                && bufferPartitionKeys.contains(SERVICE_NAME_KEY)));
+    }
+
     @AssertTrue(message = "path should start with /")
     boolean isPathValid() {
         return path == null || path.startsWith("/");
@@ -168,6 +189,10 @@ public boolean hasHealthCheck() {
         return healthCheck;
     }
 
+    public Set<String> getBufferPartitionKeys() {
+        return bufferPartitionKeys;
+    }
+
     public boolean enableHttpHealthCheck() {
         return enableUnframedRequests() && hasHealthCheck();
     }

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

@@ -5,6 +5,7 @@
 
 package org.opensearch.dataprepper.plugins.source.otelmetrics;
 
+import org.apache.commons.lang3.tuple.Pair;
 import com.linecorp.armeria.server.ServiceRequestContext;
 import io.grpc.stub.StreamObserver;
 import io.micrometer.core.instrument.Counter;
@@ -44,6 +45,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Stream;
 
@@ -85,6 +87,10 @@ public class OTelMetricsGrpcServiceTest {
     private static final int TEST_RESOURCE_DROPPED_ATTRIBUTES_COUNT = 11;
     private static final int TEST_SCOPE_DROPPED_ATTRIBUTES_COUNT = 22;
 
+    private static final String SERVICE_NAME_KEY="service.name";
+    private static final String TEST_SERVICE_NAME1="testService1";
+    private static final String TEST_SERVICE_NAME2="testService2";
+
     private static final String IS_SCOPE_NAME="testISScopeName";
     private static final String IS_SCOPE_VERSION="testISScopeVersion";
     private static final String IS_SCOPE_ATTR_KEY="scope.attr";
@@ -167,6 +173,9 @@ public class OTelMetricsGrpcServiceTest {
     @Captor
     private ArgumentCaptor<Collection<Record>> recordCaptor;
 
+    @Captor
+    private ArgumentCaptor<byte[]> byteBufferRecordCaptor;
+
     @Captor
     ArgumentCaptor<byte[]> bytesCaptor;
 
@@ -195,8 +204,12 @@ public void setup() {
 
     }
 
+    OTelMetricsGrpcService createObjectUnderTestWithKeys(final OTelProtoCodec.OTelProtoDecoder decoder, final Set<String> keys) {
+        return new OTelMetricsGrpcService(bufferWriteTimeoutInMillis, decoder, buffer, keys, mockPluginMetrics, null);
+    }
+
     OTelMetricsGrpcService createObjectUnderTest(OTelProtoCodec.OTelProtoDecoder decoder) {
-        return new OTelMetricsGrpcService(bufferWriteTimeoutInMillis, decoder, buffer, mockPluginMetrics, null);
+        return new OTelMetricsGrpcService(bufferWriteTimeoutInMillis, decoder, buffer, null, mockPluginMetrics, null);
     }
 
     @ParameterizedTest
@@ -273,6 +286,43 @@ public void export_BufferTimeout_responseObserverOnError(OTelProtoCodec.OTelProt
         verify(requestProcessDuration, times(1)).record(ArgumentMatchers.<Runnable>any());
     }
 
+    @Test
+    public void test_MetricsSource_output_with_PartitionKeyName() throws Exception {
+        final ExportMetricsServiceRequest METRIC_REQUEST = createMetricsRequest();
+        final OTelProtoStandardCodec.OTelProtoDecoder decoder = new OTelProtoStandardCodec.OTelProtoDecoder();
+        when(buffer.isByteBuffer()).thenReturn(true);
+
+        sut = createObjectUnderTestWithKeys(decoder, Set.of("name"));
+        try (MockedStatic<ServiceRequestContext> mockedStatic = mockStatic(ServiceRequestContext.class)) {
+            mockedStatic.when(ServiceRequestContext::current).thenReturn(serviceRequestContext);
+            sut.export(METRIC_REQUEST, responseObserver);
+        }
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_GAUGE_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SUM_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SUMMARY_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_HISTOGRAM_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_EXPONENTIAL_HISTOGRAM_METRIC_NAME), anyInt());
+    }
+
+
+    @Test
+    public void test_MetricsSource_output_with_PartitionKeysNameAndServiceName() throws Exception {
+        final ExportMetricsServiceRequest METRIC_REQUEST = createMetricsRequestWithTwoServices();
+        final OTelProtoStandardCodec.OTelProtoDecoder decoder = new OTelProtoStandardCodec.OTelProtoDecoder();
+        when(buffer.isByteBuffer()).thenReturn(true);
+
+        sut = createObjectUnderTestWithKeys(decoder, Set.of("name", "service_name"));
+        try (MockedStatic<ServiceRequestContext> mockedStatic = mockStatic(ServiceRequestContext.class)) {
+            mockedStatic.when(ServiceRequestContext::current).thenReturn(serviceRequestContext);
+            sut.export(METRIC_REQUEST, responseObserver);
+        }
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SERVICE_NAME1+":"+TEST_GAUGE_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SERVICE_NAME1+":"+TEST_SUM_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SERVICE_NAME2+":"+TEST_SUMMARY_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SERVICE_NAME2+":"+TEST_HISTOGRAM_METRIC_NAME), anyInt());
+        verify(buffer, times(1)).writeBytes(byteBufferRecordCaptor.capture(), eq(TEST_SERVICE_NAME2+":"+TEST_EXPONENTIAL_HISTOGRAM_METRIC_NAME), anyInt());
+    }
+
     @Test
     public void test_MetricsSource_output_with_OpensearchFormat() throws Exception {
         final ExportMetricsServiceRequest METRIC_REQUEST = createMetricsRequest();
@@ -336,19 +386,12 @@ private static Stream<Arguments> getDecoderArguments() {
         );
     }
 
-    private ExportMetricsServiceRequest createMetricsRequest() {
+    private Pair<ScopeMetrics, ScopeMetrics> createScopeMetrics() {
         Instant currentTime = Instant.now();
         startTime = currentTime.toString();
         final long currentUnixTimeNano = ((long)currentTime.getEpochSecond() * 1000_000_000L) + currentTime.getNano();
         final long endUnixTimeNano = currentUnixTimeNano + (TIME_DELTA*1000_000_000L);
         endTime = currentTime.plusSeconds(TIME_DELTA).toString();
-        final Resource resource = Resource.newBuilder()
-                .setDroppedAttributesCount(TEST_RESOURCE_DROPPED_ATTRIBUTES_COUNT)
-                .addAttributes(KeyValue.newBuilder()
-                        .setKey(TEST_RESOURCE_ATTR_KEY)
-                        .setValue(AnyValue.newBuilder().setStringValue(TEST_RESOURCE_ATTR_VALUE).build())
-                ).build();
-
         final InstrumentationScope instrumentationScope = InstrumentationScope.newBuilder()
                 .setName(IS_SCOPE_NAME)
                 .setVersion(IS_SCOPE_VERSION)
@@ -357,7 +400,7 @@ private ExportMetricsServiceRequest createMetricsRequest() {
                         .setKey(IS_SCOPE_ATTR_KEY)
                         .setValue(AnyValue.newBuilder().setStringValue(IS_SCOPE_ATTR_VALUE).build())
                 ).build();
-        ScopeMetrics scopeMetrics = ScopeMetrics.newBuilder()
+        ScopeMetrics scopeMetrics1 = ScopeMetrics.newBuilder()
                 .setScope(instrumentationScope)
                 .setSchemaUrl(TEST_SCHEMA_URL)
                 .addMetrics(io.opentelemetry.proto.metrics.v1.Metric.newBuilder()
@@ -395,6 +438,11 @@ private ExportMetricsServiceRequest createMetricsRequest() {
                         .setAggregationTemporality(io.opentelemetry.proto.metrics.v1.AggregationTemporality.AGGREGATION_TEMPORALITY_CUMULATIVE)
                         .build())
                     .build())
+                .build();
+
+        ScopeMetrics scopeMetrics2 = ScopeMetrics.newBuilder()
+                .setScope(instrumentationScope)
+                .setSchemaUrl(TEST_SCHEMA_URL)
                 .addMetrics(io.opentelemetry.proto.metrics.v1.Metric.newBuilder()
                     .setName(TEST_EXPONENTIAL_HISTOGRAM_METRIC_NAME)
                     .setDescription(TEST_EXPONENTIAL_HISTOGRAM_DESCRIPTION)
@@ -473,14 +521,63 @@ private ExportMetricsServiceRequest createMetricsRequest() {
                         .build())
                     .build())
                 .build();
+        return Pair.of(scopeMetrics1, scopeMetrics2);
+    }
 
+    private ExportMetricsServiceRequest createMetricsRequest() {
+        final Resource resource = Resource.newBuilder()
+                .setDroppedAttributesCount(TEST_RESOURCE_DROPPED_ATTRIBUTES_COUNT)
+                .addAttributes(KeyValue.newBuilder()
+                        .setKey(TEST_RESOURCE_ATTR_KEY)
+                        .setValue(AnyValue.newBuilder().setStringValue(TEST_RESOURCE_ATTR_VALUE).build())
+                ).build();
+
+        Pair<ScopeMetrics, ScopeMetrics> scopeMetrics = createScopeMetrics();
         ResourceMetrics resourceMetrics = ResourceMetrics.newBuilder()
                 .setResource(resource)
-                .addScopeMetrics(scopeMetrics)
+                .addScopeMetrics(scopeMetrics.getLeft())
+                .addScopeMetrics(scopeMetrics.getRight())
                 .build();
 
         return ExportMetricsServiceRequest.newBuilder()
                 .addResourceMetrics(resourceMetrics)
                 .build();
     }
+
+    private ExportMetricsServiceRequest createMetricsRequestWithTwoServices() {
+        Instant currentTime = Instant.now();
+        startTime = currentTime.toString();
+        final long currentUnixTimeNano = ((long)currentTime.getEpochSecond() * 1000_000_000L) + currentTime.getNano();
+        final long endUnixTimeNano = currentUnixTimeNano + (TIME_DELTA*1000_000_000L);
+        endTime = currentTime.plusSeconds(TIME_DELTA).toString();
+        final Resource resource1 = Resource.newBuilder()
+                .setDroppedAttributesCount(TEST_RESOURCE_DROPPED_ATTRIBUTES_COUNT)
+                .addAttributes(KeyValue.newBuilder()
+                        .setKey(SERVICE_NAME_KEY)
+                        .setValue(AnyValue.newBuilder().setStringValue(TEST_SERVICE_NAME1).build())
+                ).build();
+        final Resource resource2 = Resource.newBuilder()
+                .setDroppedAttributesCount(TEST_RESOURCE_DROPPED_ATTRIBUTES_COUNT)
+                .addAttributes(KeyValue.newBuilder()
+                        .setKey(SERVICE_NAME_KEY)
+                        .setValue(AnyValue.newBuilder().setStringValue(TEST_SERVICE_NAME2).build())
+                ).build();
+
+        Pair<ScopeMetrics, ScopeMetrics> scopeMetrics = createScopeMetrics();
+
+        ResourceMetrics resourceMetrics1 = ResourceMetrics.newBuilder()
+                .setResource(resource1)
+                .addScopeMetrics(scopeMetrics.getLeft())
+                .build();
+
+        ResourceMetrics resourceMetrics2 = ResourceMetrics.newBuilder()
+                .setResource(resource2)
+                .addScopeMetrics(scopeMetrics.getRight())
+                .build();
+
+        return ExportMetricsServiceRequest.newBuilder()
+                .addResourceMetrics(resourceMetrics1)
+                .addResourceMetrics(resourceMetrics2)
+                .build();
+    }
 }