Merge branch 'cassandra-4.0' into cassandra-4.1

CHANGES.txt

@@ -6,6 +6,7 @@
  * Allow empty keystore_password in encryption_options (CASSANDRA-18778)
  * Skip ColumnFamilyStore#topPartitions initialization when client or tool mode (CASSANDRA-18697)
 Merged from 4.0:
+ * Support max SSTable size in sorted CQLSSTableWriter (CASSANDRA-18941)
  * Fix nodetool repair_admin summarize-pending command to not throw exception (CASSANDRA-19014)
  * Fix cassandra-stress in simplenative mode with prepared statements (CASSANDRA-18744)
  * Fix filtering system ks sstables for relocation on startup (CASSANDRA-18963)

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

@@ -58,6 +58,8 @@ import org.apache.cassandra.schema.*;
 import org.apache.cassandra.service.ClientState;
 import org.apache.cassandra.transport.ProtocolVersion;
 import org.apache.cassandra.utils.ByteBufferUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import static org.apache.cassandra.utils.Clock.Global.currentTimeMillis;
 
@@ -364,19 +366,22 @@ public class CQLSSTableWriter implements Closeable
      */
     public static class Builder
     {
-        private File directory;
+        private static final Logger logger = LoggerFactory.getLogger(Builder.class);
+        private static final long DEFAULT_BUFFER_SIZE_IN_MIB_FOR_UNSORTED = 128L;
+
 
         protected SSTableFormat.Type formatType = null;
 
-        private CreateTableStatement.Raw schemaStatement;
         private final List<CreateTypeStatement.Raw> typeStatements;
+        private CreateTableStatement.Raw schemaStatement;
         private ModificationStatement.Parsed modificationStatement;
         private IPartitioner partitioner;
-
+        private File directory;
         private boolean sorted = false;
-        private long bufferSizeInMiB = 128;
+        private long maxSSTableSizeInMiB = -1L;
 
-        protected Builder() {
+        protected Builder()
+        {
             this.typeStatements = new ArrayList<>();
         }
 
@@ -482,23 +487,42 @@ public class CQLSSTableWriter implements Closeable
             return this;
         }
 
+        /**
+         * Defines the maximum SSTable size in mebibytes when using the sorted writer.
+         * By default, i.e. not specified, there is no maximum size limit for the produced SSTable
+         *
+         * @param size the maximum sizein mebibytes of each individual SSTable allowed
+         * @return this builder
+         */
+        public Builder withMaxSSTableSizeInMiB(int size)
+        {
+            if (size <= 0)
+            {
+                logger.warn("A non-positive value for maximum SSTable size is specified, " +
+                            "which disables the size limiting effectively. Please supply a positive value in order " +
+                            "to enforce size limiting for the produced SSTables.");
+            }
+            this.maxSSTableSizeInMiB = size;
+            return this;
+        }
+
         /**
          * The size of the buffer to use.
          * <p>
          * This defines how much data will be buffered before being written as
-         * a new SSTable. This correspond roughly to the data size that will have the created
+         * a new SSTable. This corresponds roughly to the data size that will have the created
          * sstable.
          * <p>
          * The default is 128MiB, which should be reasonable for a 1GiB heap. If you experience
          * OOM while using the writer, you should lower this value.
          *
+         * @deprecated This method is deprecated in favor of the new withMaxSSTableSizeInMiB(int size)
          * @param size the size to use in MiB.
          * @return this builder.
          */
         public Builder withBufferSizeInMiB(int size)
         {
-            this.bufferSizeInMiB = size;
+            return withMaxSSTableSizeInMiB(size);
-            return this;
         }
 
         /**
@@ -506,12 +530,13 @@ public class CQLSSTableWriter implements Closeable
          * The size of the buffer to use.
          * <p>
          * This defines how much data will be buffered before being written as
-         * a new SSTable. This correspond roughly to the data size that will have the created
+         * a new SSTable. This corresponds roughly to the data size that will have the created
          * sstable.
          * <p>
          * The default is 128MiB, which should be reasonable for a 1GiB heap. If you experience
          * OOM while using the writer, you should lower this value.
          *
+         * @deprecated This method is deprecated in favor of the new withMaxSSTableSizeInMiB(int size)
          * @param size the size to use in MiB.
          * @return this builder.
          */
@@ -557,6 +582,13 @@ public class CQLSSTableWriter implements Closeable
             Preconditions.checkState(Sets.difference(SchemaConstants.LOCAL_SYSTEM_KEYSPACE_NAMES, Schema.instance.getKeyspaces()).isEmpty(),
                                      "Local keyspaces were not loaded. If this is running as a client, please make sure to add %s=true system property.",
                                      CassandraRelevantProperties.FORCE_LOAD_LOCAL_KEYSPACES.getKey());
+
+            // Assign the default max SSTable size if not defined in builder
+            if (isMaxSSTableSizeUnset())
+            {
+                maxSSTableSizeInMiB = sorted ? -1L : DEFAULT_BUFFER_SIZE_IN_MIB_FOR_UNSORTED;
+            }
+
             synchronized (CQLSSTableWriter.class)
             {
 
@@ -584,8 +616,8 @@ public class CQLSSTableWriter implements Closeable
 
                 TableMetadataRef ref = TableMetadataRef.forOfflineTools(tableMetadata);
                 AbstractSSTableSimpleWriter writer = sorted
-                                                     ? new SSTableSimpleWriter(directory, ref, preparedModificationStatement.updatedColumns())
+                                                     ? new SSTableSimpleWriter(directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB)
-                                                     : new SSTableSimpleUnsortedWriter(directory, ref, preparedModificationStatement.updatedColumns(), bufferSizeInMiB);
+                                                     : new SSTableSimpleUnsortedWriter(directory, ref, preparedModificationStatement.updatedColumns(), maxSSTableSizeInMiB);
 
                 if (formatType != null)
                     writer.setSSTableFormatType(formatType);
@@ -594,6 +626,11 @@ public class CQLSSTableWriter implements Closeable
             }
         }
 
