Fix Data Prepper router to send records through routing strategy before sending to the sinks (opensearch-project#6370)

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

Author: kkondaka

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

@@ -9,31 +9,44 @@
 import org.junit.jupiter.api.BeforeEach;
 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.event.JacksonEvent;
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.plugins.InMemorySinkAccessor;
 import org.opensearch.dataprepper.plugins.InMemorySourceAccessor;
 import org.opensearch.dataprepper.test.framework.DataPrepperTestRunner;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.ArgumentMatchers.any;
 import static org.awaitility.Awaitility.await;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 
 class Router_ThreeRoutesDefaultIT {
     private static final String TESTING_KEY = "ConditionalRoutingIT";
     private static final String ALL_SOURCE_KEY = TESTING_KEY + "_all";
+    private static final int NUM_ALPHA_EVENTS = 10;
+    private static final int NUM_BETA_EVENTS = 20;
+    private static final int NUM_GAMMA_EVENTS = 20;
+    private static final int NUM_DEFAULT_EVENTS = 20;
     private static final String ALPHA_SOURCE_KEY = TESTING_KEY + "_alpha";
     private static final String BETA_SOURCE_KEY = TESTING_KEY + "_beta";
     private static final String ALPHA_DEFAULT_SOURCE_KEY = TESTING_KEY + "_alpha_default";
@@ -47,9 +60,43 @@ class Router_ThreeRoutesDefaultIT {
     private DataPrepperTestRunner dataPrepperTestRunner;
     private InMemorySourceAccessor inMemorySourceAccessor;
     private InMemorySinkAccessor inMemorySinkAccessor;
+    private Map<Object, AtomicInteger> numReleases;
+    private Map<String, DefaultEventHandle> eventHandles;
 
     @BeforeEach
     void setUp() {
+        eventHandles = new HashMap<>();
+        Map<String, Integer> testEventTypesMap = Map.of(ALPHA_VALUE, NUM_ALPHA_EVENTS, BETA_VALUE, NUM_BETA_EVENTS, GAMMA_VALUE, NUM_GAMMA_EVENTS, DEFAULT_VALUE, NUM_DEFAULT_EVENTS);
+        // Setup check on event handles such a way that all acquires happen before release
+        for (Map.Entry<String, Integer> testEventType: testEventTypesMap.entrySet()) {
+            for (int i = 0; i < testEventType.getValue(); i++) {
+                DefaultEventHandle eventHandle = mock(DefaultEventHandle.class);
+                eventHandles.put(testEventType.getKey()+i, eventHandle);
+                when(eventHandle.getInternalOriginationTime()).thenReturn(Instant.now());
+                when(eventHandle.getExternalOriginationTime()).thenReturn(Instant.now());
+                doAnswer(invocation -> {
+                    Object mock = invocation.getMock();
+                    synchronized (numReleases) {
+                        AtomicInteger releases = numReleases.computeIfAbsent(mock, k -> new AtomicInteger(0));
+                        assertThat(releases.get(), equalTo(0));
+                    }
+                    return null;
+                }).when(eventHandle).acquireReference();
+
+                doAnswer(invocation -> {
+                    Object mock = invocation.getMock();
+                    synchronized (numReleases) {
+                        AtomicInteger releases = numReleases.get(mock);
+                        if (releases != null) {
+                            releases.incrementAndGet();
+                        }
+                    }
+                    return null;
+                }).when(eventHandle).release(any(Boolean.class));
+            }
+        }
+        numReleases = new HashMap<>();
+
         dataPrepperTestRunner = DataPrepperTestRunner.builder()
                 .withPipelinesDirectoryOrFile("route/three-route-with-default-route.yaml")
                 .build();
@@ -66,10 +113,10 @@ void tearDown() {
 
     @Test
     void test_default_route() {
-        final List<Record<Event>> alphaEvents = createEvents(ALPHA_VALUE, 10);
-        final List<Record<Event>> betaEvents = createEvents(BETA_VALUE, 20);
-        final List<Record<Event>> gammaEvents = createEvents(GAMMA_VALUE, 20);
-        final List<Record<Event>> defaultEvents = createEvents(DEFAULT_VALUE, 20);
+        final List<Record<Event>> alphaEvents = createEvents(ALPHA_VALUE, NUM_ALPHA_EVENTS);
+        final List<Record<Event>> betaEvents = createEvents(BETA_VALUE, NUM_BETA_EVENTS);
+        final List<Record<Event>> gammaEvents = createEvents(GAMMA_VALUE, NUM_GAMMA_EVENTS);
+        final List<Record<Event>> defaultEvents = createEvents(DEFAULT_VALUE, NUM_DEFAULT_EVENTS);
 
         final List<Record<Event>> allEvents = new ArrayList<>(alphaEvents);
         allEvents.addAll(betaEvents);
@@ -120,8 +167,12 @@ void test_default_route() {
 
     private List<Record<Event>> createEvents(final String value, final int numberToCreate) {
         return IntStream.range(0, numberToCreate)
-                .mapToObj(i -> Map.of(KNOWN_CONDITIONAL_KEY, value, "arbitrary_field", UUID.randomUUID().toString()))
-                .map(map -> JacksonEvent.builder().withData(map).withEventType("TEST").build())
+                .mapToObj(i -> Map.of(KNOWN_CONDITIONAL_KEY, value, "id", i, "arbitrary_field", UUID.randomUUID().toString()))
+                .map(map -> {
+                    DefaultEventHandle eventHandle = eventHandles.get(value+map.get("id"));
+                    assertNotNull(eventHandle);
+                    return JacksonEvent.builder().withData(map).withEventType("TEST").withEventHandle(eventHandle).build();
+                })
                 .map(jacksonEvent -> (Event) jacksonEvent)
                 .map(Record::new)
                 .collect(Collectors.toList());

data-prepper-core/src/integrationTest/resources/org/opensearch/dataprepper/pipeline/route/three-route-with-default-route.yaml

@@ -4,6 +4,7 @@ routing-pipeline:
   source:
     in_memory:
       testing_key: ConditionalRoutingIT
+      acknowledgments: true
   buffer:
     bounded_blocking:
       # Use a small batch size to help ensure that multiple threads
@@ -15,27 +16,33 @@ routing-pipeline:
     - gamma: '/value == "g"'
   sink:
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_alpha
         routes:
           - alpha
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_beta
         routes:
           - beta
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_alpha_default
         routes:
           - alpha
           - _default
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_alpha_beta_gamma
         routes:
           - alpha
           - beta
           - gamma
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_default
         routes:
           - _default
     - in_memory:
+        acknowledgments: true
         testing_key: ConditionalRoutingIT_all

data-prepper-core/src/main/java/org/opensearch/dataprepper/core/pipeline/router/Router.java

@@ -10,6 +10,7 @@
 import org.opensearch.dataprepper.model.record.Record;
 
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
@@ -41,6 +42,8 @@ public <C> void route(
         Objects.requireNonNull(dataFlowComponents);
         Objects.requireNonNull(componentRecordsConsumer);
 
+        final Map<Object, Collection<Record>> componentRecords = new HashMap<>();
+
         final Map<Record, Set<String>> recordsToRoutes = routeEventEvaluator.evaluateEventRoutes(allRecords);
 
         boolean allRecordsRouted = false;
@@ -61,10 +64,14 @@ public <C> void route(
                         recordsUnRouted.remove(record);
                     }
                 }
-                componentRecordsConsumer.accept(component, records);
+                componentRecords.put((C)component, records);
             });
         }
 
+        for (Map.Entry<Object, Collection<Record>> entry : componentRecords.entrySet()) {
+            componentRecordsConsumer.accept((C)entry.getKey(), entry.getValue());
+        }
+
         if (recordsUnRouted != null) {
             for (Record record: recordsUnRouted) {
                 if (record.getData() instanceof Event) {