CNDB-14237-main-cherry-pick: register composite compaction observer as additional obverser on top of UCS (#1783)

### What is the issue

UCS didn't clear pending compaction tasks in
`BackgroundCompactions#compactions` for parallel background compaction

### What does this PR fix and why was it fixed

Register both UCS and composite compaction observer for parallel
compaction task: both UCS and CNDB are notified

Author: jasonstack

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

@@ -18,11 +18,19 @@
 
 package org.apache.cassandra.db.compaction;
 
+import java.util.List;
 import java.util.UUID;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
+import javax.annotation.Nullable;
+
+import com.google.common.collect.ImmutableList;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.utils.Throwables;
+
 /// Utility class to share a compaction observer among multiple compaction tasks and only report start and completion
 /// once when the first task starts and completion when all tasks complete (successfully or not, where the passed
 /// `isSuccess` state is a logical and of the subtasks').
@@ -34,14 +42,25 @@
 /// if assertions are enabled.
 public class SharedCompactionObserver implements CompactionObserver
 {
+    private static final Logger logger = LoggerFactory.getLogger(SharedCompactionObserver.class);
+
     private final AtomicInteger toReportOnComplete = new AtomicInteger(0);
     private final AtomicReference<Throwable> onCompleteException = new AtomicReference(null);
     private final AtomicReference<CompactionProgress> inProgressReported = new AtomicReference<>(null);
-    private final CompactionObserver observer;
+
+    private final List<CompactionObserver> compObservers;
 
     public SharedCompactionObserver(CompactionObserver observer)
     {
-        this.observer = observer;
+        this(observer, null);
+    }
+
+    public SharedCompactionObserver(CompactionObserver primary, @Nullable CompactionObserver secondary)
+    {
+        if (primary == null)
+            throw new IllegalArgumentException("Primary observer cannot be null");
+
+        this.compObservers = secondary != null ? ImmutableList.of(primary, secondary) : ImmutableList.of(primary);
     }
 
     public void registerExpectedSubtask()
@@ -55,7 +74,13 @@ public void registerExpectedSubtask()
     public void onInProgress(CompactionProgress progress)
     {
         if (inProgressReported.compareAndSet(null, progress))
-            observer.onInProgress(progress);
+        {
+            Throwable err = null;
+            for (CompactionObserver compObserver : compObservers)
+                err = Throwables.perform(err, () -> compObserver.onInProgress(progress));
+
+            Throwables.maybeFail(err);
+        }
         else
             assert inProgressReported.get() == progress; // progress object must also be shared
     }
@@ -69,6 +94,12 @@ public void onCompleted(UUID id, Throwable err)
         assert remainingToComplete >= 0 : "onCompleted called without corresponding registerExpectedSubtask";
         // The individual operation ID given here may be different from the shared ID. Pass on the shared one.
         if (remainingToComplete == 0)
-            observer.onCompleted(inProgressReported.get().operationId(), onCompleteException.get());
+        {
+            Throwable error = null;
+            for (CompactionObserver compObserver : compObservers)
+                error = Throwables.perform(error, () -> compObserver.onCompleted(inProgressReported.get().operationId(), onCompleteException.get()));
+
+            Throwables.maybeFail(error);
+        }
     }
 }

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

@@ -32,6 +32,8 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import javax.annotation.Nullable;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterables;
@@ -342,7 +344,7 @@ public synchronized Collection<AbstractCompactionTask> getNextBackgroundTasks(in
         if (expirationTasks != null)
             return expirationTasks;
 
-        return getNextBackgroundTasks(getNextCompactionAggregates(), gcBefore, this);
+        return getNextBackgroundTasks(getNextCompactionAggregates(), gcBefore, null);
     }
 
     /// Check for fully expired sstables and return a collection of expiration tasks if found.
@@ -403,20 +405,20 @@ public Set<CompactionSSTable> getFullyExpiredSSTables(int gcBefore)
     /// 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,