+        private boolean isMaxSSTableSizeUnset()
+        {
+            return maxSSTableSizeInMiB <= 0;
+        }
+
         private Types createTypes(String keyspace)
         {
             Types.RawBuilder builder = Types.rawBuilder(keyspace);

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

@@ -44,6 +44,7 @@ public interface SSTableMultiWriter extends Transactional
 
     String getFilename();
     long getFilePointer();
+    long getOnDiskBytesWritten();
     TableId getTableId();
 
     static void abortOrDie(SSTableMultiWriter writer)

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

@@ -56,7 +56,7 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
     private static final Buffer SENTINEL = new Buffer();
 
     private Buffer buffer = new Buffer();
-    private final long bufferSize;
+    private final long maxSSTableSizeInBytes;
     private long currentSize;
 
     // Used to compute the row serialized size
@@ -66,15 +66,16 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
     private final BlockingQueue<Buffer> writeQueue = newBlockingQueue(0);
     private final DiskWriter diskWriter = new DiskWriter();
 
-    SSTableSimpleUnsortedWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long bufferSizeInMB)
+    SSTableSimpleUnsortedWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB)
     {
         super(directory, metadata, columns);
-        this.bufferSize = bufferSizeInMB * 1024L * 1024L;
+        this.maxSSTableSizeInBytes = maxSSTableSizeInMiB * 1024L * 1024L;
         this.header = new SerializationHeader(true, metadata.get(), columns, EncodingStats.NO_STATS);
         this.helper = new SerializationHelper(this.header);
         diskWriter.start();
     }
 
+    @Override
     PartitionUpdate.Builder getUpdateFor(DecoratedKey key)
     {
         assert key != null;
@@ -95,7 +96,7 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
         // Note that the accounting of a row is a bit inaccurate (it doesn't take some of the file format optimization into account)
         // and the maintaining of the bufferSize is in general not perfect. This has always been the case for this class but we should
         // improve that. In particular, what we count is closer to the serialized value, but it's debatable that it's the right thing
-        // to count since it will take a lot more space in memory and the bufferSize if first and foremost used to avoid OOM when
+        // to count since it will take a lot more space in memory and the bufferSize is first and foremost used to avoid OOM when
         // using this writer.
         currentSize += UnfilteredSerializer.serializer.serializedSize(row, helper, 0, formatType.info.getLatestVersion().correspondingMessagingVersion());
     }
@@ -104,7 +105,7 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
     {
         try
         {
-            if (currentSize > bufferSize)
+            if (currentSize > maxSSTableSizeInBytes)
                 sync();
         }
         catch (IOException e)

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

@@ -19,6 +19,7 @@ package org.apache.cassandra.io.sstable;
 
 import java.io.IOException;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 
 import org.apache.cassandra.db.*;
@@ -35,57 +36,114 @@ import org.apache.cassandra.schema.TableMetadataRef;
  * means that rows should be added by increasing md5 of the row key. This is
  * rarely possible and SSTableSimpleUnsortedWriter should most of the time be
  * prefered.
+ * <p>
+ * Optionally, the writer can be configured with a max SSTable size for SSTables.
+ * The output will be a series of SSTables that do not exceed a specified size.
+ * By default, all sorted data are written into a single SSTable.
  */
 class SSTableSimpleWriter extends AbstractSSTableSimpleWriter
 {
+    private final long maxSSTableSizeInBytes;
+
     protected DecoratedKey currentKey;
     protected PartitionUpdate.Builder update;
 
     private SSTableTxnWriter writer;
 
-    protected SSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns)
+    /**
+     * Create a SSTable writer for sorted input data.
+     * When a positive {@param maxSSTableSizeInMiB} is defined, the writer outputs a sequence of SSTables,
+     * whose sizes do not exceed the specified value.
+     *
+     * @param directory directory to store the sstable files
+     * @param metadata table metadata
+     * @param columns columns to update
+     * @param maxSSTableSizeInMiB defines the max SSTable size if the value is positive.
+     *                            Any non-positive value indicates the sstable size is unlimited.
+     */
+    protected SSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns, long maxSSTableSizeInMiB)
     {
         super(directory, metadata, columns);
+        this.maxSSTableSizeInBytes = maxSSTableSizeInMiB * 1024L * 1024L;
     }
 
