CNDB-14725: add SSTableFlushObserver#onSSTableWriterSwitched to flush SAI segment builder for written shards without waiting for entire transaction to complete (#1859)

This reduces memory usage during sharded compaction.

### What is the issue

riptano/cndb#14725 OOM during SAI compaction
with large num of shards

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

Flush segment builder when sstable writer is switched to free memory
without waiting full compaction to complete

Author: jasonstack

src/java/org/apache/cassandra/index/sai/disk/PerIndexWriter.java

@@ -46,6 +46,12 @@ public interface PerIndexWriter
      */
     void complete(Stopwatch stopwatch) throws IOException;
 
+    /**
+     * Called when current sstable writer is switched during sharded compaction to free any in-memory resources associated
+     * with the sstable for current index without waiting for full transaction to complete
+     */
+    void onSSTableWriterSwitched(Stopwatch stopwatch) throws IOException;
+
     /**
      * Aborts accumulating data. Allows to clean up resources on error.
      * 

src/java/org/apache/cassandra/index/sai/disk/StorageAttachedIndexWriter.java

@@ -184,6 +184,29 @@ public void staticRow(Row staticRow)
         }
     }
 
+    @Override
+    public void onSSTableWriterSwitched()
+    {
+        if (aborted) return;
+
+        try
+        {
+            long start = Clock.Global.nanoTime();
+            for (PerIndexWriter w : perIndexWriters)
+            {
+                w.onSSTableWriterSwitched(stopwatch);
+            }
+            totalTimeSpent += (Clock.Global.nanoTime() - start);
+        }
+        catch (Throwable t)
+        {
+            logger.error(indexDescriptor.logMessage("Failed to flush segment on sstable writer switched"), t);
+            abort(t, true);
+            // fail compaction task or index build task if SAI failed
+            throw Throwables.unchecked(t);
+        }
+    }
+
     @Override
     public void complete(SSTable sstable)
     {

src/java/org/apache/cassandra/index/sai/disk/v1/MemtableIndexWriter.java

@@ -101,6 +101,13 @@ public void abort(Throwable cause)
         perIndexComponents.forceDeleteAllComponents();
     }
 
+    @Override
+    public void onSSTableWriterSwitched(Stopwatch stopwatch) throws IOException
+    {
+        // no-op for memtable index where all terms are already inside memory index, we can't get rid of memory index
+        // until full flush are completed
+    }
+
     @Override
     public void complete(Stopwatch stopwatch) throws IOException
     {

src/java/org/apache/cassandra/index/sai/disk/v1/SSTableIndexWriter.java

@@ -142,6 +142,18 @@ public void addRow(PrimaryKey key, Row row, long sstableRowId) throws IOExceptio
 
     }
 
+    @Override
+    public void onSSTableWriterSwitched(Stopwatch stopwatch) throws IOException
+    {
+        if (maybeAbort())
+            return;
+
+        boolean emptySegment = currentBuilder == null || currentBuilder.isEmpty();
+        logger.debug("Flushing index {} with {}buffered data on sstable writer switched...", indexContext.getIndexName(), emptySegment ? "no " : "");
+        if (!emptySegment)
+            flushSegment();
+    }
+
     @Override
     public void complete(Stopwatch stopwatch) throws IOException
     {
@@ -322,8 +334,9 @@ private void flushSegment() throws IOException
                 if (indexContext.getIndexMetrics() != null)
                     indexContext.getIndexMetrics().compactionSegmentBytesPerSecond.update((long)(segmentBytes / flushMillis * 1000.0));
 
-                logger.debug("Flushed segment with {} cells for a total of {} in {} ms",
-                             (long) rowCount, FBUtilities.prettyPrintMemory((long) segmentBytes), flushMillis);
+                logger.debug("Flushed segment with {} cells for a total of {} in {} ms for index {} with starting row id {} for sstable {}",
+                             (long) rowCount, FBUtilities.prettyPrintMemory((long) segmentBytes), flushMillis, indexContext.getIndexName(),
+                             segmentMetadata.minSSTableRowId, perIndexComponents.descriptor());
             }
 
             // Builder memory is released against the limiter at the conclusion of a successful
@@ -402,9 +415,10 @@ else if (indexContext.isLiteral())
         }
 
         long globalBytesUsed = limiter.increment(builder.totalBytesAllocated());
-        logger.debug("Created new segment builder while flushing SSTable {}. Global segment memory usage now at {}",
+        logger.debug("Created new segment builder while flushing SSTable {}. Global segment memory usage now at {} with {} active segment builders",
                      perIndexComponents.descriptor(),
-                     FBUtilities.prettyPrintMemory(globalBytesUsed));
+                     FBUtilities.prettyPrintMemory(globalBytesUsed),
+                     SegmentBuilder.ACTIVE_BUILDER_COUNT.get() - 1);
 
         return builder;
     }

src/java/org/apache/cassandra/io/sstable/SSTableFlushObserver.java

@@ -67,6 +67,14 @@ public interface SSTableFlushObserver
      */
     void complete(SSTable ssTable);
 
+    /**
+     * Called when current sstable writer is switched during sharded compaction to free any in-memory resources associated
+     * with the sstable without waiting for full transaction to complete
+     */
+    default void onSSTableWriterSwitched()
+    {
+    }
+
     /**
      * Clean up resources on error. There should be no side effects if called multiple times.
      */

src/java/org/apache/cassandra/io/sstable/SSTableRewriter.java

