cndb-14201: allow UCS#getNextBackgroundTasks to config custom compaction observer for composite compaction (#1767)

CNDB: riptano/cndb#14203

Author: jasonstack

src/java/org/apache/cassandra/db/compaction/UnifiedCompactionStrategy.java

@@ -31,7 +31,6 @@
 import java.util.function.BiPredicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
-import java.util.stream.Collectors;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -343,7 +342,7 @@ public synchronized Collection<AbstractCompactionTask> getNextBackgroundTasks(in
         if (expirationTasks != null)
             return expirationTasks;
 
-        return getNextBackgroundTasks(getNextCompactionAggregates(), gcBefore);
+        return getNextBackgroundTasks(getNextCompactionAggregates(), gcBefore, this);
     }
 
     /// Check for fully expired sstables and return a collection of expiration tasks if found.
@@ -401,21 +400,23 @@ public Set<CompactionSSTable> getFullyExpiredSSTables(int gcBefore)
                                                                    controller.getIgnoreOverlapsInExpirationCheck());
     }
 
-    /// Used by CNDB where compaction aggregates come from etcd rather than the strategy
+    /// Used by CNDB where compaction aggregates come from etcd rather than the strategy.
     /// @return collection of `AbstractCompactionTask`, which could be either a `CompactionTask` or a `UnifiedCompactionTask`
-    public synchronized Collection<AbstractCompactionTask> getNextBackgroundTasks(Collection<CompactionAggregate> aggregates, int gcBefore)
+    public synchronized Collection<AbstractCompactionTask> getNextBackgroundTasks(Collection<CompactionAggregate> aggregates, int gcBefore,
+                                                                                  CompactionObserver compositeCompactionObserver)
     {
         controller.onStrategyBackgroundTaskRequest();
-        return createCompactionTasks(aggregates, gcBefore);
+        return createCompactionTasks(aggregates, gcBefore, compositeCompactionObserver);
     }
 