-    private SSTableTxnWriter getOrCreateWriter() throws IOException
+    @Override
-    {
-        if (writer == null)
-            writer = createWriter();
-
-        return writer;
-    }
-
     PartitionUpdate.Builder getUpdateFor(DecoratedKey key) throws IOException
     {
-        assert key != null;
+        Preconditions.checkArgument(key != null, "Partition update cannot have null key");
 
-        // If that's not the current key, write the current one if necessary and create a new
+        // update for the first partition or a new partition
-        // update for the new key.
+        if (update == null || !key.equals(currentKey))
-        if (!key.equals(currentKey))
         {
+            // write the previous update if not absent
             if (update != null)
-                writePartition(update.build());
+                writePartition(update.build()); // might switch to a new sstable writer and reset currentSize
+
             currentKey = key;
             update = new PartitionUpdate.Builder(metadata.get(), currentKey, columns, 4);
         }
 
-        assert update != null;
+        Preconditions.checkState(update != null, "Partition update to write cannot be null");
         return update;
     }
 
+    @Override
     public void close()
+    {
+        writeLastPartitionUpdate(update);
+        maybeCloseWriter(writer);
+    }
+
+    /**
+     * Switch to a new writer when writer is absent or the file size has exceeded the configured max
+     */
+    private boolean shouldSwitchToNewWriter()
+    {
+        return writer == null || (maxSSTableSizeInBytes > 0 && writer.getOnDiskBytesWritten() > maxSSTableSizeInBytes);
+    }
+
+    /**
+     * Get the current writer, or create a new writer if needed, e.g. writer does not exist
+     * or writer has reached to the configured max size.
+     */
+    private SSTableTxnWriter getOrCreateWriter() throws IOException
+    {
+        if (shouldSwitchToNewWriter())
+        {
+            maybeCloseWriter(writer);
+            writer = createWriter();
+        }
+
+        return writer;
+    }
+
+    private void writeLastPartitionUpdate(PartitionUpdate.Builder update)
     {
         try
         {
             if (update != null)
                 writePartition(update.build());
+        }
+        catch (Throwable t)
+        {
+            Throwable e = writer == null ? t : writer.abort(t);
+            Throwables.throwIfUnchecked(e);
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void maybeCloseWriter(SSTableTxnWriter writer)
+    {
+        try
+        {
             if (writer != null)
                 writer.finish(false);
         }
         catch (Throwable t)
         {
-            throw Throwables.propagate(writer == null ? t : writer.abort(t));
+            Throwable e = writer.abort(t);
+            Throwables.throwIfUnchecked(e);
+            throw new RuntimeException(e);
         }
     }
 

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

@@ -67,6 +67,17 @@ public class SSTableTxnWriter extends Transactional.AbstractTransactional implem
         return writer.getFilePointer();
     }
 
+    /**
+     * Get the amount of data written to disk. Unlike {@link #getFilePointer()}, which returns the position in the
+     * _uncompressed_ data, this method returns the actual file pointer position of the on disk file.
+     *
+     * @return the amount of data already written to disk
+     */
+    public long getOnDiskBytesWritten()
+    {
+        return writer.getOnDiskBytesWritten();
+    }
+
     protected Throwable doCommit(Throwable accumulate)
     {
         return writer.commit(txn.commit(accumulate));

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

@@ -80,6 +80,12 @@ public class SimpleSSTableMultiWriter implements SSTableMultiWriter
         return writer.getFilePointer();
     }
 
+    @Override
+    public long getOnDiskBytesWritten()
+    {
+        return writer.getEstimatedOnDiskBytesWritten();
+    }
+
     public TableId getTableId()
     {
         return writer.metadata().id;

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

@@ -164,6 +164,12 @@ public class RangeAwareSSTableWriter implements SSTableMultiWriter
        return currentWriter != null ? currentWriter.getFilePointer() : 0L;
     }
 
+    @Override
+    public long getOnDiskBytesWritten()
+    {
+        return currentWriter != null ? currentWriter.getOnDiskBytesWritten() : 0L;
+    }
+
     @Override
     public TableId getTableId()
     {

src/java/org/apache/cassandra/io/sstable/format/big/BigTableZeroCopyWriter.java

@@ -167,6 +167,12 @@ public class BigTableZeroCopyWriter extends SSTable implements SSTableMultiWrite
         return 0;
     }
 
+    @Override
+    public long getOnDiskBytesWritten()
+    {
+        return 0;
+    }
+
     @Override
     public TableId getTableId()
     {

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

@@ -59,6 +59,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.psjava.util.AssertStatus.assertNotNull;
 
 public class CQLSSTableWriterTest
 {
@@ -1116,6 +1117,80 @@ public class CQLSSTableWriterTest
         assertEquals(0, filtered.size());
     }
 
+    @Test
+    public void testWriteWithSorted() throws Exception
+    {
+        String schema = "CREATE TABLE " + qualifiedTable + " ("
+                        + "  k int PRIMARY KEY,"
+                        + "  v blob )";
+        CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                  .inDirectory(dataDir)
+                                                  .forTable(schema)
+                                                  .using("INSERT INTO " + qualifiedTable +
+                                                         " (k, v) VALUES (?, textAsBlob(?))" )
+                                                  .sorted()
+                                                  .build();
+        int rowCount = 10_000;
+        for (int i = 0; i < rowCount; i++)
+        {
+            writer.addRow(i, UUID.randomUUID().toString());
+        }
+        writer.close();
+        loadSSTables(dataDir, keyspace);
+
+        UntypedResultSet resultSet = QueryProcessor.executeInternal("SELECT * FROM " + qualifiedTable);
+        assertEquals(rowCount, resultSet.size());
+        Iterator<UntypedResultSet.Row> iter = resultSet.iterator();
+        for (int i = 0; i < rowCount; i++)
+        {
+            UntypedResultSet.Row row = iter.next();
+            assertEquals(i, row.getInt("k"));
+        }
+    }
+
+    @Test
+    public void testWriteWithSortedAndMaxSize() throws Exception
+    {
+        String schema = "CREATE TABLE " + qualifiedTable + " ("
+                        + "  k int PRIMARY KEY,"
+                        + "  v blob )";
+        CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                  .inDirectory(dataDir)
+                                                  .forTable(schema)
+                                                  .using("INSERT INTO " + qualifiedTable +
+                                                         " (k, v) VALUES (?, textAsBlob(?))")
+                                                  .sorted()
+                                                  .withMaxSSTableSizeInMiB(1)
+                                                  .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
+        for (int i = 0; i < rowCount; i++)
+        {
+            writer.addRow(i, UUID.randomUUID().toString());
+        }
+        writer.close();
+
+        File[] dataFiles = dataDir.list(f -> f.name().endsWith(Component.DATA.type.repr));
+        assertNotNull(dataFiles);
+        assertEquals("The sorted writer should produce 2 sstables when max sstable size is configured",
+                     2, dataFiles.length);
+        long closeTo1MiBFileSize = Math.max(dataFiles[0].length(), dataFiles[1].length());
+        assertTrue("The file size should be close to 1MiB (with at most 50KiB error rate for the test)",
+                   Math.abs(1024 * 1024 - closeTo1MiBFileSize) < 50 * 1024);
+
+        loadSSTables(dataDir, keyspace);
+
+        UntypedResultSet resultSet = QueryProcessor.executeInternal("SELECT * FROM " + qualifiedTable);
+        assertEquals(rowCount, resultSet.size());
+        Iterator<UntypedResultSet.Row> iter = resultSet.iterator();
+        for (int i = 0; i < rowCount; i++)
+        {
+            UntypedResultSet.Row row = iter.next();
+            assertEquals(i, row.getInt("k"));
+        }
+    }
+
     private static void loadSSTables(File dataDir, String ks) throws ExecutionException, InterruptedException
     {
         SSTableLoader loader = new SSTableLoader(dataDir, new SSTableLoader.Client()

tools/stress/src/org/apache/cassandra/io/sstable/StressCQLSSTableWriter.java

@@ -578,7 +578,7 @@ public class StressCQLSSTableWriter implements Closeable
 
                 UpdateStatement preparedInsert = prepareInsert();
                 AbstractSSTableSimpleWriter writer = sorted
-                                                     ? new SSTableSimpleWriter(cfs.getDirectories().getDirectoryForNewSSTables(), cfs.metadata, preparedInsert.updatedColumns())
+                                                     ? new SSTableSimpleWriter(cfs.getDirectories().getDirectoryForNewSSTables(), cfs.metadata, preparedInsert.updatedColumns(), -1)
                                                      : new SSTableSimpleUnsortedWriter(cfs.getDirectories().getDirectoryForNewSSTables(), cfs.metadata, preparedInsert.updatedColumns(), bufferSizeInMiB);
 
                 if (formatType != null)