Synchronization fix for aggregate processor and aggregate event handles when attaching events to the aggregate group. (opensearch-project#6431)

There is a possible synchronization issue in the aggregate processor. It currently calls attachToEventAcknowledgementSet on the aggregate group outside of any locks. It is possible that one thread gets this group. Then thread two gets the closes the group. If thread 1 then attaches the event to that group, thread 2 may still reset it. The solution is to move attachToEventAcknowledgementSet into the locks.

Signed-off-by: David Venable <dlv@amazon.com>

Author: dlvenable

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

@@ -92,6 +92,7 @@ AggregateActionResponse handleEventForGroup(final Event event, final Identificat
         handleEventForGroupLock.lock();
         try {
             LOG.debug("Start critical section in handleEventForGroup");
+            aggregateGroup.attachToEventAcknowledgementSet(event);
             handleEventResponse = aggregateAction.handleEvent(event, aggregateGroup);
             aggregateGroupManager.putGroupWithHash(hash, aggregateGroup);
         } catch (final Exception e) {

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

@@ -97,10 +97,8 @@ private AggregateAction loadAggregateAction(final PluginFactory pluginFactory) {
         return pluginFactory.loadPlugin(AggregateAction.class, actionPluginSetting);
     }
 
-    AggregateGroup getAggregateGroupForEvent(final IdentificationKeysHasher.IdentificationKeysMap identificationKeysMap, final Event event) {
-        AggregateGroup aggregateGroup = aggregateGroupManager.getAggregateGroup(identificationKeysMap);
-        aggregateGroup.attachToEventAcknowledgementSet(event);
-        return aggregateGroup;
+    private AggregateGroup getAggregateGroupForEvent(final IdentificationKeysHasher.IdentificationKeysMap identificationKeysMap) {
+        return aggregateGroupManager.getAggregateGroup(identificationKeysMap);
     }
 
     @Override
@@ -134,7 +132,7 @@ public Collection<Record<Event>> doExecute(Collection<Record<Event>> records) {
                 continue;
             }
             final IdentificationKeysHasher.IdentificationKeysMap identificationKeysMap = identificationKeysHasher.createIdentificationKeysMapFromEvent(event);
-            final AggregateGroup aggregateGroupForEvent = getAggregateGroupForEvent(identificationKeysMap, event);
+            final AggregateGroup aggregateGroupForEvent = getAggregateGroupForEvent(identificationKeysMap);
 
             final AggregateActionResponse handleEventResponse = aggregateActionSynchronizer.handleEventForGroup(event, identificationKeysMap, aggregateGroupForEvent);
 

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

@@ -158,10 +158,11 @@ void handleEventForGroup_calls_expected_functions_and_returns_correct_AggregateA
 
         final AggregateActionResponse handleEventResponse = objectUnderTest.handleEventForGroup(event, identificationKeysMap, aggregateGroup);
 
-        final InOrder inOrder = Mockito.inOrder(concludeGroupLock, handleEventForGroupLock, aggregateAction, aggregateGroupManager);
+        final InOrder inOrder = Mockito.inOrder(concludeGroupLock, handleEventForGroupLock, aggregateGroup, aggregateAction, aggregateGroupManager);
         inOrder.verify(concludeGroupLock).lock();
         inOrder.verify(concludeGroupLock).unlock();
         inOrder.verify(handleEventForGroupLock).lock();
+        inOrder.verify(aggregateGroup).attachToEventAcknowledgementSet(event);
         inOrder.verify(aggregateAction).handleEvent(event, aggregateGroup);
         inOrder.verify(aggregateGroupManager).putGroupWithHash(identificationKeysMap, aggregateGroup);
         inOrder.verify(handleEventForGroupLock).unlock();

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

@@ -9,6 +9,7 @@
 
 package org.opensearch.dataprepper.plugins.processor.aggregate;
 
+import org.junit.jupiter.api.AfterEach;
 import org.opensearch.dataprepper.expression.ExpressionEvaluator;
 import org.opensearch.dataprepper.metrics.MetricNames;
 import org.opensearch.dataprepper.metrics.PluginMetrics;
@@ -50,6 +51,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
@@ -158,6 +160,11 @@ void setUp() {
         when(pluginMetrics.timer(MetricNames.TIME_ELAPSED)).thenReturn(timeElapsed);
     }
 
+    @AfterEach
+    void processorDoesNotAttachEventsDirectly() {
+        verify(aggregateGroup, never()).attachToEventAcknowledgementSet(any());
+    }
+
     @Test
     void invalid_aggregate_when_statement_throws_InvalidPluginConfigurationException() {
         final String whenCondition = UUID.randomUUID().toString();