Add parameter to acknowledge group events on conclude immediately in aggregate processor

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

Author: graytaylor0

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,

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

@@ -81,6 +81,11 @@ 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;
+
     public List<String> getIdentificationKeys() {
         return identificationKeys;
     }
@@ -112,4 +117,8 @@ boolean isValidConfig() {
 
     public PluginModel getAggregateAction() { return aggregateAction; }
 
+    public Boolean getAcknowledgeOnConclude() {
+        return acknowledgeOnConclude;
+    }
+
 }

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();
+    }
 }