-                                                                                  CompactionObserver compositeCompactionObserver)
+                                                                                  @Nullable CompactionObserver additionalObserver)
     {
         controller.onStrategyBackgroundTaskRequest();
-        return createCompactionTasks(aggregates, gcBefore, compositeCompactionObserver);
+        return createCompactionTasks(aggregates, gcBefore, additionalObserver);
     }
 
     private Collection<AbstractCompactionTask> createCompactionTasks(Collection<CompactionAggregate> aggregates, int gcBefore,
-                                                                     CompactionObserver compositeCompactionObserver)
+                                                                     @Nullable CompactionObserver additionalObserver)
     {
         Collection<AbstractCompactionTask> tasks = new ArrayList<>(aggregates.size());
         try
         {
             for (CompactionAggregate aggregate : aggregates)
-                createAndAddTasks(gcBefore, (CompactionAggregate.UnifiedAggregate) aggregate, tasks, compositeCompactionObserver);
+                createAndAddTasks(gcBefore, (CompactionAggregate.UnifiedAggregate) aggregate, tasks, additionalObserver);
 
             return tasks;
         }
@@ -428,7 +430,7 @@ 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, CompactionObserver compositeCompactionObserver)
+                                  Collection<? super UnifiedCompactionTask> tasks, @Nullable CompactionObserver additionalObserver)
     {
         CompactionPick selected = aggregate.getSelected();
         int parallelism = aggregate.getPermittedParallelism();
@@ -442,7 +444,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, compositeCompactionObserver);
+            createAndAddTasks(gcBefore, transaction, aggregate.operationRange(), aggregate.keepOriginals(), getShardingStats(aggregate), parallelism, tasks, additionalObserver);
         }
         else
         {
@@ -451,6 +453,12 @@ public void createAndAddTasks(int gcBefore, CompactionAggregate.UnifiedAggregate
         }
     }
 
+    /// Return the num of in-progress compactions tracked by UCS
+    public int getCompactionInProgress()
+    {
+        return backgroundCompactions.getCompactionsInProgress().size();
+    }
+
     private static RuntimeException rejectTasks(Iterable<? extends AbstractCompactionTask> tasks, Throwable error)
     {
         for (var task : tasks)
@@ -614,7 +622,7 @@ void createAndAddTasks(int gcBefore,
                            int parallelism,
                            Collection<? super CompactionTask> tasks)
     {
-        createAndAddTasks(gcBefore, transaction, null, false, shardingStats, parallelism, tasks, this);
+        createAndAddTasks(gcBefore, transaction, null, false, shardingStats, parallelism, tasks, null);
     }
 
     @VisibleForTesting
@@ -625,12 +633,12 @@ void createAndAddTasks(int gcBefore,
                            ShardingStats shardingStats,
                            int parallelism,
                            Collection<? super UnifiedCompactionTask> tasks,
-                           CompactionObserver compositeCompactionObserver)
+                           @Nullable CompactionObserver additionalObserver)
     {
         if (controller.parallelizeOutputShards() && parallelism > 1)
-            tasks.addAll(createParallelCompactionTasks(transaction, operationRange, keepOriginals, shardingStats, gcBefore, parallelism, compositeCompactionObserver));
+            tasks.addAll(createParallelCompactionTasks(transaction, operationRange, keepOriginals, shardingStats, gcBefore, parallelism, additionalObserver));
         else
-            tasks.add(createCompactionTask(transaction, operationRange, keepOriginals, shardingStats, gcBefore));
+            tasks.add(createCompactionTask(transaction, operationRange, keepOriginals, shardingStats, gcBefore, additionalObserver));
     }
 
     /// Create the sstable writer used for flushing.
