natas_hw
/
cassandra
mirror of https://github.com/apache/cassandra
0
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Enhance CQLSSTableWriter to notify client on sstable production 

Patch by Yifan Cai; Reviewed by Francisco Guerrero for CASSANDRA-19800
This commit is contained in:
Yifan Cai 2024-11-12 15:56:06 -08:00
parent 426eebb513
commit 659558c980
6 changed files with 168 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGES.txt
View File

@ -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)

31
src/java/org/apache/cassandra/io/sstable/AbstractSSTableSimpleWriter.java
View File

@ -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<Collection<SSTableReader>> 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<Collection<SSTableReader>> 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<SSTableReader> sstables)
{
if (sstableProducedListener == null)
return;
sstableProducedListener.accept(sstables);
}
protected SSTableTxnWriter createWriter()
{

38
src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java
View File

@ -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<Collection<SSTableReader>> 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
* <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;
}
@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());
}
}

7
src/java/org/apache/cassandra/io/sstable/SSTableSimpleUnsortedWriter.java
View File

@ -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<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)

9
src/java/org/apache/cassandra/io/sstable/SSTableSimpleWriter.java
View File

@ -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<SSTableReader> finished = writer.finish(shouldOpenSSTables());
notifySSTableProduced(finished);
}
catch (Throwable t)
{

89
test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java
View File

@ -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<SSTableReader> 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<File> notifiedDataFileSet = produced.stream()
.map(sstable -> new File(sstable.descriptor.filenameFor(Component.DATA)))
.collect(Collectors.toSet());
Set<File> 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()