Addressed review comments and fixed acknowledgements issue

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

Author: kkondaka

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkContext.java

@@ -8,6 +8,7 @@
 import org.opensearch.dataprepper.model.pipeline.HeadlessPipeline;
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.event.InternalEventHandle;
 import org.opensearch.dataprepper.model.event.EventMetadata;
 
 import java.util.Collection;
@@ -63,7 +64,20 @@ public void setForwardToPipelines(final Map<String, HeadlessPipeline> pipelines)
         }
     }
 
-    public boolean forwardRecords(final Collection<Record<Event>> records, final Map<String, Object> withData, final Map<String, Object> withMetadata) {
+    public SinkForwardRecordsContext prepareRecordsForForwarding(final Collection<Record<Event>> records) {
+        for (Map.Entry<String, HeadlessPipeline> entry: forwardToPipelines.entrySet()) {
+            records.forEach((record) -> {
+                InternalEventHandle eventHandle = (InternalEventHandle)record.getData().getEventHandle();
+                if (eventHandle != null) {
+                    eventHandle.acquireReference();
+                }
+            });
+        }
+        return new SinkForwardRecordsContext();
+
+    }
+
+    public boolean forwardRecords(final SinkForwardRecordsContext sinkForwardRecordsContext, final Collection<Record<Event>> records, final Map<String, Object> withData, final Map<String, Object> withMetadata) {
         if (forwardToPipelines.size() == 0) {
             return false;
         }

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkForwardRecordsContext.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.sink;
+
+public class SinkForwardRecordsContext {
+    public SinkForwardRecordsContext() {}
+}
+

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkContextTest.java

@@ -10,13 +10,15 @@
 
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.event.EventMetadata;
 import org.opensearch.dataprepper.model.pipeline.HeadlessPipeline;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
@@ -78,23 +80,27 @@ public void testForwardToPipelinesWithPipelineMap() {
         Record<Event> record = mock(Record.class);
         EventMetadata eventMetadata = mock(EventMetadata.class);
         Event event = mock(Event.class);
+        DefaultEventHandle eventHandle = mock(DefaultEventHandle.class);
+        doNothing().when(eventHandle).acquireReference();
+        when(event.getEventHandle()).thenReturn(eventHandle);
         when(record.getData()).thenReturn(event);
         when(event.getMetadata()).thenReturn(eventMetadata);
         Collection<Record<Event>> records = Collections.singletonList(record);
+        SinkForwardRecordsContext sinkForwardRecordsContext = sinkContext.prepareRecordsForForwarding(records);
 
-        assertThat(sinkContext.forwardRecords(records, Map.of("datakey1", "datavalue1"), Map.of("metadataKey1", "metadataValue1")), equalTo(true));
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, records, Map.of("datakey1", "datavalue1"), Map.of("metadataKey1", "metadataValue1")), equalTo(true));
         verify(forwardPipeline1, times(1)).sendEvents(eq(records));
         verify(forwardPipeline2, times(1)).sendEvents(eq(records));
         verify(event, times(1)).put(any(String.class), any(Object.class));
         verify(event, times(1)).getMetadata();
         verify(eventMetadata, times(1)).setAttribute(any(String.class), any(Object.class));
-        assertThat(sinkContext.forwardRecords(records, null, null), equalTo(true));
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, records, null, null), equalTo(true));
         verify(forwardPipeline1, times(2)).sendEvents(eq(records));
         verify(forwardPipeline2, times(2)).sendEvents(eq(records));
-        assertThat(sinkContext.forwardRecords(records, Map.of(), Map.of()), equalTo(true));
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, records, Map.of(), Map.of()), equalTo(true));
         verify(forwardPipeline1, times(3)).sendEvents(eq(records));
         verify(forwardPipeline2, times(3)).sendEvents(eq(records));
-        assertThat(sinkContext.forwardRecords(null, Map.of(), Map.of()), equalTo(true));
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, null, Map.of(), Map.of()), equalTo(true));
         verify(forwardPipeline1, times(3)).sendEvents(eq(records));
         verify(forwardPipeline2, times(3)).sendEvents(eq(records));
     }
@@ -113,7 +119,8 @@ public void testForwardToPipelinesWithPipelineMapAndFailureCases() {
         HeadlessPipeline forwardPipeline1 = mock(HeadlessPipeline.class);
         assertThrows(RuntimeException.class, () -> sinkContext.setForwardToPipelines(Map.of("forward-pipeline1", forwardPipeline1)));
         Collection<Record<Event>> records = mock(Collection.class);
-        assertThat(sinkContext.forwardRecords(records, Map.of(), Map.of()), equalTo(false));
+        SinkForwardRecordsContext sinkForwardRecordsContext = sinkContext.prepareRecordsForForwarding(records);
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, records, Map.of(), Map.of()), equalTo(false));
     }
 
     @Test
@@ -125,7 +132,8 @@ public void testWithNoForwardToPipelines() {
         final List<String> testForwardToPipelineNames = Collections.emptyList();
         sinkContext = new SinkContext(testTagsTargetKey, testRoutes, testIncludeKeys, testExcludeKeys, testForwardToPipelineNames);
         Collection<Record<Event>> records = mock(Collection.class);
-        assertThat(sinkContext.forwardRecords(records, Map.of(), Map.of()), equalTo(false));
+        SinkForwardRecordsContext sinkForwardRecordsContext = sinkContext.prepareRecordsForForwarding(records);
+        assertThat(sinkContext.forwardRecords(sinkForwardRecordsContext, records, Map.of(), Map.of()), equalTo(false));
     }
 }
 

data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/HeadlessPipelinesIT.java

@@ -78,6 +78,7 @@ void dlq_pipeline_test() {
 
         final List<Record<Event>> sinkRecords = inMemorySinkAccessor.get(IN_MEMORY_IDENTIFIER_DLQ);
 
+        assertThat(sinkRecords.size(), equalTo(recordsToCreate));
         for (int i = 0; i < sinkRecords.size(); i++) {
             final Record<Event> inputRecord = inputRecords.get(i);
             final Record<Event> sinkRecord = sinkRecords.get(i);

data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/sink/FileSink.java

@@ -68,28 +68,23 @@ public FileSink(final FileSinkConfig fileSinkConfig, final SinkContext sinkConte
 
     @Override
     public void output(final Collection<Record<Object>> records) {
-        final boolean doForward = sinkContext.getForwardToPipelines().size() > 0;
+        sinkContext.prepareForForwardPipelines(records);
         lock.lock();
         Collection<Record<Event>> events = new ArrayList<>();
+
         try {
             if (isStopRequested)
                 return;
 
             for (final Record<Object> record : records) {
                 try {
                     checkTypeAndWriteObject(record.getData(), writer);
-                    if (doForward && record.getData() instanceof Event) {
-                        Event event = (Event)record.getData();
-                        events.add(new Record<>(event));
-                    }
                 } catch (final IOException ex) {
                     throw new RuntimeException(format("Encountered exception writing to file %s", outputFilePath), ex);
                 }
             }
 
-            if (doForward) {
-                sinkContext.forwardRecords(events, null, null);
-            }
+            sinkContext.forwardRecords(events, null, null);
             try {
                 writer.flush();
             } catch (final IOException ex) {