Merge branch 'cassandra-5.0' into trunk

* cassandra-5.0:
  Enhance CQLSSTableWriter to notify client on sstable production

CHANGES.txt

@@ -121,6 +121,7 @@ Merged from 5.0:
 Merged from 4.1:
  * Fix race condition in DecayingEstimatedHistogramReservoir during rescale (CASSANDRA-19365)
 Merged from 4.0:
+ * Enhance CQLSSTableWriter to notify clients on sstable production (CASSANDRA-19800)
  * Change the resolution of AbstractCommitLogService#lastSyncedAt to nanos to be aligned with later comparisons (CASSANDRA-20074)
  * Support UDTs and vectors as clustering keys in descending order (CASSANDRA-20050)
  * Fix CQL in snapshot's schema which did not contained UDTs used as reverse clustering columns (CASSANDRA-20036)

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

@@ -25,7 +25,9 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 import org.apache.cassandra.config.DatabaseDescriptor;
@@ -36,6 +38,7 @@ import org.apache.cassandra.db.partitions.PartitionUpdate;
 import org.apache.cassandra.db.rows.EncodingStats;
 import org.apache.cassandra.index.Index;
 import org.apache.cassandra.io.sstable.format.SSTableFormat;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.schema.TableMetadataRef;
 import org.apache.cassandra.service.ActiveRepairService;
@@ -52,6 +55,8 @@ abstract class AbstractSSTableSimpleWriter implements Closeable
     protected static final AtomicReference<SSTableId> id = new AtomicReference<>(SSTableIdFactory.instance.defaultBuilder().generator(Stream.empty()).get());
     protected boolean makeRangeAware = false;
     protected final Collection<Index.Group> indexGroups;
+    protected Consumer<Collection<SSTableReader>> sstableProducedListener;
+    protected boolean openSSTableOnProduced = false;
 
     protected AbstractSSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns)
     {
@@ -76,6 +81,32 @@ abstract class AbstractSSTableSimpleWriter implements Closeable
         this.indexGroups.add(indexGroup);
     }
 
+    protected void setSSTableProducedListener(Consumer<Collection<SSTableReader>> listener)
+    {
+        this.sstableProducedListener = Objects.requireNonNull(listener, "sstableProducedListener cannot be null");
+    }
+
+    protected void setShouldOpenProducedSSTable(boolean openSSTableOnProduced)
+    {
+        this.openSSTableOnProduced = openSSTableOnProduced;
+    }
+
+    /**
+     * Indicate whether the produced sstable should be opened or not.
+     */
+    protected boolean shouldOpenSSTables()
+    {
+        return openSSTableOnProduced;
+    }
+
+    protected void notifySSTableProduced(Collection<SSTableReader> sstables)
+    {
+        if (sstableProducedListener == null)
+            return;
+
+        sstableProducedListener.accept(sstables);
+    }
+
     protected SSTableTxnWriter createWriter(SSTable.Owner owner) throws IOException
     {
         SerializationHeader header = new SerializationHeader(true, metadata.get(), columns, EncodingStats.NO_STATS);

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

@@ -24,11 +24,13 @@ import java.nio.ByteBuffer;
 import java.nio.file.NoSuchFileException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.NavigableSet;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import com.google.common.base.Preconditions;
@@ -59,6 +61,7 @@ import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.exceptions.SyntaxException;
 import org.apache.cassandra.index.sai.StorageAttachedIndexGroup;
 import org.apache.cassandra.io.sstable.format.SSTableFormat;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.schema.KeyspaceMetadata;
 import org.apache.cassandra.schema.KeyspaceParams;
@@ -407,6 +410,8 @@ public class CQLSSTableWriter implements Closeable
         private boolean sorted = false;
         private long maxSSTableSizeInMiB = -1L;
         private boolean buildIndexes = true;
+        private Consumer<Collection<SSTableReader>> sstableProducedListener;
+        private boolean openSSTableOnProduced = false;
 
         protected Builder()
         {
@@ -620,6 +625,33 @@ public class CQLSSTableWriter implements Closeable
             return this;
         }
 
+        /**
+         * Set the listener to receive notifications on sstable produced
+         * <p>
+         * Note that if listener is registered, the sstables are opened into {@link SSTableReader}.
+         * The consumer is responsible for releasing the {@link SSTableReader}
+         *
+         * @param sstableProducedListener receives the produced sstables
+         * @return this builder
+         */
+        public Builder withSSTableProducedListener(Consumer<Collection<SSTableReader>> sstableProducedListener)
+        {
+            this.sstableProducedListener = sstableProducedListener;
+            return this;
+        }
+
+        /**
+         * Whether the produced sstable should be open or not.
+         * By default, the writer does not open the produced sstables
+         *
+         * @return this builder
+         */
+        public Builder openSSTableOnProduced()
+        {
+            this.openSSTableOnProduced = true;
+            return this;
+        }
+
         public CQLSSTableWriter build()
         {
             if (directory == null)
@@ -727,6 +759,11 @@ public class CQLSSTableWriter implements Closeable
                         writer.addIndexGroup(saiGroup);
                 }
 
+                if (sstableProducedListener != null)
+                    writer.setSSTableProducedListener(sstableProducedListener);
+
+                writer.setShouldOpenProducedSSTable(openSSTableOnProduced);
+
                 return new CQLSSTableWriter(writer, preparedModificationStatement, preparedModificationStatement.getBindVariables());
             }
         }

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

