Addressed review comments by adding InternalEventHandle

Signed-off-by: Krishna Kondaka <[email protected]>
opensearch-project · Oct 31, 2023 · a9c3ebd · a9c3ebd
1 parent eeedfb8
commit a9c3ebd
Show file tree

Hide file tree

Showing 16 changed files with 124 additions and 64 deletions.
diff --git a/...-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java b/...-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java
@@ -11,7 +11,7 @@
 import java.time.Instant;
 import java.io.Serializable;
 
-public class DefaultEventHandle implements EventHandle, Serializable {
+public class DefaultEventHandle implements EventHandle, InternalEventHandle, Serializable {
     private Instant externalOriginationTime;
     private final Instant internalOriginationTime;
     private WeakReference<AcknowledgementSet> acknowledgementSetRef;
@@ -51,7 +51,6 @@ public Instant getExternalOriginationTime() {
 
     @Override
     public void release(boolean result) {
-        System.out.println("======release called==="+result);
         AcknowledgementSet acknowledgementSet = getAcknowledgementSet();
         if (acknowledgementSet != null) {
             acknowledgementSet.release(this, result);

diff --git a/...repper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java b/...repper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java
@@ -27,6 +27,8 @@ public class DefaultEventMetadata implements EventMetadata {
 
     private final Instant timeReceived;
 
+    private Instant externalOriginationTime;
+
     private Map<String, Object> attributes;
 
     private Set<String> tags;
@@ -43,13 +45,15 @@ private DefaultEventMetadata(final Builder builder) {
         this.attributes = builder.attributes == null ? new HashMap<>() : new HashMap<>(builder.attributes);
 
         this.tags = builder.tags == null ? new HashSet<>() : new HashSet(builder.tags);
+        this.externalOriginationTime = null;
     }
 
     private DefaultEventMetadata(final EventMetadata eventMetadata) {
         this.eventType = eventMetadata.getEventType();
         this.timeReceived = eventMetadata.getTimeReceived();
         this.attributes = new HashMap<>(eventMetadata.getAttributes());
         this.tags = new HashSet<>(eventMetadata.getTags());
+        this.externalOriginationTime = null;
     }
 
     @Override
@@ -62,6 +66,16 @@ public Instant getTimeReceived() {
         return timeReceived;
     }
 
+    @Override
+    public Instant getExternalOriginationTime() {
+        return externalOriginationTime;
+    }
+
+    @Override
+    public void setExternalOriginationTime(Instant externalOriginationTime) {
+        this.externalOriginationTime = externalOriginationTime;
+    }
+
     @Override
     public Map<String, Object> getAttributes() {
         return attributes;

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java
@@ -18,22 +18,6 @@ public interface EventHandle {
      */
     void release(boolean result);
 
-    /**
-     * sets acknowledgement set
-     *
-     * @param acknowledgementSet acknowledgementSet to be set in the event handle
-     * @since 2.6
-     */
-    void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet);
-
-    /**
-     * gets acknowledgement set
-     *
-     * @return returns acknowledgementSet from the event handle
-     * @since 2.6
-     */
-    AcknowledgementSet getAcknowledgementSet();
-
     /**
      * sets external origination time
      *

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java
@@ -31,6 +31,20 @@ public interface EventMetadata extends Serializable {
      */
     Instant getTimeReceived();
 
+    /**
+     * Returns the external origination time of the event
+     * @return the external origination time
+     * @since 2.6
+     */
+    Instant getExternalOriginationTime();
+
+    /**
+     * Sets the external origination time of the event
+     * @param externalOriginationTime the external origination time
+     * @since 2.6
+     */
+    void setExternalOriginationTime(Instant externalOriginationTime);
+
     /**
      * Returns the attributes
      * @return a map of attributes

diff --git a/...prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java b/...prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.event;
+
+import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
+import java.time.Instant;
+
+public interface InternalEventHandle {
+    /**
+     * sets acknowledgement set
+     *
+     * @param acknowledgementSet acknowledgementSet to be set in the event handle
+     * @since 2.6
+     */
+    void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet);
+
+    /**
+     * gets acknowledgement set
+     *
+     * @return returns acknowledgementSet from the event handle
+     * @since 2.6
+     */
+    AcknowledgementSet getAcknowledgementSet();
+
+}
+
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
@@ -101,11 +101,10 @@ protected JacksonEvent(final JacksonEvent otherEvent) {
     }
 
     public static Event fromMessage(String message) {
-        JacksonEvent event = JacksonEvent.builder()
+        return JacksonEvent.builder()
                 .withEventType(EVENT_TYPE)
                 .withData(Collections.singletonMap(MESSAGE_KEY, message))
                 .build();
-        return event;
     }
 
     private JsonNode getInitialJsonNode(final Object data) {

diff --git a/...er-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java b/...er-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java
@@ -24,6 +24,7 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.CoreMatchers.sameInstance;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.anEmptyMap;
@@ -80,6 +81,16 @@ public void testGetTimeReceived() {
         assertThat(timeReceived, is(equalTo(testTimeReceived)));
     }
 
+    @Test
+    public void testExternalOriginationTime() {
+        Instant externalOriginationTime = eventMetadata.getExternalOriginationTime();
+        assertThat(externalOriginationTime, is(nullValue()));
+        Instant now = Instant.now();
+        eventMetadata.setExternalOriginationTime(now);
+        externalOriginationTime = eventMetadata.getExternalOriginationTime();
+        assertThat(externalOriginationTime, is(equalTo(now)));
+    }
+
     @Test
     public void testGetAttributes() {
         final Map<String, Object> attributes = eventMetadata.getAttributes();

diff --git a/.../src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java b/.../src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java
@@ -6,6 +6,8 @@
 package org.opensearch.dataprepper.acknowledgements;
 
 import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 
 import java.util.concurrent.locks.ReentrantLock;
@@ -32,7 +34,12 @@ class AcknowledgementSetMonitor implements Runnable {
     private final AtomicInteger numNullHandles;
 
     private DefaultAcknowledgementSet getAcknowledgementSet(final EventHandle eventHandle) {
-        return (DefaultAcknowledgementSet)eventHandle.getAcknowledgementSet();
+        if (eventHandle instanceof DefaultEventHandle) {
+            InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+            return (DefaultAcknowledgementSet)internalEventHandle.getAcknowledgementSet();
+        } else {
+            throw new RuntimeException("Unsupported event handle");
+        }
     }
 
     public AcknowledgementSetMonitor() {

diff --git a/.../src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java b/.../src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java
@@ -8,6 +8,8 @@
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,8 +55,11 @@ public void add(Event event) {
         try {
             if (event instanceof JacksonEvent) {
                 EventHandle eventHandle = event.getEventHandle();
-                eventHandle.setAcknowledgementSet(this);
-                pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
+                if (eventHandle instanceof DefaultEventHandle) {
+                    InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+                    internalEventHandle.setAcknowledgementSet(this);
+                    pendingAcknowledgments.put(eventHandle, new AtomicInteger(1));
+                }
             }
         } finally {
             lock.unlock();

diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java
@@ -13,6 +13,8 @@
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
 import org.opensearch.dataprepper.pipeline.common.FutureHelper;
 import org.opensearch.dataprepper.pipeline.common.FutureHelperResult;
 import org.slf4j.Logger;
@@ -97,10 +99,15 @@ private void processAcknowledgements(List<Event> inputEvents, Collection outputR
         // For each event in the input events list that is not present in the output events, send positive acknowledgement, if acknowledgements are enabled for it
         inputEvents.forEach(event -> {
             EventHandle eventHandle = event.getEventHandle();
-            if (Objects.nonNull(eventHandle) && eventHandle.getAcknowledgementSet() != null && !outputEventsSet.contains(event)) {
-                eventHandle.release(true);
-            } else if (acknowledgementsEnabled && Objects.isNull(eventHandle)) {
-                invalidEventHandlesCounter.increment();
+            if (eventHandle != null && eventHandle instanceof DefaultEventHandle) {
+                InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle;
+                if (internalEventHandle.getAcknowledgementSet() != null && !outputEventsSet.contains(event)) {
+                    eventHandle.release(true);
+                } else if (acknowledgementsEnabled) {
+                    invalidEventHandlesCounter.increment();
+                }
+            } else if (eventHandle != null) {
+                throw new RuntimeException("Unexpected EventHandle");
             }
         });
     }

diff --git a/...re/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java b/...re/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java
@@ -65,7 +65,7 @@ private void acquireEventReference(final Record record) {
         }
         if (referencedRecords.contains(record) || ((routedRecords != null) && routedRecords.contains(record))) {
             EventHandle eventHandle = ((JacksonEvent)record.getData()).getEventHandle();
-            if (eventHandle != null && eventHandle.getAcknowledgementSet() != null) {
+            if (eventHandle != null && eventHandle instanceof DefaultEventHandle) {
                 acknowledgementSetManager.acquireEventReference(eventHandle);
             }
         } else if (!referencedRecords.contains(record)) {

diff --git a/...va/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java b/...va/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java
@@ -6,7 +6,7 @@
 package org.opensearch.dataprepper.acknowledgements;
 
 import org.opensearch.dataprepper.model.event.JacksonEvent;
-import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 import org.junit.jupiter.api.BeforeEach;
@@ -38,20 +38,20 @@ class DefaultAcknowledgementSetManagerTests {
     @Mock
     JacksonEvent event3;
 
-    EventHandle eventHandle1;
-    EventHandle eventHandle2;
-    EventHandle eventHandle3;
+    DefaultEventHandle eventHandle1;
+    DefaultEventHandle eventHandle2;
+    DefaultEventHandle eventHandle3;
     Boolean result;
 
     @BeforeEach
     void setup() {
         callbackExecutor = Executors.newFixedThreadPool(2);
         event1 = mock(JacksonEvent.class);
-        eventHandle1 = mock(EventHandle.class);
+        eventHandle1 = mock(DefaultEventHandle.class);
         lenient().when(event1.getEventHandle()).thenReturn(eventHandle1);
 
         event2 = mock(JacksonEvent.class);
-        eventHandle2 = mock(EventHandle.class);
+        eventHandle2 = mock(DefaultEventHandle.class);
         lenient().when(event2.getEventHandle()).thenReturn(eventHandle2);
 
         acknowledgementSetManager = createObjectUnderTest();
@@ -91,7 +91,7 @@ void testExpirations() throws InterruptedException {
     @Test
     void testMultipleAcknowledgementSets() {
         event3 = mock(JacksonEvent.class);
-        eventHandle3 = mock(EventHandle.class);
+        eventHandle3 = mock(DefaultEventHandle.class);
         lenient().when(event3.getEventHandle()).thenReturn(eventHandle3);
 
         AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, TEST_TIMEOUT);

diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java
@@ -14,7 +14,7 @@
 import org.opensearch.dataprepper.model.buffer.Buffer;
 import org.opensearch.dataprepper.model.configuration.PluginSetting;
 import org.opensearch.dataprepper.model.event.EventFactory;
-import org.opensearch.dataprepper.model.event.EventHandle;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 import org.opensearch.dataprepper.model.processor.Processor;
 import org.opensearch.dataprepper.model.record.Record;
@@ -84,7 +84,7 @@ class PipelineTests {
     private Duration peerForwarderDrainTimeout;
     private EventFactory eventFactory;
     private JacksonEvent event;
-    private EventHandle eventHandle;
+    private DefaultEventHandle eventHandle;
     private AcknowledgementSetManager acknowledgementSetManager;
 
     @BeforeEach
@@ -428,7 +428,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM
                 RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2);
                 Record rec = records.get(0);
                 event = mock(JacksonEvent.class);
-                eventHandle = mock(EventHandle.class);
+                eventHandle = mock(DefaultEventHandle.class);
                 acknowledgementSet = mock(AcknowledgementSet.class);
                 when(event.getEventHandle()).thenReturn(eventHandle);
                 when(eventHandle.getAcknowledgementSet()).thenReturn(acknowledgementSet);
@@ -441,7 +441,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM
             Pipeline pipeline = createObjectUnderTest();
             when(mockSource.areAcknowledgementsEnabled()).thenReturn(true);
             pipeline.publishToSinks(records);
-            verify(acknowledgementSetManager).acquireEventReference(any(EventHandle.class));
+            verify(acknowledgementSetManager).acquireEventReference(any(DefaultEventHandle.class));
 
             verify(router)
                     .route(anyCollection(), eq(dataFlowComponents), any(RouterGetRecordStrategy.class), any(BiConsumer.class));
@@ -455,7 +455,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_InactiveAcknowledge
                 RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2);
                 Record rec = records.get(0);
                 event = mock(JacksonEvent.class);
-                eventHandle = mock(EventHandle.class);
+                eventHandle = mock(DefaultEventHandle.class);
                 when(event.getEventHandle()).thenReturn(null);
                 when(rec.getData()).thenReturn(event);
                 routerCopyRecordStrategy.getRecord(rec);