Merge branch 'cassandra-5.0' into trunk

* cassandra-5.0:
  Enhance CQLSSTableWriter to notify client on sstable production
This commit is contained in:
Yifan Cai 2024-11-22 14:59:35 -08:00
commit 8b126e97bb
6 changed files with 163 additions and 3 deletions

View File

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

View File

@ -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<SSTableId> id = new AtomicReference<>(SSTableIdFactory.instance.defaultBuilder().generator(Stream.empty()).get());
protected boolean makeRangeAware = false;
protected final Collection<Index.Group> indexGroups;
protected Consumer<Collection<SSTableReader>> 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<Collection<SSTableReader>> 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<SSTableReader> 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);

View File

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

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;
@ -219,7 +221,8 @@ class SSTableSimpleUnsortedWriter extends AbstractSSTableSimpleWriter
{
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,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<SSTableReader> finished = writer.finish(shouldOpenSSTables());
notifySSTableProduced(finished);
}
catch (Throwable t)
{

View File

@ -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<SSTableReader> 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<File> notifiedDataFileSet = produced.stream()
.map(sstable -> sstable.descriptor.fileFor(BigFormat.Components.DATA))
.collect(Collectors.toSet());
Set<File> listedDataFileSet = Arrays.stream(dataFiles)
.map(File::toCanonical)
.collect(Collectors.toSet());
assertEquals(notifiedDataFileSet, listedDataFileSet);
}
@Test
public void testMultipleWritersWithDistinctTables() throws IOException
{