From 659558c980c67a80287ca7ccdfc8a70b1a56b7e2 Mon Sep 17 00:00:00 2001 From: Yifan Cai Date: Tue, 12 Nov 2024 15:56:06 -0800 Subject: [PATCH] Enhance CQLSSTableWriter to notify client on sstable production Patch by Yifan Cai; Reviewed by Francisco Guerrero for CASSANDRA-19800 --- CHANGES.txt | 1 + .../sstable/AbstractSSTableSimpleWriter.java | 31 +++++++ .../io/sstable/CQLSSTableWriter.java | 38 +++++++- .../sstable/SSTableSimpleUnsortedWriter.java | 7 +- .../io/sstable/SSTableSimpleWriter.java | 9 +- .../io/sstable/CQLSSTableWriterTest.java | 89 ++++++++++++++++++- 6 files changed, 168 insertions(+), 7 deletions(-) diff --git a/CHANGES.txt b/CHANGES.txt index 1e25031aeb..23dcc154f8 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -3,6 +3,7 @@ 4.0.15 + * 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) * Backport of CASSANDRA-17812: Rate-limit new client connection auth setup to avoid overwhelming bcrypt (CASSANDRA-20057) * Support UDTs and vectors as clustering keys in descending order (CASSANDRA-20050) diff --git a/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java b/src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java index 4eaf1fe64a..74ebc00b71 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.IOException; import java.io.Closeable; import java.nio.ByteBuffer; import java.util.Collections; +import java.util.Collection; import java.util.HashSet; +import java.util.Objects; import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; 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.schema.TableMetadataRef; import org.apache.cassandra.service.ActiveRepairService; @@ -45,6 +49,8 @@ abstract class AbstractSSTableSimpleWriter implements Closeable protected SSTableFormat.Type formatType = SSTableFormat.Type.current(); protected static AtomicInteger generation = new AtomicInteger(0); protected boolean makeRangeAware = false; + protected Consumer> sstableProducedListener; + protected boolean openSSTableOnProduced = false; protected AbstractSSTableSimpleWriter(File directory, TableMetadataRef metadata, RegularAndStaticColumns columns) { @@ -63,6 +69,31 @@ 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() { diff --git a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java index 4c299d071e..3b20b75651 100644 --- a/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java @@ -24,10 +24,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.SortedSet; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import java.util.stream.Collectors; import org.apache.cassandra.config.DatabaseDescriptor; @@ -50,6 +52,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.schema.*; import org.apache.cassandra.service.ClientState; import org.apache.cassandra.transport.ProtocolVersion; @@ -355,9 +358,10 @@ public class CQLSSTableWriter implements Closeable private CreateTableStatement.Raw schemaStatement; private ModificationStatement.Parsed insertStatement; private IPartitioner partitioner; - private boolean sorted = false; private long maxSSTableSizeInMiB = -1L; + private Consumer> sstableProducedListener; + private boolean openSSTableOnProduced = false; protected Builder() { this.typeStatements = new ArrayList<>(); @@ -525,6 +529,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() { @@ -581,6 +612,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, preparedInsert, preparedInsert.getBindVariables()); } } diff --git a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java index 87e87e95c2..457621b65b 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java @@ -19,6 +19,7 @@ package org.apache.cassandra.io.sstable; import java.io.File; 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.SerializationHelper; import org.apache.cassandra.db.rows.UnfilteredSerializer; import org.apache.cassandra.db.partitions.PartitionUpdate; +import org.apache.cassandra.io.sstable.format.SSTableReader; import org.apache.cassandra.schema.TableMetadataRef; import org.apache.cassandra.utils.JVMStabilityInspector; @@ -208,11 +210,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 1cbff412ad..cd51e1ebb4 100644 --- a/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java +++ b/src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java @@ -19,12 +19,14 @@ package org.apache.cassandra.io.sstable; import java.io.File; 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.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 13b6e2d927..8295ac2468 100644 --- a/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java +++ b/test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java @@ -25,13 +25,12 @@ import java.util.*; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; -import java.util.stream.Stream; import java.util.stream.StreamSupport; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.io.Files; +import org.apache.cassandra.io.sstable.format.SSTableReader; import org.junit.Before; import org.junit.BeforeClass; import org.junit.Rule; @@ -919,6 +918,92 @@ 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.listFiles((file, name) -> 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 -> new File(sstable.descriptor.filenameFor(Component.DATA))) + .collect(Collectors.toSet()); + Set listedDataFileSet = Arrays.stream(dataFiles) + .map(f -> { + try { + return f.getCanonicalFile(); + } catch (IOException e) { + throw new RuntimeException(e); + } + }) + .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()