diff --git a/CHANGES.txt b/CHANGES.txt index 274e42c19a..8aab7a4e8c 100644 --- a/CHANGES.txt +++ b/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) diff --git a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java index 8d8543d25a..c12a76bc84 100644 --- a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java +++ b/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 id = new AtomicReference<>(SSTableIdFactory.instance.defaultBuilder().generator(Stream.empty()).get()); protected boolean makeRangeAware = false; protected final Collection indexGroups; + protected Consumer> 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> 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 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); diff --git a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java index cd2f7c79e8..8bb882071e 100644 --- a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java +++ b/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> 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 + *

+ * 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> 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()); } } diff --git a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java index ce927d9947..cf406af4c4 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java +++ b/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 entry : b.entrySet()) writer.append(entry.getValue().build().unfilteredIterator()); - writer.finish(false); + Collection finished = writer.finish(shouldOpenSSTables()); + notifySSTableProduced(finished); } } catch (Throwable e) diff --git a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java index ce423f2ec5..bd0ba94cce 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java +++ b/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) - writer.finish(false); + if (writer == null) + return; + + Collection finished = writer.finish(shouldOpenSSTables()); + notifySSTableProduced(finished); } catch (Throwable t) { diff --git a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java b/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java index b566e4e0fd..45de365a6c 100644 --- a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java +++ b/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 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 notifiedDataFileSet = produced.stream() + .map(sstable -> sstable.descriptor.fileFor(BigFormat.Components.DATA)) + .collect(Collectors.toSet()); + Set listedDataFileSet = Arrays.stream(dataFiles) + .map(File::toCanonical) + .collect(Collectors.toSet()); + assertEquals(notifiedDataFileSet, listedDataFileSet); + } + @Test public void testMultipleWritersWithDistinctTables() throws IOException {