Add parameter to acknowledge group events on conclude immediately, and a parameter to disable group acknowledgments in aggregate processor (opensearch-project#6430)

Signed-off-by: Taylor Gray <tylgry@amazon.com>

Author: graytaylor0

data-prepper-plugins/aggregate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateActionSynchronizer.java

@@ -46,11 +46,17 @@ class AggregateActionSynchronizer {
     private final AggregateAction aggregateAction;
     private final AggregateGroupManager aggregateGroupManager;
 
+    private final boolean disableGroupAcknowledgments;
+
     private static final Logger LOG = LoggerFactory.getLogger(AggregateActionSynchronizer.class);
 
-    private AggregateActionSynchronizer(final AggregateAction aggregateAction, final AggregateGroupManager aggregateGroupManager, final PluginMetrics pluginMetrics) {
+    private AggregateActionSynchronizer(final AggregateAction aggregateAction,
+                                        final AggregateGroupManager aggregateGroupManager,
+                                        final PluginMetrics pluginMetrics,
+                                        final AggregateProcessorConfig aggregateProcessorConfig) {
         this.aggregateAction = aggregateAction;
         this.aggregateGroupManager = aggregateGroupManager;
+        this.disableGroupAcknowledgments = aggregateProcessorConfig.getDisableGroupAcknowledgments();
 
         this.actionHandleEventsProcessingErrors = pluginMetrics.counter(ACTION_HANDLE_EVENTS_PROCESSING_ERRORS);
         this.actionConcludeGroupEventsProcessingErrors = pluginMetrics.counter(ACTION_CONCLUDE_GROUP_EVENTS_PROCESSING_ERRORS);
@@ -92,7 +98,9 @@ AggregateActionResponse handleEventForGroup(final Event event, final Identificat
         handleEventForGroupLock.lock();
         try {
             LOG.debug("Start critical section in handleEventForGroup");
-            aggregateGroup.attachToEventAcknowledgementSet(event);
+            if (!disableGroupAcknowledgments) {
+                aggregateGroup.attachToEventAcknowledgementSet(event);
+            }
             handleEventResponse = aggregateAction.handleEvent(event, aggregateGroup);
             aggregateGroupManager.putGroupWithHash(hash, aggregateGroup);
         } catch (final Exception e) {
@@ -107,8 +115,11 @@ AggregateActionResponse handleEventForGroup(final Event event, final Identificat
     }
 
     static class AggregateActionSynchronizerProvider {
-        public AggregateActionSynchronizer provide(final AggregateAction aggregateAction, final AggregateGroupManager aggregateGroupManager, final PluginMetrics pluginMetrics) {
-            return new AggregateActionSynchronizer(aggregateAction, aggregateGroupManager, pluginMetrics);
+        public AggregateActionSynchronizer provide(final AggregateAction aggregateAction,
+                                                   final AggregateGroupManager aggregateGroupManager,
+                                                   final PluginMetrics pluginMetrics,
+                                                   final AggregateProcessorConfig aggregateProcessorConfig) {
+            return new AggregateActionSynchronizer(aggregateAction, aggregateGroupManager, pluginMetrics, aggregateProcessorConfig);
         }
     }
 }

data-prepper-plugins/aggregate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateGroupManager.java

@@ -10,6 +10,7 @@
 package org.opensearch.dataprepper.plugins.processor.aggregate;
 
 import com.google.common.collect.Maps;
+import org.opensearch.dataprepper.model.event.EventHandle;
 import org.opensearch.dataprepper.plugins.hasher.IdentificationKeysHasher;
 
 import java.time.Duration;
@@ -21,9 +22,11 @@ class AggregateGroupManager {
 
     private final Map<IdentificationKeysHasher.IdentificationKeysMap, AggregateGroup> allGroups = Maps.newConcurrentMap();
     private final Duration groupDuration;
+    private final boolean acknowledgeOnConclude;
 
-    AggregateGroupManager(final Duration groupDuration) {
+    AggregateGroupManager(final Duration groupDuration, final boolean acknowledgeOnConclude) {
         this.groupDuration = groupDuration;
+        this.acknowledgeOnConclude = acknowledgeOnConclude;
     }
 
     AggregateGroup getAggregateGroup(final IdentificationKeysHasher.IdentificationKeysMap identificationKeysMap) {
@@ -43,6 +46,13 @@ List<Map.Entry<IdentificationKeysHasher.IdentificationKeysMap, AggregateGroup>>
 
     void closeGroup(final IdentificationKeysHasher.IdentificationKeysMap hashKeyMap, final AggregateGroup group) {
         allGroups.remove(hashKeyMap, group);
+
+        if (acknowledgeOnConclude) {
+            EventHandle handle = group.getEventHandle();
+            if (handle != null) {
+                handle.release(true);
+            }
+        }
         group.resetGroup();
     }
 

data-prepper-plugins/aggregate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateProcessor.java

@@ -60,7 +60,7 @@ public class AggregateProcessor extends AbstractProcessor<Record<Event>, Record<
 
     @DataPrepperPluginConstructor
     public AggregateProcessor(final AggregateProcessorConfig aggregateProcessorConfig, final PluginMetrics pluginMetrics, final PluginFactory pluginFactory, final ExpressionEvaluator expressionEvaluator) {
-        this(aggregateProcessorConfig, pluginMetrics, pluginFactory, new AggregateGroupManager(aggregateProcessorConfig.getGroupDuration()),
+        this(aggregateProcessorConfig, pluginMetrics, pluginFactory, new AggregateGroupManager(aggregateProcessorConfig.getGroupDuration(), aggregateProcessorConfig.getAcknowledgeOnConclude()),
                 new IdentificationKeysHasher(aggregateProcessorConfig.getIdentificationKeys()), new AggregateActionSynchronizer.AggregateActionSynchronizerProvider(), expressionEvaluator);
     }
     public AggregateProcessor(final AggregateProcessorConfig aggregateProcessorConfig, final PluginMetrics pluginMetrics, final PluginFactory pluginFactory, final AggregateGroupManager aggregateGroupManager,
@@ -73,7 +73,7 @@ public AggregateProcessor(final AggregateProcessorConfig aggregateProcessorConfi
         this.expressionEvaluator = expressionEvaluator;
         this.identificationKeysHasher = identificationKeysHasher;
         this.aggregateAction = loadAggregateAction(pluginFactory);
-        this.aggregateActionSynchronizer = aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics);
+        this.aggregateActionSynchronizer = aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics, aggregateProcessorConfig);
 
         this.actionConcludeGroupEventsOutCounter = pluginMetrics.counter(ACTION_CONCLUDE_GROUP_EVENTS_OUT);
         this.actionConcludeGroupEventsDroppedCounter = pluginMetrics.counter(ACTION_CONCLUDE_GROUP_EVENTS_DROPPED);

data-prepper-plugins/aggregate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateProcessorConfig.java

@@ -81,6 +81,14 @@ public class AggregateProcessorConfig {
     })
     private String whenCondition;
 
+    @JsonPropertyDescription("When set to true, releases the group's event handle when the group concludes. " +
+            "This sacrifices true end-to-end acknowledgments for aggregated events but prevents reprocessing.")
+    @JsonProperty("acknowledge_on_conclude")
+    private Boolean acknowledgeOnConclude = false;
+
+    @JsonProperty("disable_group_acknowledgments")
+    private Boolean disableGroupAcknowledgments = false;
+
     public List<String> getIdentificationKeys() {
         return identificationKeys;
     }
@@ -112,4 +120,12 @@ boolean isValidConfig() {
 
     public PluginModel getAggregateAction() { return aggregateAction; }
 
+    public Boolean getAcknowledgeOnConclude() {
+        return acknowledgeOnConclude;
+    }
+
+    public Boolean getDisableGroupAcknowledgments() {
+        return disableGroupAcknowledgments;
+    }
+
 }

data-prepper-plugins/aggregate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateActionSynchronizerTest.java

@@ -36,6 +36,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -75,6 +76,9 @@ public class AggregateActionSynchronizerTest {
     @Mock
     private Counter actionConcludeGroupEventsProcessingErrors;
 
+    @Mock
+    private AggregateProcessorConfig aggregateProcessorConfig;
+
     @BeforeEach
     void setup() {
         doNothing().when(handleEventForGroupLock).lock();
@@ -93,7 +97,7 @@ void setup() {
 
     private AggregateActionSynchronizer createObjectUnderTest() {
         final AggregateActionSynchronizer.AggregateActionSynchronizerProvider aggregateActionSynchronizerProvider = new AggregateActionSynchronizer.AggregateActionSynchronizerProvider();
-        return aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics);
+        return aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics, aggregateProcessorConfig);
     }
 
     @Test
@@ -170,6 +174,29 @@ void handleEventForGroup_calls_expected_functions_and_returns_correct_AggregateA
         assertThat(handleEventResponse, equalTo(aggregateActionResponse));
     }
 
+    @Test
+    void handleEventForGroup_calls_expected_functions_and_returns_correct_AggregateActionResponse_with_disable_group_acknowledgments_true() {
+        when(aggregateProcessorConfig.getDisableGroupAcknowledgments()).thenReturn(true);
+
+        final AggregateActionSynchronizer objectUnderTest = createObjectUnderTest();
+        when(aggregateAction.handleEvent(event, aggregateGroup)).thenReturn(aggregateActionResponse);
+
+        final AggregateActionResponse handleEventResponse = objectUnderTest.handleEventForGroup(event, identificationKeysMap, aggregateGroup);
+
+        final InOrder inOrder = Mockito.inOrder(concludeGroupLock, handleEventForGroupLock, aggregateGroup, aggregateAction, aggregateGroupManager);
+        inOrder.verify(aggregateGroup).getConcludeGroupLock();
+        inOrder.verify(aggregateGroup).getHandleEventForGroupLock();
+        inOrder.verify(concludeGroupLock).lock();
+        inOrder.verify(concludeGroupLock).unlock();
+        inOrder.verify(handleEventForGroupLock).lock();
+        inOrder.verify(aggregateAction).handleEvent(event, aggregateGroup);
+        inOrder.verify(aggregateGroupManager).putGroupWithHash(identificationKeysMap, aggregateGroup);
+        inOrder.verify(handleEventForGroupLock).unlock();
+
+        verifyNoMoreInteractions(aggregateGroup);
+        assertThat(handleEventResponse, equalTo(aggregateActionResponse));
+    }
+
     @Test
     void locks_are_unlocked_and_event_returned_when_aggregateAction_handleEvent_throws_exception() {
         final AggregateActionSynchronizer objectUnderTest = createObjectUnderTest();

data-prepper-plugins/aggregate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateGroupManagerTest.java

@@ -9,6 +9,7 @@
 
 package org.opensearch.dataprepper.plugins.processor.aggregate;
 
+import org.opensearch.dataprepper.model.event.EventHandle;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.opensearch.dataprepper.plugins.hasher.IdentificationKeysHasher;
@@ -27,6 +28,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class AggregateGroupManagerTest {
@@ -46,7 +48,7 @@ void setup() {
     }
 
     private AggregateGroupManager createObjectUnderTest() {
-        return new AggregateGroupManager(TEST_GROUP_DURATION);
+        return new AggregateGroupManager(TEST_GROUP_DURATION, false);
     }
 
     @Test
@@ -145,4 +147,29 @@ void getGroupsToConclude_with_force_conclude_return_all() {
             assertThat(groupsToConclude.get(1).getValue(), equalTo(groupToConclude1));
         }
     }
+
+    @Test
+    void closeGroup_with_acknowledge_on_conclude_releases_event_handle() {
+        aggregateGroupManager = new AggregateGroupManager(TEST_GROUP_DURATION, true);
+
+        final AggregateGroup group = mock(AggregateGroup.class);
+        final EventHandle eventHandle = mock(EventHandle.class);
+        when(group.getEventHandle()).thenReturn(eventHandle);
+
+        aggregateGroupManager.closeGroup(identificationKeysMap, group);
+
+        verify(eventHandle).release(true);
+        verify(group).resetGroup();
+    }
+
+    @Test
+    void closeGroup_without_acknowledge_on_conclude_does_not_release_event_handle() {
+        aggregateGroupManager = new AggregateGroupManager(TEST_GROUP_DURATION, false);
+
+        final AggregateGroup group = mock(AggregateGroup.class);
+
+        aggregateGroupManager.closeGroup(identificationKeysMap, group);
+
+        verify(group).resetGroup();
+    }
 }

data-prepper-plugins/aggregate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateProcessorTest.java

@@ -148,7 +148,7 @@ void setUp() {
                 .withEventType("event")
                 .build();
 
-        when(aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics)).thenReturn(aggregateActionSynchronizer);
+        when(aggregateActionSynchronizerProvider.provide(aggregateAction, aggregateGroupManager, pluginMetrics, aggregateProcessorConfig)).thenReturn(aggregateActionSynchronizer);
 
         when(pluginMetrics.counter(AggregateProcessor.ACTION_HANDLE_EVENTS_OUT)).thenReturn(actionHandleEventsOutCounter);
         when(pluginMetrics.counter(AggregateProcessor.ACTION_HANDLE_EVENTS_DROPPED)).thenReturn(actionHandleEventsDroppedCounter);