-    private Collection<AbstractCompactionTask> createCompactionTasks(Collection<CompactionAggregate> aggregates, int gcBefore)
+    private Collection<AbstractCompactionTask> createCompactionTasks(Collection<CompactionAggregate> aggregates, int gcBefore,
+                                                                     CompactionObserver compositeCompactionObserver)
     {
         Collection<AbstractCompactionTask> tasks = new ArrayList<>(aggregates.size());
         try
         {
             for (CompactionAggregate aggregate : aggregates)
-                createAndAddTasks(gcBefore, (CompactionAggregate.UnifiedAggregate) aggregate, tasks);
+                createAndAddTasks(gcBefore, (CompactionAggregate.UnifiedAggregate) aggregate, tasks, compositeCompactionObserver);
 
             return tasks;
         }
@@ -426,7 +427,8 @@ private Collection<AbstractCompactionTask> createCompactionTasks(Collection<Comp
     }
 
     /// Create compaction tasks for the given aggregate and add them to the given tasks list.
-    public void createAndAddTasks(int gcBefore, CompactionAggregate.UnifiedAggregate aggregate, Collection<? super UnifiedCompactionTask> tasks)
+    public void createAndAddTasks(int gcBefore, CompactionAggregate.UnifiedAggregate aggregate,
+                                  Collection<? super UnifiedCompactionTask> tasks, CompactionObserver compositeCompactionObserver)
     {
         CompactionPick selected = aggregate.getSelected();
         int parallelism = aggregate.getPermittedParallelism();
@@ -440,7 +442,7 @@ public void createAndAddTasks(int gcBefore, CompactionAggregate.UnifiedAggregate
         {
             // This will ignore the range of the operation, which is fine.
             backgroundCompactions.setSubmitted(this, transaction.opId(), aggregate);
-            createAndAddTasks(gcBefore, transaction, aggregate.operationRange(), aggregate.keepOriginals(), getShardingStats(aggregate), parallelism, tasks);
+            createAndAddTasks(gcBefore, transaction, aggregate.operationRange(), aggregate.keepOriginals(), getShardingStats(aggregate), parallelism, tasks, compositeCompactionObserver);
         }
         else
         {
@@ -612,7 +614,7 @@ void createAndAddTasks(int gcBefore,
                            int parallelism,
                            Collection<? super CompactionTask> tasks)
     {
-        createAndAddTasks(gcBefore, transaction, null, false, shardingStats, parallelism, tasks);
+        createAndAddTasks(gcBefore, transaction, null, false, shardingStats, parallelism, tasks, this);
     }
 
     @VisibleForTesting
@@ -622,10 +624,11 @@ void createAndAddTasks(int gcBefore,
                            boolean keepOriginals,
                            ShardingStats shardingStats,
                            int parallelism,
-                           Collection<? super UnifiedCompactionTask> tasks)
+                           Collection<? super UnifiedCompactionTask> tasks,
+                           CompactionObserver compositeCompactionObserver)
     {
         if (controller.parallelizeOutputShards() && parallelism > 1)
-            tasks.addAll(createParallelCompactionTasks(transaction, operationRange, keepOriginals, shardingStats, gcBefore, parallelism));
+            tasks.addAll(createParallelCompactionTasks(transaction, operationRange, keepOriginals, shardingStats, gcBefore, parallelism, compositeCompactionObserver));
         else
             tasks.add(createCompactionTask(transaction, operationRange, keepOriginals, shardingStats, gcBefore));
     }
@@ -697,7 +700,8 @@ private Collection<UnifiedCompactionTask> createParallelCompactionTasks(Lifecycl
                                                                             boolean keepOriginals,
                                                                             ShardingStats shardingStats,
                                                                             int gcBefore,
-                                                                            int parallelism)
+                                                                            int parallelism,
+                                                                            CompactionObserver compositeCompactionObserver)
     {
         final int coveredShardCount = shardingStats.coveredShardCount;
         assert parallelism > 1;
@@ -706,7 +710,7 @@ private Collection<UnifiedCompactionTask> createParallelCompactionTasks(Lifecycl
         ShardManager shardManager = getShardManager();
         CompositeLifecycleTransaction compositeTransaction = new CompositeLifecycleTransaction(transaction);
         SharedCompactionProgress sharedProgress = new SharedCompactionProgress(transaction.opId(), transaction.opType(), TableOperation.Unit.BYTES);
-        SharedCompactionObserver sharedObserver = new SharedCompactionObserver(this);
+        SharedCompactionObserver sharedObserver = new SharedCompactionObserver(compositeCompactionObserver);
         SharedTableOperation sharedOperation = new SharedTableOperation(sharedProgress);
         List<UnifiedCompactionTask> tasks = shardManager.splitSSTablesInShardsLimited(
             sstables,

test/unit/org/apache/cassandra/db/compaction/UnifiedCompactionStrategyTest.java

@@ -29,11 +29,14 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import javax.annotation.Nullable;
+
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -1753,6 +1756,59 @@ public void testDontCreateParallelTasks()
         assertEquals(1, tasks.size());
         assertEquals(allSSTables, tasks.get(0).inputSSTables());
     }
+
+    @Test
+    public void testCustomCompositeCompactionObserver()
+    {
+        int numShards = 5;
+        Set<SSTableReader> allSSTables = new HashSet<>();
+        allSSTables.addAll(mockNonOverlappingSSTables(10, 0, 100 << 20));
+        allSSTables.addAll(mockNonOverlappingSSTables(15, 1, 200 << 20));
+        allSSTables.addAll(mockNonOverlappingSSTables(25, 2, 400 << 20));
+        dataTracker.addInitialSSTables(allSSTables);
+        Controller controller = Mockito.mock(Controller.class);
+        when(controller.getNumShards(anyDouble())).thenReturn(numShards);
+        when(controller.parallelizeOutputShards()).thenReturn(true);
+        UnifiedCompactionStrategy strategy = new UnifiedCompactionStrategy(strategyFactory, controller);
+        strategy.startup();
+        LifecycleTransaction txn = dataTracker.tryModify(allSSTables, OperationType.COMPACTION);
+        var tasks = new ArrayList<CompactionTask>();
+
+        AtomicInteger compositeInProgress = new AtomicInteger(0);
+        AtomicInteger compositedCompleted = new AtomicInteger(0);
+        CompactionObserver compositeCompactionObserver = new CompactionObserver()
+        {
+            @Override
+            public void onInProgress(CompactionProgress progress)
+            {
+                compositeInProgress.incrementAndGet();
+            }
+
+            @Override
+            public void onCompleted(UUID id, @Nullable Throwable error)
+            {
+                compositedCompleted.incrementAndGet();
+            }
+        };
+        strategy.createAndAddTasks(0, txn, null, false, strategy.makeShardingStats(txn), 1000, tasks, compositeCompactionObserver);
+        assertEquals(numShards, tasks.size());
+
+        // move all tasks to in-progress
+        CompactionProgress progress = mockProgress(strategy, txn.opId());
+        for (CompactionTask task : tasks)
+            task.getCompObservers().forEach(o -> o.onInProgress(progress));
+
+        assertThat(compositeInProgress).hasValue(1);
+        assertThat(compositedCompleted).hasValue(0);
+
+        // move all tasks to complete
+        for (CompactionTask task : tasks)
+            task.getCompObservers().forEach(o -> o.onCompleted(task.transaction.opId(), null));
+
+        assertThat(compositeInProgress).hasValue(1);
+        assertThat(compositedCompleted).hasValue(1);
+    }
+
     @Test
     public void testMaximalSelection()
     {

test/unit/org/apache/cassandra/db/compaction/unified/RangedAggregatesTest.java

@@ -122,7 +122,7 @@ private void testRangedAggregates(int numShards, int aggregateParallelism, int r
         {
             // execute each partial aggregate separately because we can only mark the inputs compacting once
             List<AbstractCompactionTask> tasks = new ArrayList<>();
-            strategy.createAndAddTasks(0, maximal, tasks);
+            strategy.createAndAddTasks(0, maximal, tasks, strategy);
             totalTaskCount += tasks.size();
             List<Future<?>> futures = tasks.stream()
                                            .map(t -> ForkJoinPool.commonPool()