Merge branch 'cassandra-4.0' into cassandra-4.1

* cassandra-4.0:
  Enhance CQLSSTableWriter to notify client on sstable production
This commit is contained in:
Yifan Cai 2024-11-22 14:50:51 -08:00
commit c44008e9e4
6 changed files with 163 additions and 4 deletions

View File

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

View File

@ -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<SSTableId> id = new AtomicReference<>(SSTableIdFactory.instance.defaultBuilder().generator(Stream.empty()).get());
protected boolean makeRangeAware = false;
protected Consumer<Collection<SSTableReader>> 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<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() throws IOException
{
SerializationHeader header = new SerializationHeader(true, metadata.get(), columns, EncodingStats.NO_STATS);

View File

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

View File

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

View File

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

View File

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