@@ -18,6 +18,7 @@
 package org.apache.cassandra.io.sstable;
 
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.concurrent.BlockingQueue;
@@ -34,6 +35,7 @@ import org.apache.cassandra.db.rows.EncodingStats;
 import org.apache.cassandra.db.rows.Row;
 import org.apache.cassandra.db.rows.SerializationHelper;
 import org.apache.cassandra.db.rows.UnfilteredSerializer;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.schema.TableMetadataRef;
 import org.apache.cassandra.utils.JVMStabilityInspector;
@@ -219,7 +221,8 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
                     {
                         for (Map.Entry<DecoratedKey, PartitionUpdate.Builder> entry : b.entrySet())
                             writer.append(entry.getValue().build().unfilteredIterator());
-                        writer.finish(false);
+                        Collection<SSTableReader> finished = writer.finish(shouldOpenSSTables());
+                        notifySSTableProduced(finished);
                     }
                 }
                 catch (Throwable e)

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

@@ -18,6 +18,7 @@
 package org.apache.cassandra.io.sstable;
 
 import java.io.IOException;
+import java.util.Collection;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
@@ -25,6 +26,7 @@ import com.google.common.base.Throwables;
 import org.apache.cassandra.db.DecoratedKey;
 import org.apache.cassandra.db.RegularAndStaticColumns;
 import org.apache.cassandra.db.partitions.PartitionUpdate;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.schema.TableMetadataRef;
 
@@ -137,8 +139,11 @@ class SSTableSimpleWriter extends AbstractSSTableSimpleWriter
     {
         try
         {
-            if (writer != null)
+            if (writer == null)
-                writer.finish(false);
+                return;
+
+            Collection<SSTableReader> finished = writer.finish(shouldOpenSSTables());
+            notifySSTableProduced(finished);
         }
         catch (Throwable t)
         {

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

@@ -20,6 +20,7 @@ package org.apache.cassandra.io.sstable;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
@@ -27,6 +28,7 @@ import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -60,6 +62,7 @@ import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
 import org.apache.cassandra.index.sai.utils.IndexIdentifier;
+import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.sstable.format.big.BigFormat;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.io.util.PathUtils;
@@ -1362,6 +1365,86 @@ public abstract class CQLSSTableWriterTest
         }
     }
 
+    @Test
+    public void testNotifySSTableFinishedForSorted() throws Exception
+    {
+        testNotifySSTableFinished(true, false);
+    }
+
+    @Test
+    public void testNotifySSTableFinishedForUnsorted() throws Exception
+    {
+        testNotifySSTableFinished(false, false);
+    }
+
+    @Test
+    public void testCloseSortedWriterOnFirstPorducedShouldStillResultInTwoSSTables() throws Exception
+    {
+        // Writing a new partition (and exceeding the size limit) lead to closing the current writer and buffering the last partition update.
+        // Since there is a last partition buffered, closing the sstable writer flushes to a new sstable.
+        // Therefore, even though the test closes the writer immediately, there are still 2 sstables produced.
+        testNotifySSTableFinished(true, true);
+    }
+
+    private void testNotifySSTableFinished(boolean sorted, boolean closeWriterOnFirstProduced) throws Exception
+    {
+        List<SSTableReader> produced = new ArrayList<>();
+        String schema = "CREATE TABLE " + qualifiedTable + " ("
+                        + "  k int PRIMARY KEY,"
+                        + "  v blob )";
+        CQLSSTableWriter.Builder builder = CQLSSTableWriter
+                                           .builder()
+                                           .inDirectory(dataDir)
+                                           .forTable(schema)
+                                           .using("INSERT INTO " + qualifiedTable +
+                                                  " (k, v) VALUES (?, text_as_blob(?))")
+                                           .withMaxSSTableSizeInMiB(1)
+                                           .openSSTableOnProduced()
+                                           .withSSTableProducedListener(produced::addAll);
+        if (sorted)
+        {
+            builder.sorted();
+        }
+        CQLSSTableWriter writer = builder.build();
+
+        int rowCount = 30_000;
+        // Max SSTable size is 1 MiB
+        // 30_000 rows should take 30_000 * (4 + 37) = 1.17 MiB > 1 MiB, i.e. producing 2 sstables
+        for (int i = 0; i < rowCount; i++)
+        {
+            writer.addRow(i, UUID.randomUUID().toString());
+            if (closeWriterOnFirstProduced && !produced.isEmpty())
+            {
+                // on closing writer, it flushes the last update to the new sstable
+                writer.close();
+                break;
+            }
+        }
+        // the assertion is only performed for sorted because unsorted writer writes asynchrously; avoid flakiness
+        if (!closeWriterOnFirstProduced && sorted)
+        {
+            // while writing, one sstable should be finished
+            assertEquals(1, produced.size());
+        }
+
+        if (!closeWriterOnFirstProduced)
+            writer.close();
+        // another sstable is finished on closing the writer
+        assertEquals(2, produced.size());
+
+        File[] dataFiles = dataDir.list(f -> f.name().endsWith(BigFormat.Components.DATA.type.repr));
+        assertNotNull(dataFiles);
+        assertEquals("The sorted writer should produce 2 sstables when max sstable size is configured",
+                     2, dataFiles.length);
+        Set<File> notifiedDataFileSet = produced.stream()
+                                                .map(sstable -> sstable.descriptor.fileFor(BigFormat.Components.DATA))
+                                                .collect(Collectors.toSet());
+        Set<File> listedDataFileSet = Arrays.stream(dataFiles)
+                                            .map(File::toCanonical)
+                                            .collect(Collectors.toSet());
+        assertEquals(notifiedDataFileSet, listedDataFileSet);
+    }
+
     @Test
     public void testMultipleWritersWithDistinctTables() throws IOException
     {