@@ -677,9 +685,13 @@ private UnifiedCompactionTask createCompactionTask(LifecycleTransaction transact
     /// where we produce outputs but cannot delete the input sstables until all components of the operation are complete.
     ///
     /// @return a sharded compaction task that in turn will create a sharded compaction writer.
-    private UnifiedCompactionTask createCompactionTask(LifecycleTransaction transaction, Range<Token> operationRange, boolean keepOriginals, ShardingStats shardingStats, int gcBefore)
+    private UnifiedCompactionTask createCompactionTask(LifecycleTransaction transaction, Range<Token> operationRange, boolean keepOriginals, ShardingStats shardingStats, int gcBefore,
+                                                       @Nullable CompactionObserver additionalObserver)
     {
-        return new UnifiedCompactionTask(realm, this, transaction, gcBefore, keepOriginals, getShardManager(), shardingStats, operationRange, transaction.originals(), null, null, null);
+        UnifiedCompactionTask task = new UnifiedCompactionTask(realm, this, transaction, gcBefore, keepOriginals, getShardManager(), shardingStats, operationRange, transaction.originals(), null, null, null);
+        if (additionalObserver != null)
+            task.addObserver(additionalObserver);
+        return task;
     }
 
     @Override
@@ -701,7 +713,7 @@ private Collection<UnifiedCompactionTask> createParallelCompactionTasks(Lifecycl
                                                                             ShardingStats shardingStats,
                                                                             int gcBefore,
                                                                             int parallelism,
-                                                                            CompactionObserver compositeCompactionObserver)
+                                                                            @Nullable CompactionObserver additionalObserver)
     {
         final int coveredShardCount = shardingStats.coveredShardCount;
         assert parallelism > 1;
@@ -710,7 +722,8 @@ 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(compositeCompactionObserver);
+        SharedCompactionObserver sharedObserver = new SharedCompactionObserver(this, additionalObserver);
+
         SharedTableOperation sharedOperation = new SharedTableOperation(sharedProgress);
         List<UnifiedCompactionTask> tasks = shardManager.splitSSTablesInShardsLimited(
             sstables,
@@ -741,7 +754,7 @@ private Collection<UnifiedCompactionTask> createParallelCompactionTasks(Lifecycl
         if (tasks.size() == 1) // if there's just one range, make it a non-ranged task (to apply early open etc.)
         {
             assert tasks.get(0).inputSSTables().equals(sstables);
-            return Collections.singletonList(createCompactionTask(transaction, operationRange, keepOriginals, shardingStats, gcBefore));
+            return Collections.singletonList(createCompactionTask(transaction, operationRange, keepOriginals, shardingStats, gcBefore, additionalObserver));
         }
         else
             return tasks;

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

@@ -36,6 +36,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
 
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.*;
 
 public class SharedCompactionObserverTest
@@ -185,4 +186,57 @@ public void testErrorWrongProgress()
         when(mockProgress2.operationId()).thenReturn(UUID.randomUUID());
         Assert.assertThrows(AssertionError.class, () -> sharedCompactionObserver.onInProgress(mockProgress2));
     }
+
+    @Test
+    public void testMultipleObserver()
+    {
+        CompactionObserver mockObserver1 = Mockito.mock(CompactionObserver.class);
+        CompactionObserver mockObserver2 = Mockito.mock(CompactionObserver.class);
+        SharedCompactionObserver sharedCompactionObserver = new SharedCompactionObserver(mockObserver1, mockObserver2);
+
+        sharedCompactionObserver.registerExpectedSubtask();
+        sharedCompactionObserver.onInProgress(mockProgress);
+        verify(mockObserver1, times(1)).onInProgress(mockProgress);
+        verify(mockObserver2, times(1)).onInProgress(mockProgress);
+
+        sharedCompactionObserver.onCompleted(operationId, null);
+        verify(mockObserver1, times(1)).onCompleted(operationId, null);
+        verify(mockObserver2, times(1)).onCompleted(operationId, null);
+    }
+
+    @Test
+    public void testMultipleObserverWithOnInProgressError()
+    {
+        CompactionObserver mockObserver1 = Mockito.mock(CompactionObserver.class);
+        CompactionObserver mockObserver2 = Mockito.mock(CompactionObserver.class);
+        Mockito.doThrow(new RuntimeException("Injected Exception")).when(mockObserver1).onInProgress(any());
+
+        SharedCompactionObserver sharedCompactionObserver = new SharedCompactionObserver(mockObserver1, mockObserver2);
+
+        sharedCompactionObserver.registerExpectedSubtask();
+        assertThatThrownBy(() -> sharedCompactionObserver.onInProgress(mockProgress)).isInstanceOf(RuntimeException.class);
+
+        // both onInProgress are called
+        verify(mockObserver1, times(1)).onInProgress(mockProgress);
+        verify(mockObserver2, times(1)).onInProgress(mockProgress);
+    }
+
+    @Test
+    public void testMultipleObserverWithOnCompleteError()
+    {
+        CompactionObserver mockObserver1 = Mockito.mock(CompactionObserver.class);
+        CompactionObserver mockObserver2 = Mockito.mock(CompactionObserver.class);
+        Mockito.doThrow(new RuntimeException("Injected Exception")).when(mockObserver1).onCompleted(any(), any());
+
+        SharedCompactionObserver sharedCompactionObserver = new SharedCompactionObserver(mockObserver1, mockObserver2);
+
+        sharedCompactionObserver.registerExpectedSubtask();
+        sharedCompactionObserver.onInProgress(mockProgress);
+        verify(mockObserver1, times(1)).onInProgress(mockProgress);
+        verify(mockObserver2, times(1)).onInProgress(mockProgress);
+
+        assertThatThrownBy(() -> sharedCompactionObserver.onCompleted(operationId, null)).isInstanceOf(RuntimeException.class);
+        verify(mockObserver1, times(1)).onCompleted(operationId, null);
+        verify(mockObserver2, times(1)).onCompleted(operationId, null);
+    }
 }

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

@@ -81,6 +81,7 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.when;
 
 /**
@@ -1758,9 +1759,25 @@ public void testDontCreateParallelTasks()
     }
 
     @Test
-    public void testCustomCompositeCompactionObserver()
+    public void testAdditionalCompactionObserverForParallelCompaction()
+    {
+        testAdditionalCompactionObserver(5, 1000);
+    }
+
+    @Test
+    public void testAdditionalCompactionObserverForSingleCompactionSingleShard()
+    {
+        testAdditionalCompactionObserver(1, 1000);
+    }
+
+    @Test
+    public void testAdditionalCompactionObserverForSingleCompactionLimitedParallelism()
+    {
+        testAdditionalCompactionObserver(5, 1);
+    }
+
+    private void testAdditionalCompactionObserver(int numShards, int parallelism)
     {
-        int numShards = 5;
         Set<SSTableReader> allSSTables = new HashSet<>();
         allSSTables.addAll(mockNonOverlappingSSTables(10, 0, 100 << 20));
         allSSTables.addAll(mockNonOverlappingSSTables(15, 1, 200 << 20));
@@ -1769,7 +1786,9 @@ public void testCustomCompositeCompactionObserver()
         Controller controller = Mockito.mock(Controller.class);
         when(controller.getNumShards(anyDouble())).thenReturn(numShards);
         when(controller.parallelizeOutputShards()).thenReturn(true);
-        UnifiedCompactionStrategy strategy = new UnifiedCompactionStrategy(strategyFactory, controller);
+
+        BackgroundCompactions backgroundCompactions = Mockito.mock(BackgroundCompactions.class);
+        UnifiedCompactionStrategy strategy = new UnifiedCompactionStrategy(strategyFactory, backgroundCompactions, controller);
         strategy.startup();
         LifecycleTransaction txn = dataTracker.tryModify(allSSTables, OperationType.COMPACTION);
         var tasks = new ArrayList<CompactionTask>();
@@ -1790,8 +1809,15 @@ 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());
+        strategy.createAndAddTasks(0, txn, null, false, strategy.makeShardingStats(txn), parallelism, tasks, compositeCompactionObserver);
+        if (parallelism > 1)
+            assertEquals(numShards, tasks.size());
+        else
+            assertEquals(1, tasks.size());
+
+        assertThat(compositedCompleted).hasValue(0);
+        Mockito.verify(backgroundCompactions, times(0)).onInProgress(Mockito.any());
+        Mockito.verify(backgroundCompactions, times(0)).onCompleted(Mockito.any(), Mockito.any());
 
         // move all tasks to in-progress
         CompactionProgress progress = mockProgress(strategy, txn.opId());
@@ -1800,15 +1826,20 @@ public void onCompleted(UUID id, @Nullable Throwable error)
 
         assertThat(compositeInProgress).hasValue(1);
         assertThat(compositedCompleted).hasValue(0);
+        Mockito.verify(backgroundCompactions, times(1)).onInProgress(Mockito.any());
+        Mockito.verify(backgroundCompactions, times(0)).onCompleted(Mockito.any(), Mockito.any());
 
         // 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);
+        Mockito.verify(backgroundCompactions, times(1)).onInProgress(Mockito.any());
+        Mockito.verify(backgroundCompactions, times(1)).onCompleted(Mockito.any(), Mockito.any());
     }
 
+
     @Test
     public void testMaximalSelection()
     {