Fix pipeline latency to calculate correct latency when persistent buf…

…fer is used (#4187) * Fix pipeline latency to calculate correct latency when persistent buffer is used Signed-off-by: Krishna Kondaka <[email protected]> * Fixed checkstyle error and addressed comments Signed-off-by: Krishna Kondaka <[email protected]> * Fixed failing tests Signed-off-by: Krishna Kondaka <[email protected]> * Fixed failing tests Signed-off-by: Krishna Kondaka <[email protected]> * Fixed failing tests Signed-off-by: Krishna Kondaka <[email protected]> --------- Signed-off-by: Krishna Kondaka <[email protected]> Co-authored-by: Krishna Kondaka <[email protected]>
opensearch-project · Feb 29, 2024 · 42e763e · 42e763e
1 parent 6255784
commit 42e763e
Show file tree

Hide file tree

Showing 35 changed files with 318 additions and 73 deletions.
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java
@@ -11,16 +11,18 @@
 import java.io.InputStream;
 import java.io.Serializable;
 import java.util.function.Consumer;
+import java.time.Instant;
 
 public interface ByteDecoder extends Serializable {
     /**
      * Parses an {@link InputStream}. Implementors should call the {@link Consumer} for each
      * {@link Record} loaded from the {@link InputStream}.
      *
      * @param inputStream   The input stream for code to process
+     * @param timeReceived  The time received value to be populated in the Record
      * @param eventConsumer The consumer which handles each event from the stream
      * @throws IOException throws IOException when invalid input is received or incorrect codec name is provided
      */
-    void parse(InputStream inputStream, Consumer<Record<Event>> eventConsumer) throws IOException;
+    void parse(InputStream inputStream, Instant timeReceived, Consumer<Record<Event>> eventConsumer) throws IOException;
 
 }
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java
@@ -16,6 +16,7 @@
 import org.opensearch.dataprepper.model.record.Record;
 import java.io.IOException;
 import java.io.InputStream;
+import java.time.Instant;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
@@ -24,33 +25,36 @@ public class JsonDecoder implements ByteDecoder {
     private final ObjectMapper objectMapper = new ObjectMapper();
     private final JsonFactory jsonFactory = new JsonFactory();
 
-    public void parse(InputStream inputStream, Consumer<Record<Event>> eventConsumer) throws IOException {
+    public void parse(InputStream inputStream, Instant timeReceived, Consumer<Record<Event>> eventConsumer) throws IOException {
         Objects.requireNonNull(inputStream);
         Objects.requireNonNull(eventConsumer);
 
         final JsonParser jsonParser = jsonFactory.createParser(inputStream);
 
         while (!jsonParser.isClosed() && jsonParser.nextToken() != JsonToken.END_OBJECT) {
             if (jsonParser.getCurrentToken() == JsonToken.START_ARRAY) {
-                parseRecordsArray(jsonParser, eventConsumer);
+                parseRecordsArray(jsonParser, timeReceived, eventConsumer);
             }
         }
     }
 
-    private void parseRecordsArray(final JsonParser jsonParser, final Consumer<Record<Event>> eventConsumer) throws IOException {
+    private void parseRecordsArray(final JsonParser jsonParser, final Instant timeReceived, final Consumer<Record<Event>> eventConsumer) throws IOException {
         while (jsonParser.nextToken() != JsonToken.END_ARRAY) {
             final Map<String, Object> innerJson = objectMapper.readValue(jsonParser, Map.class);
 
-            final Record<Event> record = createRecord(innerJson);
+            final Record<Event> record = createRecord(innerJson, timeReceived);
             eventConsumer.accept(record);
         }
     }
 
-    private Record<Event> createRecord(final Map<String, Object> json) {
-        final JacksonEvent event = (JacksonEvent)JacksonLog.builder()
+    private Record<Event> createRecord(final Map<String, Object> json, final Instant timeReceived) {
+        final JacksonLog.Builder logBuilder = JacksonLog.builder()
                 .withData(json)
-                .getThis()
-                .build();
+                .getThis();
+        if (timeReceived != null) {
+            logBuilder.withTimeReceived(timeReceived);
+        }
+        final JacksonEvent event = (JacksonEvent)logBuilder.build();
 
         return new Record<>(event);
     }

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java
@@ -10,6 +10,7 @@
 import org.opensearch.dataprepper.model.event.EventType;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
@@ -161,6 +162,18 @@ public Builder withAttributes(final Map<String, Object> attributes) {
             return getThis();
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public Builder withTimeReceived(final Instant timeReceived) {
+            return (Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Sets the observed time of the log event
          *

diff --git a/...pi/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java b/...pi/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java
@@ -7,6 +7,7 @@
 
 import org.opensearch.dataprepper.model.event.EventType;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -237,6 +238,18 @@ public JacksonExponentialHistogram.Builder withPositiveOffset(int offset) {
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public JacksonExponentialHistogram.Builder withTimeReceived(final Instant timeReceived) {
+            return (JacksonExponentialHistogram.Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Sets the offset for the negative buckets
          *

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java
@@ -7,6 +7,7 @@
 
 import org.opensearch.dataprepper.model.event.EventType;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -67,6 +68,18 @@ public Builder withValue(final Double value) {
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public Builder withTimeReceived(final Instant timeReceived) {
+            return (Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Returns a newly created {@link JacksonGauge}
          * @return a JacksonGauge

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java
@@ -7,6 +7,7 @@
 
 import org.opensearch.dataprepper.model.event.EventType;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -190,6 +191,18 @@ public  JacksonHistogram.Builder withAggregationTemporality(String aggregationTe
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public JacksonHistogram.Builder withTimeReceived(final Instant timeReceived) {
+            return (JacksonHistogram.Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Sets the buckets for this histogram
          * @param buckets a list of buckets

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java
@@ -9,6 +9,7 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -234,6 +235,19 @@ public T withSchemaUrl(final String schemaUrl) {
             return getThis();
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public T withTimeReceived(final Instant timeReceived) {
+            return (T)super.withTimeReceived(timeReceived);
+        }
+
+
         /**
          * Sets the exemplars that are associated with this metric event
          * @param  exemplars sets the exemplars for this metric

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java
@@ -7,6 +7,7 @@
 
 import org.opensearch.dataprepper.model.event.EventType;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -101,6 +102,18 @@ public Builder withIsMonotonic(final boolean isMonotonic) {
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public Builder withTimeReceived(final Instant timeReceived) {
+            return (Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Returns a newly created {@link JacksonSum}
          * @return a JacksonSum

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java
@@ -7,6 +7,7 @@
 
 import org.opensearch.dataprepper.model.event.EventType;
 
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -106,6 +107,18 @@ public Builder withSum(double sum) {
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public Builder withTimeReceived(final Instant timeReceived) {
+            return (Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Sets the count
          * @param count the count of this summary

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java
@@ -15,6 +15,7 @@
 import org.opensearch.dataprepper.model.event.EventType;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -444,6 +445,18 @@ public Builder withTraceGroup(final String traceGroup) {
             return this;
         }
 
+        /**
+         * Sets the time received for populating event origination time in event handle
+         *
+         * @param timeReceived time received
+         * @return the builder
+         * @since 2.7
+         */
+        @Override
+        public Builder withTimeReceived(final Instant timeReceived) {
+            return (Builder)super.withTimeReceived(timeReceived);
+        }
+
         /**
          * Sets the duration of the span
          *

diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java
@@ -2,9 +2,11 @@
 
 import org.junit.jupiter.api.Test;
 import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.record.Record;
 
 import java.io.ByteArrayInputStream;
+import java.time.Instant;
 import java.util.Map;
 import java.util.Random;
 import java.util.UUID;
@@ -18,12 +20,13 @@
 public class JsonDecoderTest {
     private JsonDecoder jsonDecoder;
     private Record<Event> receivedRecord;
+    private Instant receivedTime;
 
     private JsonDecoder createObjectUnderTest() {
         return new JsonDecoder();
     }
 
-    @BeforeEach
+@BeforeEach
     void setup() {
         jsonDecoder = createObjectUnderTest();
         receivedRecord = null;
@@ -36,7 +39,7 @@ void test_basicJsonDecoder() {
         int intValue = r.nextInt();
         String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]";
         try {
-            jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), (record) -> {
+            jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), null, (record) -> {
                 receivedRecord = record;
             });
         } catch (Exception e){}
@@ -47,4 +50,25 @@ void test_basicJsonDecoder() {
         assertThat(map.get("key2"), equalTo(intValue));
     }
 
+    @Test
+    void test_basicJsonDecoder_withTimeReceived() {
+        String stringValue = UUID.randomUUID().toString();
+        Random r = new Random();
+        int intValue = r.nextInt();
+        String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]";
+        final Instant now = Instant.now();
+        try {
+            jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), now, (record) -> {
+                receivedRecord = record;
+                receivedTime = ((DefaultEventHandle)(((Event)record.getData()).getEventHandle())).getInternalOriginationTime();
+            });
+        } catch (Exception e){}
+
+        assertNotEquals(receivedRecord, null);
+        Map<String, Object> map = receivedRecord.getData().toMap();
+        assertThat(map.get("key1"), equalTo(stringValue));
+        assertThat(map.get("key2"), equalTo(intValue));
+        assertThat(receivedTime, equalTo(now));
+    }
+
 }
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java
@@ -12,7 +12,9 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.skyscreamer.jsonassert.JSONAssert;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 
+import java.time.Instant;
 import java.util.Date;
 import java.util.Map;
 import java.util.UUID;
@@ -83,6 +85,14 @@ public void testGetServiceName() {
         assertThat(name, is(equalTo(TEST_SERVICE_NAME)));
     }
 
+    @Test
+    public void testGetTimeReceived() {
+        Instant now = Instant.now();
+        builder.withTimeReceived(now);
+        log = builder.build();
+        assertThat(((DefaultEventHandle)log.getEventHandle()).getInternalOriginationTime(), is(now));
+    }
+
     @Test
     public void testGetSchemaUrl() {
         final String schemaUrl = log.getSchemaUrl();