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.

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

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

Author: pranavshenoy

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

@@ -41,6 +41,12 @@ public interface PerColumnIndexWriter
      */
     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.
      * <p> 

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

@@ -164,6 +164,25 @@ public void staticRow(Row staticRow)
         }
     }
 
+    @Override
+    public void onSSTableWriterSwitched()
+    {
+        if (aborted) return;
+
+        try
+        {
+            for (PerColumnIndexWriter w : perIndexWriters)
+            {
+                w.onSSTableWriterSwitched(stopwatch);
+            }
+        }
+        catch (Throwable t)
+        {
+            logger.error(indexDescriptor.logMessage("Failed to flush segment on sstable writer switched"), t);
+            abort(t, true);
+        }
+    }
+
     @Override
     public void complete()
     {

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

@@ -166,6 +166,13 @@ public void complete(Stopwatch stopwatch) throws IOException
         }
     }
 
+    @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
+    }
+
     private long flush(MemtableTermsIterator terms) throws IOException
     {
         SegmentWriter writer = indexTermType.isLiteral() ? new LiteralIndexWriter(indexDescriptor, indexIdentifier)

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

@@ -100,6 +100,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(index.identifier().logMessage("Flushing index with {}buffered data on SSTable writer switched..."), emptySegment ? "no " : "");
+        if (!emptySegment)
+            flushSegment();
+    }
+
     @Override
     public void complete(Stopwatch stopwatch) throws IOException
     {

src/java/org/apache/cassandra/index/sai/disk/v1/segment/SegmentBuilder.java

@@ -46,7 +46,7 @@ public abstract class SegmentBuilder
     public static final long LAST_VALID_SEGMENT_ROW_ID = (Integer.MAX_VALUE / 2) - 1L;
     private static long testLastValidSegmentRowId = -1;
 
-    /** The number of column indexes being built globally. (Starts at one to avoid divide by zero.) */
+    /** The number of column indexes being built globally. */
     private static final AtomicInteger ACTIVE_BUILDER_COUNT = new AtomicInteger(0);
 
     /** Minimum flush size, dynamically updated as segment builds are started and completed/aborted. */

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

@@ -68,6 +68,12 @@ public interface SSTableFlushObserver
      */
     void complete();
 
+    /**
+     * 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

@@ -282,6 +282,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

@@ -282,6 +282,12 @@ public final void prepareToCommit()
         txnProxy.prepareToCommit();
     }
 
+    // notify sstable flush observer about sstable writer switched
+    public final void onSSTableWriterSwitched()
+    {
+        observers.forEach(SSTableFlushObserver::onSSTableWriterSwitched);
+    }
+
     public final Throwable commit(Throwable accumulate)
     {
         try

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 com.google.common.collect.Lists;
 import org.junit.Assert;
@@ -42,6 +46,7 @@
 import org.apache.cassandra.index.sai.SAITester;
 import org.apache.cassandra.index.sai.utils.IndexIdentifier;
 import org.apache.cassandra.index.sai.disk.v1.SSTableIndexWriter;
+import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;
 import org.apache.cassandra.index.sai.utils.IndexTermType;
 import org.apache.cassandra.inject.ActionBuilder;
 import org.apache.cassandra.inject.Expression;
@@ -57,10 +62,12 @@
 import org.apache.cassandra.service.ActiveRepairService;
 import org.apache.cassandra.streaming.PreviewKind;
 import org.apache.cassandra.utils.ByteBufferUtil;
+import org.apache.cassandra.utils.FBUtilities;
 import org.apache.cassandra.utils.Throwables;
 import org.apache.cassandra.utils.TimeUUID;
 import org.apache.cassandra.utils.concurrent.Refs;
 
+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.assertNull;
@@ -352,4 +359,53 @@ public void testConcurrentIndexDropWithCompaction() throws Throwable
         .isInstanceOf(InvalidQueryException.class)
         .hasMessage(StatementRestrictions.REQUIRES_ALLOW_FILTERING_MESSAGE);
     }
+
+    @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.getActiveBuilderCount()).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

@@ -32,10 +32,12 @@
 import org.apache.commons.lang3.tuple.ImmutableTriple;
 import org.apache.commons.lang3.tuple.Triple;
 import org.junit.Assert;
+import static org.junit.Assert.assertFalse;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
+
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.db.Clustering;
 import org.apache.cassandra.db.DecoratedKey;
@@ -73,6 +75,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+
 public class SSTableFlushObserverTest
 {
 
@@ -167,14 +170,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);
             }
             finally
             {
                 FileUtils.closeQuietly(writer);
             }
-
+            
+            Assert.assertTrue(observer.isWriterSwitched);
             Assert.assertTrue(observer.isComplete);
             Assert.assertEquals(expected.size(), observer.rows.size());
 
@@ -265,6 +269,7 @@ private static class FlushObserver implements SSTableFlushObserver
         private boolean beginCalled;
         private boolean failOnBegin;
         private boolean abortCalled;
+        private boolean isWriterSwitched;
 
         @Override
         public void begin()
@@ -274,6 +279,12 @@ public void begin()
                 throw new RuntimeException("Failed to initialize");
         }
 
+        @Override
+        public void onSSTableWriterSwitched()
+        {
+            isWriterSwitched = true;
+        }
+
         @Override
         public void startPartition(DecoratedKey key, long dataPosition, long indexPosition)
         {