gh-4879 Improve locking when closing aggregates

Author: at055612

stroom-proxy/stroom-proxy-app/src/main/java/stroom/proxy/app/handler/PreAggregator.java

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -452,15 +453,30 @@ private List<Part> calculateOverflowingParts(final FileGroup fileGroup) throws I
         return Collections.singletonList(new Part(partItems, partBytes, List.copyOf(partEntries)));
     }
 
-    private void closeAggregate(final FeedKey feedKey,
-                                final AggregateState aggregateState) {
-        LOGGER.debug("Closing aggregate: {}", aggregateState);
+    private boolean closeAggregate(final FeedKey feedKey,
+                                   final AggregateState aggregateState) {
+        LOGGER.debug(() -> LogUtil.message("closeAggregate() - feedKey: {}, {}, waiting for lock",
+                feedKey, aggregateState));
         final Lock lock = feedKeyLock.get(feedKey);
         lock.lock();
         try {
-            destination.accept(aggregateState.aggregateDir);
-            aggregateStateMap.remove(feedKey);
-            LOGGER.debug("Closed aggregate: {}", aggregateState);
+            // Now we hold the feedKey lock, re-check the aggregateStateMap
+            if (aggregateStateMap.containsKey(feedKey)) {
+                LOGGER.debug(() -> LogUtil.message("closeAggregate() - feedKey: {}, {}, acquired lock",
+                        feedKey, aggregateState));
+
+                destination.accept(aggregateState.aggregateDir);
+                aggregateStateMap.remove(feedKey);
+                LOGGER.debug(() -> LogUtil.message("closeAggregate() - feedKey: {}, {}, closed aggregate",
+                        feedKey, aggregateState));
+                return true;
+            } else {
+                LOGGER.debug(() -> LogUtil.message(
+                        "closeAggregate() - feedKey: {}, {}, " +
+                        "feedKey not in aggregateStateMap, another thread must have closed it.",
+                        feedKey, aggregateState));
+                return false;
+            }
         } finally {
             lock.unlock();
         }
@@ -570,13 +586,23 @@ private AggregateState createAggregate(final FeedKey feedKey) {
      */
     private synchronized void closeOldAggregates() {
         final AtomicInteger count = new AtomicInteger();
-        aggregateStateMap.forEach((feedKey, aggregateState) -> {
-            if (aggregateState.isAggregateTooOld(aggregatorConfig)) {
-                // Close the current aggregate.
-                count.incrementAndGet();
-                closeAggregate(feedKey, aggregateState);
+        final Set<FeedKey> feedKeys = aggregateStateMap.keySet();
+        for (final FeedKey feedKey : feedKeys) {
+            // It's possible another thread may have removed it
+            final AggregateState aggregateState = aggregateStateMap.get(feedKey);
+            if (aggregateState != null
+                && aggregateState.isAggregateTooOld(aggregatorConfig)) {
+                // Close the current aggregate, under a feedKey lock, so again,
+                // another thread may beat us
+                final boolean didClose = closeAggregate(feedKey, aggregateState);
+                if (didClose) {
+                    count.incrementAndGet();
+                } else {
+                    LOGGER.debug("closeAggregate() - feedKey: {}, aggregateState: {}, didn't close",
+                            feedKey, aggregateState);
+                }
             }
-        });
+        }
         if (LOGGER.isDebugEnabled()) {
             if (count.get() > 0) {
                 LOGGER.debug("closeOldAggregates() - closed {} old aggregates", count);

unreleased_changes/20250415_171225_150__4879.md

@@ -0,0 +1,24 @@
+* Issue **#4879** : Proxy - Improve concurrency protection when closing aggregates.
+
+
+```sh
+# ********************************************************************************
+# Issue title: Proxy - exception when closing aggregates
+# Issue link:  https://github.com/gchq/stroom/issues/4879
+# ********************************************************************************
+
+# ONLY the top line will be included as a change entry in the CHANGELOG.
+# The entry should be in GitHub flavour markdown and should be written on a SINGLE
+# line with no hard breaks. You can have multiple change files for a single GitHub issue.
+# The  entry should be written in the imperative mood, i.e. 'Fix nasty bug' rather than
+# 'Fixed nasty bug'.
+#
+# Examples of acceptable entries are:
+#
+#
+# * Issue **123** : Fix bug with an associated GitHub issue in this repository
+#
+# * Issue **namespace/other-repo#456** : Fix bug with an associated GitHub issue in another repository
+#
+# * Fix bug with no associated GitHub issue.
+```