diff --git a/CHANGES.txt b/CHANGES.txt index 0bf763ff26..94e0a333c3 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -4,6 +4,7 @@ * Equality check for Paxos.Electorate should not depend on collection types (CASSANDRA-19935) * Fix race condition in DecayingEstimatedHistogramReservoir during rescale (CASSANDRA-19365) Merged from 4.0: + * Enhance CQLSSTableWriter to notify clients on sstable production (CASSANDRA-19800) * Fix gossip issue with gossip-only and bootstrapping nodes missing DC/Rack/Host ID endpoint state (CASSANDRA-19983) * Change the resolution of AbstractCommitLogService#lastSyncedAt to nanos to be aligned with later comparisons (CASSANDRA-20074) * Backport of CASSANDRA-17812: Rate-limit new client connection auth setup to avoid overwhelming bcrypt (CASSANDRA-20057) diff --git a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java index 49ceb764de..1e0700f635 100644 --- a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java @@ -23,14 +23,18 @@ import java.io.Closeable; import java.nio.ByteBuffer; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Collection; import java.util.Collections; +import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.cassandra.db.*; import org.apache.cassandra.db.partitions.PartitionUpdate; import org.apache.cassandra.db.rows.EncodingStats; 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; @@ -46,6 +50,8 @@ abstract class AbstractSSTableSimpleWriter implements Closeable protected SSTableFormat.Type formatType = SSTableFormat.Type.current(); protected static final AtomicReference id = new AtomicReference<>(SSTableIdFactory.instance.defaultBuilder().generator(Stream.empty()).get()); protected boolean makeRangeAware = false; + protected Consumer> sstableProducedListener; + protected boolean openSSTableOnProduced = false; protected AbstractSSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns) { @@ -64,6 +70,32 @@ abstract class AbstractSSTableSimpleWriter implements Closeable this.makeRangeAware = makeRangeAware; } + protected void setSSTableProducedListener(Consumer> listener) + { + this.sstableProducedListener = Objects.requireNonNull(listener); + } + + 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() 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 c2602a8586..e8e025ba65 100644 --- a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java @@ -23,10 +23,12 @@ import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.Collection; 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; @@ -53,6 +55,7 @@ import org.apache.cassandra.dht.Murmur3Partitioner; import org.apache.cassandra.exceptions.InvalidRequestException; import org.apache.cassandra.exceptions.SyntaxException; 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.*; import org.apache.cassandra.service.ClientState; @@ -379,6 +382,8 @@ public class CQLSSTableWriter implements Closeable private File directory; private boolean sorted = false; private long maxSSTableSizeInMiB = -1L; + private Consumer> sstableProducedListener; + private boolean openSSTableOnProduced = false; protected Builder() { @@ -569,6 +574,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; + } + @SuppressWarnings("resource") public CQLSSTableWriter build() { @@ -622,6 +654,11 @@ public class CQLSSTableWriter implements Closeable if (formatType != null) writer.setSSTableFormatType(formatType); + 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 160debaab0..e52f09862d 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; @@ -212,11 +214,12 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter if (b == SENTINEL) return; - try (SSTableTxnWriter writer = createWriter()) + try (SSTableTxnWriter writer = createWriter()) { 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 921c117a57..009951cdc8 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java @@ -18,12 +18,14 @@ 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; import org.apache.cassandra.db.*; 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; @@ -136,8 +138,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 ffa8c60b23..0dd2e05da3 100644 --- a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java +++ b/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java @@ -30,6 +30,7 @@ import java.util.stream.StreamSupport; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import org.apache.cassandra.io.sstable.format.SSTableReader; import org.apache.cassandra.io.util.File; import org.junit.Before; import org.junit.BeforeClass; @@ -1191,6 +1192,86 @@ public 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) leads 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 text )"; + CQLSSTableWriter.Builder builder = CQLSSTableWriter + .builder() + .inDirectory(dataDir) + .forTable(schema) + .using("INSERT INTO " + qualifiedTable + + " (k, v) VALUES (?, ?)") + .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(Component.DATA.name())); + 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(Component.DATA)) + .collect(Collectors.toSet()); + Set listedDataFileSet = Arrays.stream(dataFiles) + .map(File::toCanonical) + .collect(Collectors.toSet()); + assertEquals(notifiedDataFileSet, listedDataFileSet); + } + private static void loadSSTables(File dataDir, String ks) throws ExecutionException, InterruptedException { SSTableLoader loader = new SSTableLoader(dataDir, new SSTableLoader.Client()