@@ -306,6 +306,7 @@ public void switchWriter(SSTableWriter newWriter)
 
         currentlyOpenedEarlyAt = 0;
         bytesWritten += writer.getFilePointer();
+        writer.onSSTableWriterSwitched();
         writer = newWriter;
     }
 

src/java/org/apache/cassandra/io/sstable/format/SSTableWriter.java

@@ -337,6 +337,12 @@ public final void abort()
         }
     }
 
+    // notify sstable flush observer about sstable writer switched
+    public final void onSSTableWriterSwitched()
+    {
+        observers.forEach(SSTableFlushObserver::onSSTableWriterSwitched);
+    }
+
     protected Map<MetadataType, MetadataComponent> finalizeMetadata()
     {
         return metadataCollector.finalizeMetadata(getPartitioner().getClass().getCanonicalName(),

test/unit/org/apache/cassandra/index/sai/functional/CompactionTest.java

@@ -22,6 +22,10 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 
 import javax.management.InstanceNotFoundException;
 
@@ -46,6 +50,7 @@
 import org.apache.cassandra.index.IndexNotAvailableException;
 import org.apache.cassandra.index.sai.IndexContext;
 import org.apache.cassandra.index.sai.disk.v1.SSTableIndexWriter;
+import org.apache.cassandra.index.sai.disk.v1.SegmentBuilder;
 import org.apache.cassandra.index.sai.metrics.AbstractMetricsTest;
 import org.apache.cassandra.inject.ActionBuilder;
 import org.apache.cassandra.inject.Expression;
@@ -65,12 +70,13 @@
 import org.apache.cassandra.utils.TimeUUID;
 import org.apache.cassandra.utils.concurrent.Refs;
 
-import static org.apache.cassandra.utils.TimeUUID.Generator.nextTimeUUID;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.apache.cassandra.utils.TimeUUID.Generator.nextTimeUUID;
 
 public class CompactionTest extends AbstractMetricsTest
 {
@@ -403,4 +409,53 @@ protected long getDiskUsage()
     {
         return (long) getMetricValue(objectNameNoIndex("DiskUsedBytes", KEYSPACE, currentTable(), "IndexGroupMetrics"));
     }
+
+    @Test
+    public void testSegmentBuilderFlushWithShardedCompaction() throws Throwable
+    {
+        int shards = 64;
+        String createTable = "CREATE TABLE %s (id1 TEXT PRIMARY KEY, v1 INT, v2 TEXT) WITH compaction = " +
+                             "{'class' : 'UnifiedCompactionStrategy', 'enabled' : false, 'base_shard_count': " + shards + ", 'min_sstable_size': '1KiB' }";
+        createTable(createTable);
+        createIndex(String.format(CREATE_INDEX_TEMPLATE, "v1"));
+        createIndex(String.format(CREATE_INDEX_TEMPLATE, "v2"));
+        disableCompaction(keyspace(), currentTable());
+
+        int rowsPerSSTable = 2000;
+        int numSSTables = 4;
+        int key = 0;
+        for (int s = 0; s < numSSTables; s++)
+        {
+            for (int i = 0; i < rowsPerSSTable; i++)
+            {
+                execute("INSERT INTO %s (id1, v1, v2) VALUES (?, 0, '01e2wefnewirui32e21e21wre')", Integer.toString(key++));
+            }
+            flush();
+        }
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try
+        {
+            Future<?> future = executor.submit(() -> {
+                getCurrentColumnFamilyStore().forceMajorCompaction(false, 1);
+                waitForCompactions();
+            });
+
+            // verify that it's not accumulating segment builders
+            while (!future.isDone())
+            {
+                // ACTIVE_BUILDER_COUNT starts from 1. There are 2 segments for 2 indexes
+                assertThat(SegmentBuilder.ACTIVE_BUILDER_COUNT.get()).isGreaterThanOrEqualTo(1).isLessThanOrEqualTo(3);
+            }
+            future.get(30, TimeUnit.SECONDS);
+
+            // verify results are sharded
+            assertThat(getCurrentColumnFamilyStore().getLiveSSTables()).hasSize(shards);
+        }
+        finally
+        {
+            executor.shutdown();
+            executor.awaitTermination(30, TimeUnit.SECONDS);
+        }
+    }
 }

test/unit/org/apache/cassandra/io/sstable/SSTableFlushObserverTest.java

@@ -168,14 +168,15 @@ public void testFlushObserver() throws Exception
                                                    BufferCell.live(getColumn(cfm, "height"), now, LongType.instance.decompose(178L))));
 
                 writer.append(new RowIterator(cfm, key, Collections.singletonList(buildRow(expected.get(key)))));
-
+                writer.onSSTableWriterSwitched();
                 reader = writer.finish(true, null);
             }
             finally
             {
                 FileUtils.closeQuietly(writer);
             }
 
+            Assert.assertTrue(observer.isWriterSwitched);
             Assert.assertTrue(observer.isComplete);
             Assert.assertEquals(expected.size(), observer.rows.size());
 
@@ -269,6 +270,7 @@ private static class FlushObserver implements SSTableFlushObserver
         private boolean beginCalled;
         private boolean failOnBegin;
         private boolean abortCalled;
+        private boolean isWriterSwitched;
 
         @Override
         public void begin()
@@ -303,6 +305,12 @@ public void staticRow(Row staticRow)
             staticRow.forEach((c) -> staticRows.put(currentKey, (Cell<?>) c));
         }
 
+        @Override
+        public void onSSTableWriterSwitched()
+        {
+            isWriterSwitched = true;
+        }
+
         @Override
         public void abort(Throwable accumulate)
         {