diff --git a/fluss-common/src/main/java/org/apache/fluss/row/encode/ValueEncoder.java b/fluss-common/src/main/java/org/apache/fluss/row/encode/ValueEncoder.java index 13069a16562..7e0a0478205 100644 --- a/fluss-common/src/main/java/org/apache/fluss/row/encode/ValueEncoder.java +++ b/fluss-common/src/main/java/org/apache/fluss/row/encode/ValueEncoder.java @@ -20,9 +20,12 @@ import org.apache.fluss.record.BinaryValue; import org.apache.fluss.row.BinaryRow; +import javax.annotation.Nullable; + import java.util.function.ToLongFunction; import static org.apache.fluss.utils.Preconditions.checkNotNull; +import static org.apache.fluss.utils.Preconditions.checkState; /** An encoder to encode {@link BinaryRow} with a schema id as value to be stored in kv store. */ public final class ValueEncoder { @@ -30,21 +33,28 @@ public final class ValueEncoder { private static final ValueEncoder PLAIN_ENCODER = new ValueEncoder(KvValueLayout.PLAIN, null); private final KvValueLayout kvValueLayout; - private final ToLongFunction valueTagProvider; - private ValueEncoder(KvValueLayout kvValueLayout, ToLongFunction valueTagProvider) { + /** + * Generates tags for {@link #encodeValue(BinaryValue)}. {@code null} for plain values or when + * callers supply each tag to {@link #encodeValue(BinaryValue, long)}. + */ + @Nullable private final ToLongFunction valueTagProvider; + + private ValueEncoder( + KvValueLayout kvValueLayout, @Nullable ToLongFunction valueTagProvider) { this.kvValueLayout = kvValueLayout; this.valueTagProvider = valueTagProvider; } - /** Returns an encoder for a layout without an internal value tag. */ + /** + * Returns an encoder for the given layout. Tagged values must supply their tag to {@link + * #encodeValue(BinaryValue, long)}. + */ public static ValueEncoder forLayout(KvValueLayout kvValueLayout) { checkNotNull(kvValueLayout, "kvValueLayout must not be null."); - if (kvValueLayout != KvValueLayout.PLAIN) { - throw new IllegalArgumentException( - "A value tag provider is required for this KV value layout."); - } - return PLAIN_ENCODER; + return kvValueLayout == KvValueLayout.PLAIN + ? PLAIN_ENCODER + : new ValueEncoder(kvValueLayout, null); } /** Returns an encoder for a layout with an internal value tag. */ @@ -66,12 +76,29 @@ public boolean hasValueTag() { /** Encodes a binary value using the layout bound to this encoder. */ public byte[] encodeValue(BinaryValue value) { + if (kvValueLayout.hasValueTag()) { + checkState( + valueTagProvider != null, + "An explicit value tag is required for this KV value encoder."); + return encodeValue(value, valueTagProvider.applyAsLong(value.row)); + } + return encodeValueBody(value); + } + + /** Encodes a binary value with the supplied opaque value tag. */ + public byte[] encodeValue(BinaryValue value, long valueTag) { + checkState( + kvValueLayout.hasValueTag(), + "An explicit value tag is not supported for this KV value layout."); + byte[] values = encodeValueBody(value); + kvValueLayout.writeValueTag(values, valueTag); + return values; + } + + private byte[] encodeValueBody(BinaryValue value) { int rowPayloadOffset = kvValueLayout.rowPayloadOffset(); byte[] values = new byte[rowPayloadOffset + value.row.getSizeInBytes()]; kvValueLayout.writeSchemaId(values, value.schemaId); - if (valueTagProvider != null) { - kvValueLayout.writeValueTag(values, valueTagProvider.applyAsLong(value.row)); - } value.row.copyTo(values, rowPayloadOffset); return values; } diff --git a/fluss-common/src/test/java/org/apache/fluss/row/encode/KvValueLayoutTest.java b/fluss-common/src/test/java/org/apache/fluss/row/encode/KvValueLayoutTest.java index a884fd05d16..98ba190f36f 100644 --- a/fluss-common/src/test/java/org/apache/fluss/row/encode/KvValueLayoutTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/row/encode/KvValueLayoutTest.java @@ -68,6 +68,22 @@ void testLongTagKeepsRpcValueBodyAsSuffix() { assertThat(KvValueLayout.TAGGED.toValueBodySlice(null)).isNull(); } + @Test + void testEncodeValueWithCallerProvidedTag() { + BinaryRow row = compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); + byte[] value = + ValueEncoder.forLayout(KvValueLayout.TAGGED) + .encodeValue(new BinaryValue(DEFAULT_SCHEMA_ID, row), 42L); + + assertThat(KvValueLayout.TAGGED.readValueTag(MemorySegment.wrap(value))).isEqualTo(42L); + assertThatThrownBy( + () -> + ValueEncoder.forLayout(KvValueLayout.TAGGED) + .encodeValue(new BinaryValue(DEFAULT_SCHEMA_ID, row))) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("explicit value tag"); + } + @Test void testTwoArgumentValueDecoderUsesPlainLayout() { BinaryRow row = compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java index a3afeec4bbf..643687806ce 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvRecoverHelper.java @@ -35,10 +35,10 @@ import org.apache.fluss.row.RowPartitionGetter; import org.apache.fluss.row.encode.KeyEncoder; import org.apache.fluss.row.encode.RowEncoder; -import org.apache.fluss.row.encode.ValueEncoder; import org.apache.fluss.row.indexed.IndexedRow; import org.apache.fluss.server.kv.autoinc.AutoIncIDRange; import org.apache.fluss.server.kv.historical.HistoricalKvKeyEncoder; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; import org.apache.fluss.server.log.FetchIsolation; import org.apache.fluss.server.log.LogTablet; import org.apache.fluss.server.zk.ZooKeeperClient; @@ -56,7 +56,6 @@ import java.util.List; import static org.apache.fluss.server.TabletManagerBase.getTableInfo; -import static org.apache.fluss.server.kv.KvStateAccessor.HISTORICAL_TOMBSTONE; import static org.apache.fluss.utils.Preconditions.checkArgument; import static org.apache.fluss.utils.Preconditions.checkNotNull; @@ -81,7 +80,7 @@ public class KvRecoverHelper { private KeyEncoder keyEncoder; private RowEncoder rowEncoder; - private final ValueEncoder valueEncoder; + private final KvStateValueEncoder stateValueEncoder; @Nullable private final RowTtlTimestampProvider rowTtlTimestampProvider; private final SchemaGetter schemaGetter; @@ -109,7 +108,7 @@ public KvRecoverHelper( this.kvFormat = kvFormat; this.logFormat = logFormat; this.schemaGetter = schemaGetter; - this.valueEncoder = kvTablet.getValueEncoder(); + this.stateValueEncoder = kvTablet.getStateValueEncoder(); this.rowTtlTimestampProvider = kvTablet.getRowTtlTimestampProvider(); this.remoteLogFetcher = remoteLogFetcher; this.historicalPartition = historicalPartition; @@ -149,7 +148,9 @@ public void recover() throws Exception { (resumeRecord) -> { if (resumeRecord.value == null) { if (historicalPartition) { - kvBatchWriter.put(resumeRecord.key, HISTORICAL_TOMBSTONE); + kvBatchWriter.put( + resumeRecord.key, + HistoricalKvTombstone.encode(resumeRecord.logOffset)); } else { kvBatchWriter.delete(resumeRecord.key); } @@ -296,9 +297,9 @@ private long applyLogRecordBatch( // the log row format may not compatible with kv row format, // e.g, arrow vs. compacted, thus needs a conversion here. BinaryRow row = toKvRow(logRow); - value = - valueEncoder.encodeValue( - new BinaryValue(currentSchemaId.shortValue(), row)); + BinaryValue binaryValue = + new BinaryValue(currentSchemaId.shortValue(), row); + value = stateValueEncoder.encodeValue(binaryValue, logRecord.logOffset()); } resumeRecordConsumer.accept( new KeyValueAndLogOffset( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateAccessor.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateAccessor.java index 1e0cd39962e..9891defc6cf 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateAccessor.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateAccessor.java @@ -18,7 +18,10 @@ package org.apache.fluss.server.kv; import org.apache.fluss.annotation.Internal; +import org.apache.fluss.record.BinaryValue; +import org.apache.fluss.row.encode.ValueDecoder; import org.apache.fluss.server.kv.historical.HistoricalKvKeyEncoder; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.Key; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.TruncateReason; @@ -27,6 +30,7 @@ import javax.annotation.Nullable; import java.io.IOException; +import java.util.function.Supplier; import static org.apache.fluss.utils.Preconditions.checkArgument; import static org.apache.fluss.utils.Preconditions.checkNotNull; @@ -35,9 +39,6 @@ @Internal public final class KvStateAccessor { - /** Encoded RocksDB value marking a deleted key in historical KV state. */ - static final byte[] HISTORICAL_TOMBSTONE = new byte[0]; - private final KvPreWriteBuffer preWriteBuffer; private final RocksDBKv rocksDBKv; private final boolean historicalPartition; @@ -71,25 +72,41 @@ public Key encodeKey(byte[] primaryKey, @Nullable String originalPartitionName) return Key.of(HistoricalKvKeyEncoder.encode(partitionName, primaryKey)); } - /** Looks up an encoded key from the local prewrite buffer and RocksDB state. */ - public KvStateLookupResult lookup(Key key) throws IOException { - KvPreWriteBuffer.Value bufferedValue = preWriteBuffer.get(key); - if (bufferedValue != null) { - byte[] value = bufferedValue.get(); - return value == null - ? KvStateLookupResult.deleted() - : KvStateLookupResult.present(value); - } + /** Looks up local state, preserving the distinction between missing keys and deletes. */ + public KvStateLookupResult lookupLocal(Key key) throws IOException { + KvStateLookupResult preWriteResult = lookupPreWriteBuffer(key); + return preWriteResult.status() == KvStateLookupResult.Status.NOT_FOUND + ? lookupRocksDB(key) + : preWriteResult; + } - byte[] value = rocksDBKv.get(key.get()); - if (value == null) { - return KvStateLookupResult.notFound(); + /** + * Looks up a decoded value from the prewrite buffer, then uses saved results or RocksDB on a + * miss. + * + *

A delete in the prewrite buffer stops fallback. When a fallback is supplied, its saved + * value is returned directly without reading RocksDB or decoding it again. + * + *

Callers must hold the KV lock. Write batches must keep the write lock throughout apply so + * their prewrite mutations remain visible to later records in the batch. + * + * @param key the encoded physical key to look up + * @param valueDecoder decodes values read from local state + * @param fallbackLookup an in-memory lookup of previously saved results, or null to read + * RocksDB on a prewrite buffer miss + */ + @Nullable + public BinaryValue lookup( + Key key, ValueDecoder valueDecoder, @Nullable Supplier fallbackLookup) + throws IOException { + KvStateLookupResult localResult = lookupPreWriteBuffer(key); + if (localResult.status() == KvStateLookupResult.Status.NOT_FOUND) { + if (fallbackLookup != null) { + return fallbackLookup.get(); + } + localResult = lookupRocksDB(key); } - // Historical KV tablets persist deletes as empty values so that a local miss does not - // expose a stale value from lake storage after the buffered delete has been flushed. - return value.length == 0 - ? KvStateLookupResult.deleted() - : KvStateLookupResult.present(value); + return localResult.isPresent() ? valueDecoder.decodeValue(localResult.value()) : null; } /** Adds an insert mutation to the prewrite buffer. */ @@ -118,4 +135,26 @@ public void delete(Key key, long logOffset) { public void truncateTo(long logOffset, TruncateReason reason) { preWriteBuffer.truncateTo(logOffset, reason); } + + /** Looks up the prewrite buffer, distinguishing missing keys from pending deletes. */ + private KvStateLookupResult lookupPreWriteBuffer(Key key) { + KvPreWriteBuffer.Value preWriteValue = preWriteBuffer.get(key); + if (preWriteValue == null) { + return KvStateLookupResult.notFound(); + } + byte[] value = preWriteValue.get(); + return value == null ? KvStateLookupResult.deleted() : KvStateLookupResult.present(value); + } + + private KvStateLookupResult lookupRocksDB(Key key) throws IOException { + byte[] value = rocksDBKv.get(key.get()); + if (value == null) { + return KvStateLookupResult.notFound(); + } + // Historical KV tablets persist deletes as offset-tagged tombstones so that a local miss + // does not expose a stale value from lake storage after the buffered delete is flushed. + return historicalPartition && HistoricalKvTombstone.isTombstone(value) + ? KvStateLookupResult.deleted() + : KvStateLookupResult.present(value); + } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateValueEncoder.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateValueEncoder.java new file mode 100644 index 00000000000..63de06154a4 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvStateValueEncoder.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.record.BinaryValue; + +/** Encodes KV state values using the encoding policy bound to their tablet. */ +@Internal +@FunctionalInterface +public interface KvStateValueEncoder { + + /** + * Encodes a state value at its producing WAL offset. + * + *

Historical state uses the offset as its value tag. Normal state ignores the offset and + * uses its configured plain or row TTL encoding. + */ + byte[] encodeValue(BinaryValue value, long logOffset); +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java index 76c603c8ec0..03814d6f3e9 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvTablet.java @@ -41,7 +41,9 @@ import org.apache.fluss.rpc.protocol.MergeMode; import org.apache.fluss.server.kv.autoinc.AutoIncIDRange; import org.apache.fluss.server.kv.autoinc.AutoIncrementManager; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; import org.apache.fluss.server.kv.historical.HistoricalValueLookup; +import org.apache.fluss.server.kv.historical.LocalValueLookupResult; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.PreparedFlush; import org.apache.fluss.server.kv.rocksdb.RocksDBKv; @@ -92,11 +94,12 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import static org.apache.fluss.server.kv.KvStateAccessor.HISTORICAL_TOMBSTONE; import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; +import static org.apache.fluss.utils.Preconditions.checkArgument; import static org.apache.fluss.utils.Preconditions.checkNotNull; import static org.apache.fluss.utils.Preconditions.checkState; import static org.apache.fluss.utils.concurrent.LockUtils.inReadLock; @@ -122,6 +125,10 @@ public final class KvTablet { private static final long ROW_COUNT_DISABLED = -1; + // Retain recent historical KV state within this WAL offset distance of lake progress. + // TODO: Consider time-based retention after lake coverage is confirmed. + private static final long HISTORICAL_KV_RETENTION_OFFSET_DISTANCE = 60_000L; + /** * Max records per native write of the asynchronous flush; mirrors the batching capacity of * {@code RocksDBWriteBatchWrapper} (hundreds of keys per write batch is RocksDB best practice). @@ -148,7 +155,8 @@ public final class KvTablet { // A lock that guards all modifications to the kv. private final ReadWriteLock kvLock = new ReentrantReadWriteLock(); private final KvValueLayout kvValueLayout; - private final ValueEncoder valueEncoder; + private final KvStateValueEncoder stateValueEncoder; + private final AtomicLong historicalCleanupOffset; @Nullable private final RowTtlTimestampProvider rowTtlTimestampProvider; private final boolean rowTtlEnabled; private final AutoIncrementManager autoIncrementManager; @@ -199,6 +207,7 @@ private KvTablet( KvValueLayout kvValueLayout, ValueEncoder valueEncoder, ValueDecoder valueDecoder, + AtomicLong historicalCleanupOffset, @Nullable RocksDBStatistics rocksDBStatistics, KvFlushScheduler kvFlushScheduler, boolean closeFlushScheduler, @@ -222,7 +231,11 @@ private KvTablet( this.kvStateAccessor = new KvStateAccessor(kvPreWriteBuffer, rocksDBKv, historicalPartition); this.kvValueLayout = kvValueLayout; - this.valueEncoder = valueEncoder; + this.stateValueEncoder = + historicalPartition + ? valueEncoder::encodeValue + : (value, logOffset) -> valueEncoder.encodeValue(value); + this.historicalCleanupOffset = historicalCleanupOffset; this.rowTtlTimestampProvider = rowTtlTimestampProvider; this.rowTtlEnabled = rowTtlEnabled; this.kvWriteProcessor = @@ -237,7 +250,7 @@ private KvTablet( schemaGetter, changelogImage, autoIncrementManager, - valueEncoder, + stateValueEncoder, valueDecoder, rowTtlTimestampProvider, clock); @@ -424,11 +437,16 @@ private static KvTablet create( TableConfig tableConfig) throws IOException { checkNotNull(tableConfig, "tableConfig must not be null."); + boolean historicalPartition = + HISTORICAL_PARTITION_VALUE.equals(tablePath.getPartitionName()); Optional rowTtl = tableConfig.getKvTTL(); - KvValueLayout kvValueLayout = KvValueLayout.fromTableConfig(tableConfig); + KvValueLayout kvValueLayout = + historicalPartition + ? KvValueLayout.TAGGED + : KvValueLayout.fromTableConfig(tableConfig); @Nullable RowTtlTimestampProvider rowTtlTimestampProvider = - kvValueLayout.hasValueTag() + !historicalPartition && kvValueLayout.hasValueTag() ? RowTtlTimestampProvider.create( tableConfig, schemaGetter, ZoneId.systemDefault()) : null; @@ -437,13 +455,19 @@ private static KvTablet create( ? ValueEncoder.forLayout(kvValueLayout) : ValueEncoder.forLayout(kvValueLayout, rowTtlTimestampProvider); ValueDecoder valueDecoder = new ValueDecoder(schemaGetter, kvFormat, kvValueLayout); + AtomicLong historicalCleanupOffset = new AtomicLong(0L); @Nullable AbstractCompactionFilterFactory> compactionFilterFactory = - rowTtl.isPresent() + historicalPartition ? RowTtlCompactionFilterFactory.create( - kvValueLayout, rowTtl.get(), clock) - : null; + kvValueLayout, + HISTORICAL_KV_RETENTION_OFFSET_DISTANCE, + () -> historicalCleanupOffset.get() - 1L) + : rowTtl.isPresent() + ? RowTtlCompactionFilterFactory.create( + kvValueLayout, rowTtl.get(), clock) + : null; RocksDBKv kv = buildRocksDBKv( serverConf, @@ -483,6 +507,7 @@ private static KvTablet create( kvValueLayout, valueEncoder, valueDecoder, + historicalCleanupOffset, rocksDBStatistics, kvFlushScheduler, closeFlushScheduler, @@ -579,8 +604,8 @@ private static RocksDBKv buildRocksDBKv( } } - ValueEncoder getValueEncoder() { - return valueEncoder; + KvStateValueEncoder getStateValueEncoder() { + return stateValueEncoder; } @Nullable @@ -728,16 +753,15 @@ public LogAppendInfo putAsLeader( * Puts records for one original partition into this historical KV tablet. * *

The original partition name namespaces the physical primary keys because one historical - * bucket can contain records from multiple original partitions. The supplied fallback may only - * read lake results already resolved for this request; it must not perform lake I/O while the - * tablet lock is held. + * bucket can contain records from multiple original partitions. The supplied lookup must + * contain every previous value required by this batch and must not perform I/O. */ public LogAppendInfo putHistoricalAsLeader( KvRecordBatch kvRecords, @Nullable int[] targetColumns, MergeMode mergeMode, String originalPartitionName, - HistoricalValueLookup memoizedLakeLookup) + HistoricalValueLookup historicalValueLookup) throws Exception { checkState(historicalPartition, "%s is not a historical KV tablet", tableBucket); return putAsLeader( @@ -745,16 +769,16 @@ public LogAppendInfo putHistoricalAsLeader( targetColumns, mergeMode, checkNotNull(originalPartitionName, "originalPartitionName must not be null"), - checkNotNull(memoizedLakeLookup, "memoizedLakeLookup must not be null")); + checkNotNull(historicalValueLookup, "Historical value lookup must not be null")); } /** - * Finds keys whose historical write requires an old value that is absent from local state. + * Probes the local previous values required by a historical write. * *

This method only reads KV entries and uses the tablet read lock. Lake I/O must be * performed by the caller after this method releases the tablet lock. */ - public List findKeysRequiringLakeLookup( + public LocalValueLookupResult probeLocalPreviousValues( KvRecordBatch kvRecords, @Nullable int[] targetColumns, MergeMode mergeMode, @@ -765,7 +789,7 @@ public List findKeysRequiringLakeLookup( kvLock, () -> { rocksDBKv.checkIfRocksDBClosed(); - return kvWriteProcessor.findKeysRequiringLakeLookup( + return kvWriteProcessor.probeLocalPreviousValues( kvRecords, targetColumns, mergeMode, @@ -781,7 +805,7 @@ private LogAppendInfo putAsLeader( @Nullable int[] targetColumns, MergeMode mergeMode, @Nullable String originalPartitionName, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { return inWriteLock( kvLock, @@ -810,7 +834,7 @@ private LogAppendInfo putAsLeader( mergeMode, kvStateAccessor, originalPartitionName, - memoizedLakeLookup); + historicalValueLookup); }); } @@ -834,6 +858,21 @@ public long getFlushedLogOffset() { return flushedLogOffset; } + /** Advances the exclusive historical cleanup offset without allowing it to move backwards. */ + public boolean advanceHistoricalCleanupOffset(long cleanupOffset) { + checkState(historicalPartition, "%s is not a historical KV tablet", tableBucket); + checkArgument(cleanupOffset >= 0L, "Historical cleanup offset must be non-negative."); + long previousCleanupOffset = + historicalCleanupOffset.getAndAccumulate(cleanupOffset, Math::max); + return cleanupOffset > previousCleanupOffset; + } + + /** Returns the current exclusive cleanup offset for a historical overlay. */ + public long getHistoricalCleanupOffset() { + checkState(historicalPartition, "%s is not a historical KV tablet", tableBucket); + return historicalCleanupOffset.get(); + } + @VisibleForTesting FlushState getFlushState() { return inReadLock(kvLock, () -> flushState); @@ -979,7 +1018,9 @@ private void writePreparedFlush(PreparedFlush preparedFlush) throws Exception { if (historicalPartition) { // A physical delete would turn a local miss into a lake lookup and // could expose the stale value that this mutation deleted. - kvBatchWriter.put(entry.getKey().get(), HISTORICAL_TOMBSTONE); + kvBatchWriter.put( + entry.getKey().get(), + HistoricalKvTombstone.encode(entry.getLogSequenceNumber())); } else { kvBatchWriter.delete(entry.getKey().get()); } @@ -1171,7 +1212,7 @@ public List multiGetFromBufferOrKv(List keys) throws IOE List values = new ArrayList<>(keys.size()); for (byte[] key : keys) { KvPreWriteBuffer.Key lookupKey = kvStateAccessor.encodeKey(key, null); - byte[] rawValue = kvStateAccessor.lookup(lookupKey).value(); + byte[] rawValue = kvStateAccessor.lookupLocal(lookupKey).value(); values.add(kvValueLayout.toValueBodySlice(rawValue)); } return values; @@ -1192,7 +1233,7 @@ public KvStateLookupResult lookupHistoricalLocal(String originalPartitionName, b if (value == null) { return KvStateLookupResult.notFound(); } - return value.length == 0 + return HistoricalKvTombstone.isTombstone(value) ? KvStateLookupResult.deleted() : KvStateLookupResult.present(value); }); diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvWriteProcessor.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvWriteProcessor.java index bc0fb53fff0..f67bfaa627b 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/KvWriteProcessor.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/KvWriteProcessor.java @@ -40,11 +40,11 @@ import org.apache.fluss.row.arrow.ArrowWriterProvider; import org.apache.fluss.row.encode.KvValueLayout; import org.apache.fluss.row.encode.ValueDecoder; -import org.apache.fluss.row.encode.ValueEncoder; import org.apache.fluss.rpc.protocol.MergeMode; import org.apache.fluss.server.kv.autoinc.AutoIncrementManager; import org.apache.fluss.server.kv.autoinc.AutoIncrementUpdater; import org.apache.fluss.server.kv.historical.HistoricalValueLookup; +import org.apache.fluss.server.kv.historical.LocalValueLookupResult; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer; import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.TruncateReason; import org.apache.fluss.server.kv.rowmerger.DefaultRowMerger; @@ -66,9 +66,7 @@ import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; -import java.util.ArrayList; import java.util.HashSet; -import java.util.List; import java.util.Set; /** @@ -81,10 +79,10 @@ * *

The supplied {@link KvStateAccessor} defines how keys and state are accessed. Normal writes * use the original primary key and local state, while historical writes use partition-scoped keys. - * On a historical local miss, the processor can consult a lake result already memoized for the - * current request. Resolving that result from lake storage remains the caller's responsibility and - * must happen outside the tablet lock. The merge and WAL generation path is shared by both write - * kinds. + * Historical writes supply previous values saved for the current request as the accessor's + * fallback. The saved values include local probe results, which remain available after compaction, + * and lake results resolved by the caller outside the tablet lock. The merge and WAL generation + * path is shared by both write kinds. */ @Internal @NotThreadSafe @@ -108,7 +106,7 @@ public final class KvWriteProcessor { // the changelog image mode for this tablet private final ChangelogImage changelogImage; private final AutoIncrementManager autoIncrementManager; - private final ValueEncoder valueEncoder; + private final KvStateValueEncoder stateValueEncoder; private final ValueDecoder valueDecoder; private final ValueDecoder lakeValueDecoder; @Nullable private final RowTtlTimestampProvider rowTtlTimestampProvider; @@ -126,7 +124,7 @@ public KvWriteProcessor( SchemaGetter schemaGetter, ChangelogImage changelogImage, AutoIncrementManager autoIncrementManager, - ValueEncoder valueEncoder, + KvStateValueEncoder stateValueEncoder, ValueDecoder valueDecoder, @Nullable RowTtlTimestampProvider rowTtlTimestampProvider, Clock clock) { @@ -144,7 +142,7 @@ public KvWriteProcessor( this.schemaGetter = schemaGetter; this.changelogImage = changelogImage; this.autoIncrementManager = autoIncrementManager; - this.valueEncoder = valueEncoder; + this.stateValueEncoder = stateValueEncoder; this.valueDecoder = valueDecoder; this.lakeValueDecoder = new ValueDecoder(schemaGetter, kvFormat, KvValueLayout.PLAIN); this.rowTtlTimestampProvider = rowTtlTimestampProvider; @@ -158,7 +156,7 @@ public LogAppendInfo putAsLeader( MergeMode mergeMode, KvStateAccessor stateAccessor, @Nullable String originalPartitionName, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { WriteContext writeContext = createWriteContext(kvRecords, targetColumns, mergeMode); RowType latestRowType = writeContext.latestSchema.getRowType(); @@ -185,7 +183,7 @@ public LogAppendInfo putAsLeader( logEndOffsetOfPrevBatch, stateAccessor, originalPartitionName, - memoizedLakeLookup); + historicalValueLookup); // There will be a situation that these batches of kvRecordBatch have not // generated any CDC logs, for example, when client attempts to delete @@ -221,14 +219,13 @@ public LogAppendInfo putAsLeader( } /** - * Finds the original primary keys whose previous values must be loaded from lake storage before - * applying a historical write batch. + * Probes the local previous values required by a historical write batch. * - *

A key is returned only when its previous value is required and its partition-scoped key is - * absent from local state. Records that can establish their result without a previous value are - * skipped, and each key is returned at most once. + *

Every key whose previous value is required is probed at most once. Local values and + * deletes are decoded into the returned collection, while true local misses are exposed for + * lake lookup. Records that establish their result without a previous value are skipped. */ - List findKeysRequiringLakeLookup( + LocalValueLookupResult probeLocalPreviousValues( KvRecordBatch kvRecords, @Nullable int[] targetColumns, MergeMode mergeMode, @@ -237,9 +234,9 @@ List findKeysRequiringLakeLookup( throws Exception { WriteContext writeContext = createWriteContext(kvRecords, targetColumns, mergeMode); - List keysRequiringLakeLookup = new ArrayList<>(); - // Track keys whose lake-lookup requirement has already been evaluated. - Set keysEvaluatedForLakeLookup = new HashSet<>(); + LocalValueLookupResult localLookupResult = + new LocalValueLookupResult(valueDecoder, lakeValueDecoder); + Set probedKeys = new HashSet<>(); KvRecordBatch.ReadContext readContext = KvRecordReadContext.createReadContext(kvFormat, schemaGetter); for (KvRecord kvRecord : kvRecords.records(readContext)) { @@ -252,23 +249,20 @@ List findKeysRequiringLakeLookup( } else if (canSkipOldValueLookup( writeContext.rowMerger, writeContext.autoIncrementUpdater)) { // A full-row WAL upsert establishes the state without reading its previous value. - keysEvaluatedForLakeLookup.add(wrappedKey); + probedKeys.add(wrappedKey); continue; } // Probe local state only for the first record that needs a previous value. A true local // miss schedules one lake lookup shared by all records for this key in the batch. - if (keysEvaluatedForLakeLookup.add(wrappedKey) - && stateAccessor - .lookup( - stateAccessor.encodeKey( - primaryKey, originalPartitionName)) - .status() - == KvStateLookupResult.Status.NOT_FOUND) { - keysRequiringLakeLookup.add(primaryKey); + if (probedKeys.add(wrappedKey)) { + localLookupResult.add( + primaryKey, + stateAccessor.lookupLocal( + stateAccessor.encodeKey(primaryKey, originalPartitionName))); } } - return keysRequiringLakeLookup; + return localLookupResult; } private WriteContext createWriteContext( @@ -313,7 +307,7 @@ private void processKvRecords( long startLogOffset, KvStateAccessor stateAccessor, @Nullable String originalPartitionName, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { long logOffset = startLogOffset; @@ -337,7 +331,7 @@ private void processKvRecords( logOffset, stateAccessor, keyBytes, - memoizedLakeLookup); + historicalValueLookup); } else { logOffset = processUpsert( @@ -350,7 +344,7 @@ private void processKvRecords( logOffset, stateAccessor, keyBytes, - memoizedLakeLookup); + historicalValueLookup); } } } @@ -363,13 +357,14 @@ private long processDeletion( long logOffset, KvStateAccessor stateAccessor, byte[] primaryKey, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { if (shouldIgnoreDeletion(currentMerger)) { return logOffset; } - BinaryValue oldValue = getPreviousValue(key, primaryKey, stateAccessor, memoizedLakeLookup); + BinaryValue oldValue = + getPreviousValue(key, primaryKey, stateAccessor, historicalValueLookup); if (oldValue == null) { LOG.debug( "The specific key can't be found in kv tablet although the kv record is for deletion, " @@ -399,14 +394,15 @@ private long processUpsert( long logOffset, KvStateAccessor stateAccessor, byte[] primaryKey, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { if (canSkipOldValueLookup(currentMerger, autoIncrementUpdater)) { return applyUpdate( key, null, currentValue, walBuilder, latestSchemaRow, logOffset, stateAccessor); } - BinaryValue oldValue = getPreviousValue(key, primaryKey, stateAccessor, memoizedLakeLookup); + BinaryValue oldValue = + getPreviousValue(key, primaryKey, stateAccessor, historicalValueLookup); if (oldValue == null) { BinaryValue valueToInsert = currentMerger.merge(null, currentValue); return applyInsert( @@ -454,7 +450,7 @@ private long applyInsert( throws Exception { BinaryValue newValue = autoIncrementUpdater.updateAutoIncrementColumns(currentValue); walBuilder.append(ChangeType.INSERT, latestSchemaRow.replaceRow(newValue.row)); - stateAccessor.insert(key, valueEncoder.encodeValue(newValue), logOffset); + stateAccessor.insert(key, stateValueEncoder.encodeValue(newValue, logOffset), logOffset); return logOffset + 1; } @@ -469,42 +465,44 @@ private long applyUpdate( throws Exception { if (changelogImage == ChangelogImage.WAL) { walBuilder.append(ChangeType.UPDATE_AFTER, latestSchemaRow.replaceRow(newValue.row)); - stateAccessor.update(key, valueEncoder.encodeValue(newValue), logOffset); + stateAccessor.update( + key, stateValueEncoder.encodeValue(newValue, logOffset), logOffset); return logOffset + 1; } else { walBuilder.append(ChangeType.UPDATE_BEFORE, latestSchemaRow.replaceRow(oldValue.row)); walBuilder.append(ChangeType.UPDATE_AFTER, latestSchemaRow.replaceRow(newValue.row)); - stateAccessor.update(key, valueEncoder.encodeValue(newValue), logOffset + 1); + long updateAfterOffset = logOffset + 1; + stateAccessor.update( + key, + stateValueEncoder.encodeValue(newValue, updateAfterOffset), + updateAfterOffset); return logOffset + 2; } } /** - * Returns the previous value from local state, falling back to the memoized lake result only on - * a genuine local miss. + * Returns the previous value through the state accessor, using saved results as the fallback + * for historical writes. * * @param localStateKey the key used by the local prewrite buffer and RocksDB; it wraps {@code * primaryKey} for a normal partition and adds the original partition namespace for a * historical partition * @param primaryKey the encoded bytes of the logical primary key, without the historical - * partition namespace; used by the lake lookup + * partition namespace * @return the previous value, or null if the key is absent or locally marked as deleted */ private BinaryValue getPreviousValue( KvPreWriteBuffer.Key localStateKey, byte[] primaryKey, KvStateAccessor stateAccessor, - @Nullable HistoricalValueLookup memoizedLakeLookup) + @Nullable HistoricalValueLookup historicalValueLookup) throws Exception { - KvStateLookupResult localResult = stateAccessor.lookup(localStateKey); - if (localResult.status() != KvStateLookupResult.Status.NOT_FOUND - || memoizedLakeLookup == null) { - return localResult.value() == null - ? null - : valueDecoder.decodeValue(localResult.value()); - } - byte[] lakeValue = memoizedLakeLookup.lookup(primaryKey); - return lakeValue == null ? null : lakeValueDecoder.decodeValue(lakeValue); + return stateAccessor.lookup( + localStateKey, + valueDecoder, + historicalValueLookup == null + ? null + : () -> historicalValueLookup.lookup(primaryKey)); } private boolean canSkipOldValueLookup( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/RowTtlCompactionFilterFactory.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/RowTtlCompactionFilterFactory.java index 8cb01d0c1f2..e77e5b9b468 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/RowTtlCompactionFilterFactory.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/RowTtlCompactionFilterFactory.java @@ -17,7 +17,6 @@ package org.apache.fluss.server.kv; -import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.row.encode.KvValueLayout; import org.apache.fluss.server.utils.RowTtlUtils; import org.apache.fluss.utils.clock.Clock; @@ -26,6 +25,7 @@ import org.rocksdb.RocksDB; import java.time.Duration; +import java.util.function.LongSupplier; import static org.apache.fluss.utils.Preconditions.checkArgument; import static org.apache.fluss.utils.Preconditions.checkNotNull; @@ -40,29 +40,47 @@ private RowTtlCompactionFilterFactory() {} /** Creates a configured native compaction filter factory for row TTL cleanup. */ public static FlinkCompactionFilter.FlinkCompactionFilterFactory create( KvValueLayout kvValueLayout, Duration ttl, Clock clock) { - return create(kvValueLayout, ttl, QUERY_TIME_AFTER_NUM_ENTRIES, clock); + long ttlMillis = RowTtlUtils.validateAndConvertTtlDurationToMillis(ttl); + checkNotNull(clock, "clock must not be null."); + return create(kvValueLayout, ttlMillis, clock::milliseconds); } - @VisibleForTesting + /** Removes values using the default interval for refreshing the supplied current value. */ static FlinkCompactionFilter.FlinkCompactionFilterFactory create( - KvValueLayout kvValueLayout, Duration ttl, long queryTimeAfterNumEntries, Clock clock) { - long ttlMillis = RowTtlUtils.validateAndConvertTtlDurationToMillis(ttl); + KvValueLayout kvValueLayout, + long expirationDistance, + LongSupplier currentValueSupplier) { + return create( + kvValueLayout, + expirationDistance, + QUERY_TIME_AFTER_NUM_ENTRIES, + currentValueSupplier); + } + + /** Removes a value when {@code valueTag + expirationDistance <= currentValue}. */ + static FlinkCompactionFilter.FlinkCompactionFilterFactory create( + KvValueLayout kvValueLayout, + long expirationDistance, + long queryCurrentValueAfterNumEntries, + LongSupplier currentValueSupplier) { checkNotNull(kvValueLayout, "kvValueLayout must not be null."); - checkNotNull(clock, "clock must not be null."); - checkArgument(kvValueLayout.hasValueTag(), "Row TTL requires a tagged KV value layout."); + checkNotNull(currentValueSupplier, "currentValueSupplier must not be null."); + checkArgument(kvValueLayout.hasValueTag(), "Compaction filter requires a tagged layout."); + checkArgument(expirationDistance >= 0L, "Expiration distance must be non-negative."); checkArgument( - queryTimeAfterNumEntries > 0, - "queryTimeAfterNumEntries must be greater than zero."); + queryCurrentValueAfterNumEntries > 0L, + "queryCurrentValueAfterNumEntries must be greater than zero."); RocksDB.loadLibrary(); FlinkCompactionFilter.FlinkCompactionFilterFactory factory = - new FlinkCompactionFilter.FlinkCompactionFilterFactory(clock::milliseconds); + new FlinkCompactionFilter.FlinkCompactionFilterFactory( + currentValueSupplier::getAsLong); factory.configure( FlinkCompactionFilter.Config.createNotList( FlinkCompactionFilter.StateType.Value, kvValueLayout.valueTagOffset(), - ttlMillis, - queryTimeAfterNumEntries)); + expirationDistance, + queryCurrentValueAfterNumEntries)); return factory; } } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalKvTombstone.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalKvTombstone.java new file mode 100644 index 00000000000..a0baab3e495 --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalKvTombstone.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv.historical; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.row.encode.KvValueLayout; + +import static org.apache.fluss.utils.Preconditions.checkArgument; + +/** Utilities for tombstones stored in local historical KV state. */ +@Internal +public final class HistoricalKvTombstone { + + private HistoricalKvTombstone() {} + + /** Encodes a tombstone tagged with the WAL offset that produced the delete. */ + public static byte[] encode(long logOffset) { + checkArgument(logOffset >= 0L, "Historical KV log offset must be non-negative."); + byte[] tombstone = new byte[KvValueLayout.TAGGED.valueTagLength()]; + KvValueLayout.TAGGED.writeValueTag(tombstone, logOffset); + return tombstone; + } + + /** Returns whether the raw historical value is an offset-tagged tombstone. */ + public static boolean isTombstone(byte[] rawValue) { + return rawValue.length == KvValueLayout.TAGGED.valueTagLength(); + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalValueLookup.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalValueLookup.java index bdc723df96e..e632fd2a213 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalValueLookup.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/HistoricalValueLookup.java @@ -2,7 +2,7 @@ * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. - * The ASF licenses this file to you under the Apache License, Version 2.0 + * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * @@ -18,20 +18,21 @@ package org.apache.fluss.server.kv.historical; import org.apache.fluss.annotation.Internal; +import org.apache.fluss.record.BinaryValue; import javax.annotation.Nullable; -/** Resolves a lake value already memoized for the current historical write request. */ +/** Looks up a previous value already memoized for the current historical write request. */ @Internal @FunctionalInterface public interface HistoricalValueLookup { /** - * Returns the encoded value for the primary key, or null when it does not exist. + * Returns the decoded previous value, or null when the key was absent or deleted. * - *

This method is invoked while the KV write lock is held and must not perform lake or file + *

This method is invoked while the KV write lock is held and must not perform local or lake * I/O. */ @Nullable - byte[] lookup(byte[] primaryKey); + BinaryValue lookup(byte[] primaryKey); } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/LocalValueLookupResult.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/LocalValueLookupResult.java new file mode 100644 index 00000000000..a8340f2816e --- /dev/null +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/historical/LocalValueLookupResult.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv.historical; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.record.BinaryValue; +import org.apache.fluss.row.encode.ValueDecoder; +import org.apache.fluss.server.kv.KvStateLookupResult; +import org.apache.fluss.utils.ByteArrayWrapper; + +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.apache.fluss.utils.Preconditions.checkArgument; +import static org.apache.fluss.utils.Preconditions.checkNotNull; +import static org.apache.fluss.utils.Preconditions.checkState; + +/** Local lookup results for previous values required by a historical write request. */ +@Internal +@NotThreadSafe +public final class LocalValueLookupResult { + + /** Decodes tagged values returned by local historical KV state. */ + private final ValueDecoder localValueDecoder; + + /** Decodes plain values returned by lake storage. */ + private final ValueDecoder lakeValueDecoder; + + /** Previous values found locally, including tombstones represented by an empty optional. */ + private final Map> localValuesByKey = new HashMap<>(); + + /** True local misses in lake request order. */ + private final Set keysMissingLocally = new LinkedHashSet<>(); + + /** Creates a local lookup result with decoders for the local and lake value layouts. */ + public LocalValueLookupResult(ValueDecoder localValueDecoder, ValueDecoder lakeValueDecoder) { + this.localValueDecoder = localValueDecoder; + this.lakeValueDecoder = lakeValueDecoder; + } + + /** Records the local result for one key whose previous value is required. */ + public void add(byte[] primaryKey, KvStateLookupResult localResult) { + ByteArrayWrapper wrappedKey = new ByteArrayWrapper(primaryKey); + checkState( + !localValuesByKey.containsKey(wrappedKey) + && !keysMissingLocally.contains(wrappedKey), + "Historical write key has already been probed"); + + if (localResult.status() == KvStateLookupResult.Status.NOT_FOUND) { + keysMissingLocally.add(wrappedKey); + } else { + localValuesByKey.put(wrappedKey, decode(localResult.value(), localValueDecoder)); + } + } + + /** Returns a snapshot of true local misses in lake request order. */ + public List keysMissingLocally() { + List primaryKeys = new ArrayList<>(keysMissingLocally.size()); + for (ByteArrayWrapper keyMissingLocally : keysMissingLocally) { + primaryKeys.add(keyMissingLocally.getData()); + } + return Collections.unmodifiableList(primaryKeys); + } + + /** + * Combines lake results with the local lookup results and creates an in-memory lookup for + * apply. + * + *

Lake values must have the same order as {@link #keysMissingLocally()}. + */ + public HistoricalValueLookup createValueLookup(List lakeValues) { + checkNotNull(lakeValues, "Historical lake values must not be null"); + checkArgument( + lakeValues.size() == keysMissingLocally.size(), + "Expected %s historical lake values, but received %s", + keysMissingLocally.size(), + lakeValues.size()); + Map> previousValuesByKey = + new HashMap<>(localValuesByKey); + Iterator missingKeyIterator = keysMissingLocally.iterator(); + for (byte[] lakeValue : lakeValues) { + previousValuesByKey.put(missingKeyIterator.next(), decode(lakeValue, lakeValueDecoder)); + } + return primaryKey -> + checkNotNull( + previousValuesByKey.get(new ByteArrayWrapper(primaryKey)), + "No previous value for a historical write key") + .orElse(null); + } + + private static Optional decode( + @Nullable byte[] encodedValue, ValueDecoder valueDecoder) { + return encodedValue == null + ? Optional.empty() + : Optional.of(valueDecoder.decodeValue(encodedValue)); + } +} diff --git a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java index aae74e19b7f..e816c7f7cec 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/kv/prewrite/KvPreWriteBuffer.java @@ -412,7 +412,8 @@ public Value getValue() { return value; } - long getLogSequenceNumber() { + /** Returns the WAL offset that produced this mutation. */ + public long getLogSequenceNumber() { return logSequenceNumber; } diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java index 079e1f96502..9b4d4b8b4f5 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/Replica.java @@ -66,6 +66,7 @@ import org.apache.fluss.server.kv.RemoteLogFetcher; import org.apache.fluss.server.kv.autoinc.AutoIncIDRange; import org.apache.fluss.server.kv.historical.HistoricalValueLookup; +import org.apache.fluss.server.kv.historical.LocalValueLookupResult; import org.apache.fluss.server.kv.rocksdb.RocksDBKvBuilder; import org.apache.fluss.server.kv.scan.OpenScanResult; import org.apache.fluss.server.kv.scan.ScannerContext; @@ -114,6 +115,7 @@ import org.apache.fluss.utils.FlussPaths; import org.apache.fluss.utils.IOUtils; import org.apache.fluss.utils.clock.Clock; +import org.apache.fluss.utils.function.FunctionWithException; import org.apache.fluss.utils.types.Tuple2; import org.slf4j.Logger; @@ -908,6 +910,10 @@ private Optional initKvTablet() { } logTablet.updateMinRetainOffset(restoreStartOffset); + if (isHistoricalPartition()) { + checkNotNull(kvTablet, "kv tablet should not be null.") + .advanceHistoricalCleanupOffset(restoreStartOffset); + } recoverKvTablet(restoreStartOffset, rowCount, autoIncIDRange); } catch (Exception e) { throw new KvStorageException( @@ -1270,40 +1276,38 @@ public LogAppendInfo putRecordsToLeader( } /** - * Finds historical write keys that require lake fallback without mutating local KV state. - * - *

The caller must keep historical writes for this table bucket ordered until the subsequent - * {@link #putHistoricalRecordsToLeader} call completes. + * Looks up previous values without holding replica or KV locks during lake I/O, then writes the + * records to the local historical KV state. */ - public List findKeysRequiringLakeLookup( - KvRecordBatch kvRecords, - @Nullable int[] targetColumns, - MergeMode mergeMode, - String originalPartitionName, - int expectedLeaderEpoch, - int requiredAcks) - throws Exception { - return inReadLock( - leaderIsrUpdateLock, - () -> { - validateHistoricalWrite(expectedLeaderEpoch, requiredAcks); - KvTablet kv = this.kvTablet; - checkNotNull(kv, "KvTablet for the historical replica shouldn't be null."); - return kv.findKeysRequiringLakeLookup( - kvRecords, targetColumns, mergeMode, originalPartitionName); - }); - } - - /** Writes records to the local historical KV state of the leader replica. */ public LogAppendInfo putHistoricalRecordsToLeader( KvRecordBatch kvRecords, @Nullable int[] targetColumns, MergeMode mergeMode, String originalPartitionName, - HistoricalValueLookup memoizedLakeLookup, + FunctionWithException, List, Exception> lakeLookup, int expectedLeaderEpoch, int requiredAcks) throws Exception { + checkNotNull(lakeLookup, "Historical lake lookup must not be null"); + LocalValueLookupResult localLookupResult = + inReadLock( + leaderIsrUpdateLock, + () -> { + validateHistoricalWrite(expectedLeaderEpoch, requiredAcks); + KvTablet kv = this.kvTablet; + checkNotNull( + kv, "KvTablet for the historical replica shouldn't be null."); + return kv.probeLocalPreviousValues( + kvRecords, targetColumns, mergeMode, originalPartitionName); + }); + + List missingKeys = localLookupResult.keysMissingLocally(); + // Both the replica read lock and the KV read lock have been released before lake I/O. + List lakeValues = + missingKeys.isEmpty() ? Collections.emptyList() : lakeLookup.apply(missingKeys); + HistoricalValueLookup historicalValueLookup = + localLookupResult.createValueLookup(lakeValues); + return inReadLock( leaderIsrUpdateLock, () -> { @@ -1316,12 +1320,49 @@ public LogAppendInfo putHistoricalRecordsToLeader( targetColumns, mergeMode, originalPartitionName, - memoizedLakeLookup); + historicalValueLookup); maybeIncrementLeaderHW(logTablet, clock.milliseconds()); return appendInfo; }); } + /** + * Updates the historical cleanup offset if it is valid. + * + *

{@code beforeUpdate} runs after validation and before the new cleanup offset becomes + * visible to the RocksDB compaction filter. + */ + public void tryUpdateHistoricalCleanupOffset(long newCleanupOffset, Runnable beforeUpdate) { + checkNotNull(beforeUpdate, "beforeUpdate must not be null."); + inWriteLock( + leaderIsrUpdateLock, + () -> { + if (!isLeader() || !historicalPartition) { + return; + } + KvTablet currentKvTablet = kvTablet; + if (currentKvTablet == null) { + return; + } + long currentCleanupOffset = currentKvTablet.getHistoricalCleanupOffset(); + long localLogEndOffset = logTablet.localLogEndOffset(); + if (newCleanupOffset < currentCleanupOffset + || newCleanupOffset > localLogEndOffset) { + LOG.warn( + "Ignore invalid historical cleanup offset {} for {} with " + + "current cleanup offset {} and local log end " + + "offset {}.", + newCleanupOffset, + tableBucket, + currentCleanupOffset, + localLogEndOffset); + return; + } + beforeUpdate.run(); + currentKvTablet.advanceHistoricalCleanupOffset(newCleanupOffset); + }); + } + private void validateHistoricalWrite(int expectedLeaderEpoch, int requiredAcks) { if (!isLeader()) { throw new NotLeaderOrFollowerException( diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java index ef8ed1696fd..b6b8a1d9a34 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/ReplicaManager.java @@ -1411,7 +1411,8 @@ public void notifyLakeTableOffset( lakeBucketOffsets.entrySet()) { TableBucket tb = lakeBucketOffsetEntry.getKey(); LakeBucketOffset lakeBucketOffset = lakeBucketOffsetEntry.getValue(); - LogTablet logTablet = getReplicaOrException(tb).getLogTablet(); + Replica replica = getReplicaOrException(tb); + LogTablet logTablet = replica.getLogTablet(); logTablet.updateLakeTableSnapshotId(lakeBucketOffset.getSnapshotId()); lakeBucketOffset @@ -1420,7 +1421,16 @@ public void notifyLakeTableOffset( lakeBucketOffset .getLogEndOffset() - .ifPresent(logTablet::updateLakeLogEndOffset); + .ifPresent( + lakeLogEndOffset -> { + logTablet.updateLakeLogEndOffset(lakeLogEndOffset); + if (replica.isHistoricalPartition()) { + historicalPartitionManager.onLakeProgress( + replica, + lakeBucketOffset.getSnapshotId(), + lakeLogEndOffset); + } + }); lakeBucketOffset .getMaxTimestamp() diff --git a/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManager.java b/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManager.java index 97654e3c257..cf5302e685f 100644 --- a/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManager.java +++ b/fluss-server/src/main/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManager.java @@ -37,12 +37,10 @@ import org.apache.fluss.server.entity.PutKvDataForBucket; import org.apache.fluss.server.kv.KvStateLookupResult; import org.apache.fluss.server.kv.KvStateLookupResult.Status; -import org.apache.fluss.server.kv.historical.HistoricalValueLookup; import org.apache.fluss.server.log.LogAppendInfo; import org.apache.fluss.server.replica.Replica; import org.apache.fluss.server.storage.LocalDiskManager; import org.apache.fluss.utils.ByteArraySlice; -import org.apache.fluss.utils.ByteArrayWrapper; import org.apache.fluss.utils.concurrent.Scheduler; import javax.annotation.Nullable; @@ -50,10 +48,8 @@ import java.io.File; import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.Iterator; import java.util.List; -import java.util.Map; import java.util.concurrent.CompletableFuture; import static org.apache.fluss.utils.Preconditions.checkNotNull; @@ -190,6 +186,16 @@ public void requireLakeSnapshot(long tableId, long lakeSnapshotId) { lakeLookupManager.requireLakeSnapshot(tableId, lakeSnapshotId); } + /** Publishes lake coverage used by natural RocksDB compaction for a historical leader. */ + public void onLakeProgress(Replica replica, long lakeSnapshotId, long lakeLogEndOffset) { + replica.tryUpdateHistoricalCleanupOffset( + lakeLogEndOffset, + () -> + // Future fallback lookups must require the covering snapshot before local + // entries become eligible for physical removal. + requireLakeSnapshot(replica.getTableBucket().getTableId(), lakeSnapshotId)); + } + /** Returns the number of accepted historical operations that have not completed. */ public int numInflightRequests() { return taskExecutor.numInflightRequests(); @@ -226,56 +232,23 @@ LogAppendInfo processPut( ResolvedPartitionSpec.fromPartitionName( tableInfo.getPartitionKeys(), originalPartitionName); // The public put path holds the TableBucket ordering slot until processPut returns, so - // local state cannot be changed by a later historical write between resolve and apply. + // local state cannot be changed by a later historical write between lookup and apply. int expectedLeaderEpoch = replica.getLeaderEpoch(); - List keysRequiringLakeLookup = - replica.findKeysRequiringLakeLookup( - putData.records(), - targetColumns, - mergeMode, - originalPartitionName, - expectedLeaderEpoch, - requiredAcks); - - Map lakeResults = new HashMap<>(); - if (!keysRequiringLakeLookup.isEmpty()) { - List lakeValues = - lakeLookupManager.lookup( - new LookupDataForBucket( - putData.tableBucket(), - keysRequiringLakeLookup, - originalPartitionName), - tableInfo, - replica.getLatestSchemaInfo(), - originalPartitionSpec, - replica.tableMetrics()::recordHistoricalLakeLookup); - for (int i = 0; i < keysRequiringLakeLookup.size(); i++) { - byte[] lakeValue = lakeValues.get(i); - lakeResults.put( - new ByteArrayWrapper(keysRequiringLakeLookup.get(i)), - lakeValue == null - ? KvStateLookupResult.notFound() - : KvStateLookupResult.present(lakeValue)); - } - } - - HistoricalValueLookup memoizedLakeLookup = - primaryKey -> { - KvStateLookupResult result = - checkNotNull( - lakeResults.get(new ByteArrayWrapper(primaryKey)), - "No resolved lake value for a historical write key"); - return result.value(); - }; - - // TODO: Tag historical values and tombstones with WAL offsets for incremental cleanup; see - // https://github.com/apache/fluss/issues/4159. return replica.putHistoricalRecordsToLeader( putData.records(), targetColumns, mergeMode, originalPartitionName, - memoizedLakeLookup, + lakeLookupKeys -> + lakeLookupManager.lookup( + new LookupDataForBucket( + putData.tableBucket(), + lakeLookupKeys, + originalPartitionName), + tableInfo, + replica.getLatestSchemaInfo(), + originalPartitionSpec, + replica.tableMetrics()::recordHistoricalLakeLookup), expectedLeaderEpoch, requiredAcks); } @@ -343,15 +316,13 @@ private LookupResultForBucket lookupInternal( Iterator lakeValueIterator = lakeValues.iterator(); List values = new ArrayList<>(localResults.size()); - KvValueLayout localValueLayout = - KvValueLayout.fromTableConfig(tableInfo.getTableConfig()); for (KvStateLookupResult localResult : localResults) { // Consume one lake value for each NOT_FOUND result; local values and tombstones // keep their original positions without advancing the lake iterator. if (localResult.status() == Status.NOT_FOUND) { values.add(KvValueLayout.PLAIN.toValueBodySlice(lakeValueIterator.next())); } else { - values.add(localValueLayout.toValueBodySlice(localResult.value())); + values.add(KvValueLayout.TAGGED.toValueBodySlice(localResult.value())); } } return new LookupResultForBucket(tableBucket, values, originalPartitionName); diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/HistoricalKvCompactionFilterTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/HistoricalKvCompactionFilterTest.java new file mode 100644 index 00000000000..0b076974128 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/HistoricalKvCompactionFilterTest.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv; + +import org.apache.fluss.record.BinaryValue; +import org.apache.fluss.rocksdb.RocksDBHandle; +import org.apache.fluss.row.BinaryRow; +import org.apache.fluss.row.encode.KvValueLayout; +import org.apache.fluss.row.encode.ValueEncoder; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; + +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.rocksdb.ColumnFamilyOptions; +import org.rocksdb.DBOptions; +import org.rocksdb.FlinkCompactionFilter; +import org.rocksdb.FlushOptions; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.concurrent.atomic.AtomicLong; + +import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE; +import static org.apache.fluss.record.TestData.DEFAULT_SCHEMA_ID; +import static org.apache.fluss.testutils.DataTestUtils.compactedRow; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests offset-distance retention for historical values and tombstones. */ +class HistoricalKvCompactionFilterTest { + + @TempDir private Path tempDir; + + @ParameterizedTest + @ValueSource(longs = {0L, 60_000L}) + void testRemovesOnlyOffsetsBeyondRetentionDistance(long retentionOffsetDistance) + throws Exception { + AtomicLong cleanupOffset = new AtomicLong(0L); + byte[] valueBeforeKey = bytes("value-before"); + byte[] tombstoneBeforeKey = bytes("tombstone-before"); + byte[] valueAtKey = bytes("value-at"); + byte[] tombstoneAtKey = bytes("tombstone-at"); + byte[] valueAfterKey = bytes("value-after"); + byte[] tombstoneAfterKey = bytes("tombstone-after"); + BinaryRow row = compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}); + + try (FlinkCompactionFilter.FlinkCompactionFilterFactory filterFactory = + RowTtlCompactionFilterFactory.create( + KvValueLayout.TAGGED, + retentionOffsetDistance, + 1L, + () -> cleanupOffset.get() - 1L); + DBOptions dbOptions = new DBOptions().setCreateIfMissing(true); + ColumnFamilyOptions cfOptions = + new ColumnFamilyOptions().setCompactionFilterFactory(filterFactory); + RocksDBHandle handle = new RocksDBHandle(tempDir.toFile(), dbOptions, cfOptions); + FlushOptions flushOptions = new FlushOptions().setWaitForFlush(true)) { + handle.openDB(); + handle.getDb().put(valueBeforeKey, encodeValue(row, 4L)); + handle.getDb().put(tombstoneBeforeKey, HistoricalKvTombstone.encode(4L)); + handle.getDb().put(valueAtKey, encodeValue(row, 5L)); + handle.getDb().put(tombstoneAtKey, HistoricalKvTombstone.encode(5L)); + handle.getDb().put(valueAfterKey, encodeValue(row, 6L)); + handle.getDb().put(tombstoneAfterKey, HistoricalKvTombstone.encode(6L)); + handle.getDb().flush(flushOptions); + + handle.getDb().compactRange(); + assertThat(handle.getDb().get(valueBeforeKey)).isNotNull(); + assertThat(handle.getDb().get(tombstoneBeforeKey)).isNotNull(); + + // At this boundary, offset 4 is still inside the retained range. + cleanupOffset.set(retentionOffsetDistance + 4L); + handle.getDb().compactRange(); + assertThat(handle.getDb().get(valueBeforeKey)).isNotNull(); + assertThat(handle.getDb().get(tombstoneBeforeKey)).isNotNull(); + + // Advancing by one makes offset 4 eligible, while offsets 5 and 6 remain retained. + cleanupOffset.set(retentionOffsetDistance + 5L); + handle.getDb().compactRange(); + + assertThat(handle.getDb().get(valueBeforeKey)).isNull(); + assertThat(handle.getDb().get(tombstoneBeforeKey)).isNull(); + assertThat(handle.getDb().get(valueAtKey)).isNotNull(); + assertThat(handle.getDb().get(tombstoneAtKey)).isNotNull(); + assertThat(handle.getDb().get(valueAfterKey)).isNotNull(); + assertThat(handle.getDb().get(tombstoneAfterKey)).isNotNull(); + } + } + + private static byte[] encodeValue(BinaryRow row, long logOffset) { + return ValueEncoder.forLayout(KvValueLayout.TAGGED) + .encodeValue(new BinaryValue(DEFAULT_SCHEMA_ID, row), logOffset); + } + + private static byte[] bytes(String value) { + return value.getBytes(StandardCharsets.UTF_8); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/KvStateAccessorTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvStateAccessorTest.java new file mode 100644 index 00000000000..fb64323078e --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/KvStateAccessorTest.java @@ -0,0 +1,109 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv; + +import org.apache.fluss.metadata.KvFormat; +import org.apache.fluss.record.BinaryValue; +import org.apache.fluss.record.TestingSchemaGetter; +import org.apache.fluss.row.encode.KvValueLayout; +import org.apache.fluss.row.encode.ValueDecoder; +import org.apache.fluss.row.encode.ValueEncoder; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer; +import org.apache.fluss.server.kv.prewrite.KvPreWriteBuffer.Key; +import org.apache.fluss.server.kv.rocksdb.RocksDBKv; +import org.apache.fluss.server.metrics.group.TestingMetricGroups; + +import org.junit.jupiter.api.Test; + +import java.util.function.Supplier; + +import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE; +import static org.apache.fluss.record.TestData.DATA1_SCHEMA; +import static org.apache.fluss.record.TestData.DEFAULT_SCHEMA_ID; +import static org.apache.fluss.testutils.DataTestUtils.compactedRow; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +/** Tests for {@link KvStateAccessor}. */ +class KvStateAccessorTest { + + private final RocksDBKv rocksDBKv = mock(RocksDBKv.class); + private final KvStateAccessor stateAccessor = + new KvStateAccessor( + new KvPreWriteBuffer(TestingMetricGroups.TABLET_SERVER_METRICS), + rocksDBKv, + true); + private final ValueDecoder valueDecoder = + new ValueDecoder( + new TestingSchemaGetter(DEFAULT_SCHEMA_ID, DATA1_SCHEMA), + KvFormat.COMPACTED, + KvValueLayout.TAGGED); + private final Key key = stateAccessor.encodeKey(new byte[] {1}, "20240107"); + + @Test + void testReusesSavedValueWithoutRocksDBLookup() throws Exception { + BinaryValue savedValue = binaryValue("saved"); + + assertThat(stateAccessor.lookup(key, valueDecoder, () -> savedValue)).isSameAs(savedValue); + assertThat(stateAccessor.lookup(key, valueDecoder, () -> null)).isNull(); + verifyNoInteractions(rocksDBKv); + } + + @Test + void testBufferedUpdateAndDeleteOverrideSavedValue() throws Exception { + Supplier savedLookup = + () -> { + throw new AssertionError( + "A buffered mutation must stop fallback to saved state"); + }; + BinaryValue updatedValue = binaryValue("updated"); + stateAccessor.update( + key, + ValueEncoder.forLayout(KvValueLayout.TAGGED).encodeValue(updatedValue, 0L), + 0L); + assertThat(stateAccessor.lookup(key, valueDecoder, savedLookup)).isEqualTo(updatedValue); + assertThat(stateAccessor.lookup(key, valueDecoder, null)).isEqualTo(updatedValue); + + // A delete must mask the saved value so a later record cannot resurrect it. + stateAccessor.delete(key, 1L); + assertThat(stateAccessor.lookup(key, valueDecoder, savedLookup)).isNull(); + assertThat(stateAccessor.lookup(key, valueDecoder, null)).isNull(); + verifyNoInteractions(rocksDBKv); + } + + @Test + void testLocalLookupPreservesMissingAndDeletedStates() throws Exception { + BinaryValue value = binaryValue("local"); + byte[] encodedValue = ValueEncoder.forLayout(KvValueLayout.TAGGED).encodeValue(value, 1L); + when(rocksDBKv.get(key.get())) + .thenReturn(null, HistoricalKvTombstone.encode(0L), encodedValue); + + // The local probe must schedule lake lookup only for a true miss, not a tombstone. + assertThat(stateAccessor.lookupLocal(key)).isEqualTo(KvStateLookupResult.notFound()); + assertThat(stateAccessor.lookupLocal(key)).isEqualTo(KvStateLookupResult.deleted()); + assertThat(stateAccessor.lookup(key, valueDecoder, null)).isEqualTo(value); + } + + private static BinaryValue binaryValue(String value) { + return new BinaryValue( + DEFAULT_SCHEMA_ID, compactedRow(DATA1_ROW_TYPE, new Object[] {1, value})); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/RowTtlCompactionFilterTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/RowTtlCompactionFilterTest.java index 762b3756c2e..644922c6d5a 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/kv/RowTtlCompactionFilterTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/RowTtlCompactionFilterTest.java @@ -56,10 +56,7 @@ void testFlinkCompactionFilterReadsTimestampFromTaggedValue() throws Exception { try (FlinkCompactionFilter.FlinkCompactionFilterFactory filterFactory = RowTtlCompactionFilterFactory.create( - KvValueLayout.TAGGED, - Duration.ofHours(1L), - 1L, - new ManualClock(now)); + KvValueLayout.TAGGED, Duration.ofHours(1L), new ManualClock(now)); DBOptions dbOptions = new DBOptions().setCreateIfMissing(true); ColumnFamilyOptions cfOptions = new ColumnFamilyOptions().setCompactionFilterFactory(filterFactory); @@ -83,10 +80,7 @@ void testCreateRejectsInvalidTtlDuration() { assertThatThrownBy( () -> RowTtlCompactionFilterFactory.create( - KvValueLayout.TAGGED, - Duration.ZERO, - 1L, - new ManualClock(0L))) + KvValueLayout.TAGGED, Duration.ZERO, new ManualClock(0L))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining(ConfigOptions.TABLE_KV_TTL.key()); } diff --git a/fluss-server/src/test/java/org/apache/fluss/server/kv/historical/LocalValueLookupResultTest.java b/fluss-server/src/test/java/org/apache/fluss/server/kv/historical/LocalValueLookupResultTest.java new file mode 100644 index 00000000000..f04d17780e3 --- /dev/null +++ b/fluss-server/src/test/java/org/apache/fluss/server/kv/historical/LocalValueLookupResultTest.java @@ -0,0 +1,101 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.server.kv.historical; + +import org.apache.fluss.metadata.KvFormat; +import org.apache.fluss.record.BinaryValue; +import org.apache.fluss.record.TestingSchemaGetter; +import org.apache.fluss.row.BinaryRow; +import org.apache.fluss.row.encode.KvValueLayout; +import org.apache.fluss.row.encode.ValueDecoder; +import org.apache.fluss.row.encode.ValueEncoder; +import org.apache.fluss.server.kv.KvStateLookupResult; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE; +import static org.apache.fluss.record.TestData.DATA1_SCHEMA; +import static org.apache.fluss.record.TestData.DEFAULT_SCHEMA_ID; +import static org.apache.fluss.testutils.DataTestUtils.compactedRow; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link LocalValueLookupResult}. */ +class LocalValueLookupResultTest { + + @Test + void testCreatesLookupFromLocalAndLakeResults() { + LocalValueLookupResult localLookupResult = createLookupResult(); + byte[] localValueKey = new byte[] {1}; + byte[] localDeleteKey = new byte[] {2}; + byte[] lakeValueKey = new byte[] {3}; + byte[] lakeMissKey = new byte[] {4}; + BinaryValue localValue = binaryValue(1, "local"); + BinaryValue lakeValue = binaryValue(3, "lake"); + + localLookupResult.add( + localValueKey, + KvStateLookupResult.present( + ValueEncoder.forLayout(KvValueLayout.TAGGED).encodeValue(localValue, 10L))); + localLookupResult.add(localDeleteKey, KvStateLookupResult.deleted()); + localLookupResult.add(lakeValueKey, KvStateLookupResult.notFound()); + localLookupResult.add(lakeMissKey, KvStateLookupResult.notFound()); + + assertThat(localLookupResult.keysMissingLocally()) + .containsExactly(lakeValueKey, lakeMissKey); + + HistoricalValueLookup valueLookup = + localLookupResult.createValueLookup( + Arrays.asList( + ValueEncoder.forLayout(KvValueLayout.PLAIN).encodeValue(lakeValue), + null)); + + assertThat(valueLookup.lookup(localValueKey)).isEqualTo(localValue); + assertThat(valueLookup.lookup(localDeleteKey)).isNull(); + assertThat(valueLookup.lookup(lakeValueKey)).isEqualTo(lakeValue); + assertThat(valueLookup.lookup(lakeMissKey)).isNull(); + assertThatThrownBy(() -> valueLookup.lookup(new byte[] {5})) + .isInstanceOf(NullPointerException.class) + .hasMessageContaining("No previous value for a historical write key"); + } + + @Test + void testValidatesLakeResultCountBeforeCreatingLookup() { + LocalValueLookupResult localLookupResult = createLookupResult(); + localLookupResult.add(new byte[] {1}, KvStateLookupResult.notFound()); + + assertThatThrownBy(() -> localLookupResult.createValueLookup(Collections.emptyList())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Expected 1 historical lake values, but received 0"); + } + + private static LocalValueLookupResult createLookupResult() { + TestingSchemaGetter schemaGetter = new TestingSchemaGetter(DEFAULT_SCHEMA_ID, DATA1_SCHEMA); + return new LocalValueLookupResult( + new ValueDecoder(schemaGetter, KvFormat.COMPACTED, KvValueLayout.TAGGED), + new ValueDecoder(schemaGetter, KvFormat.COMPACTED, KvValueLayout.PLAIN)); + } + + private static BinaryValue binaryValue(int id, String value) { + BinaryRow row = compactedRow(DATA1_ROW_TYPE, new Object[] {id, value}); + return new BinaryValue(DEFAULT_SCHEMA_ID, row); + } +} diff --git a/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java index b18f96b3a7a..dbce8b94b0d 100644 --- a/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java +++ b/fluss-server/src/test/java/org/apache/fluss/server/replica/historical/HistoricalPartitionManagerTest.java @@ -25,6 +25,7 @@ import org.apache.fluss.exception.HistoricalPartitionThrottledException; import org.apache.fluss.exception.InvalidPartitionException; import org.apache.fluss.lake.lakestorage.LakeTableLookuper; +import org.apache.fluss.memory.MemorySegment; import org.apache.fluss.metadata.ChangelogImage; import org.apache.fluss.metadata.DataLakeFormat; import org.apache.fluss.metadata.PhysicalTablePath; @@ -50,16 +51,21 @@ import org.apache.fluss.rpc.entity.FetchLogResultForBucket; import org.apache.fluss.rpc.entity.LookupResultForBucket; import org.apache.fluss.rpc.entity.PutKvResultForBucket; +import org.apache.fluss.rpc.messages.NotifyLakeTableOffsetResponse; import org.apache.fluss.rpc.protocol.ApiKeys; import org.apache.fluss.rpc.protocol.Errors; import org.apache.fluss.rpc.protocol.MergeMode; import org.apache.fluss.server.entity.FetchReqInfo; +import org.apache.fluss.server.entity.LakeBucketOffset; import org.apache.fluss.server.entity.LookupDataForBucket; +import org.apache.fluss.server.entity.NotifyLakeTableOffsetData; import org.apache.fluss.server.entity.NotifyLeaderAndIsrData; import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket; import org.apache.fluss.server.entity.PutKvDataForBucket; import org.apache.fluss.server.kv.KvStateLookupResult; import org.apache.fluss.server.kv.KvTablet; +import org.apache.fluss.server.kv.historical.HistoricalKvKeyEncoder; +import org.apache.fluss.server.kv.historical.HistoricalKvTombstone; import org.apache.fluss.server.log.FetchParams; import org.apache.fluss.server.log.LogAppendInfo; import org.apache.fluss.server.metadata.BucketMetadata; @@ -83,9 +89,11 @@ import com.github.benmanes.caffeine.cache.Scheduler; import com.github.benmanes.caffeine.cache.Ticker; import org.junit.jupiter.api.Test; +import org.rocksdb.FlushOptions; import javax.annotation.Nullable; +import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; @@ -108,6 +116,7 @@ import static org.apache.fluss.testutils.DataTestUtils.genKvRecordBatch; import static org.apache.fluss.testutils.DataTestUtils.row; import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow; +import static org.apache.fluss.testutils.common.CommonTestUtils.retry; import static org.apache.fluss.utils.PartitionUtils.HISTORICAL_PARTITION_VALUE; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -123,6 +132,10 @@ class HistoricalPartitionManagerTest extends ReplicaTestBase { private static final String ANOTHER_ORIGINAL_PARTITION = "20240108"; private static final String HISTORICAL_PARTITION = HISTORICAL_PARTITION_VALUE; private static final TableBucket TABLE_BUCKET = new TableBucket(TABLE_ID, PARTITION_ID, 0); + private static final RowType KEY_TYPE = + DataTypes.ROW( + new DataField("id", DataTypes.INT()), + new DataField("region", DataTypes.STRING())); @Test void testResolvesMultipleLakeMissesWithoutPrewriteRollback() throws Exception { @@ -132,45 +145,30 @@ void testResolvesMultipleLakeMissesWithoutPrewriteRollback() throws Exception { assertThat(kvTablet).isNotNull(); TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); - HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); - - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); - CompactedKeyEncoder keyEncoder = new CompactedKeyEncoder(keyType); - byte[] firstKey = keyEncoder.encodeKey(row(1, "us")); - byte[] secondKey = keyEncoder.encodeKey(row(2, "eu")); + + byte[] firstKey = key(1, "us"); + byte[] secondKey = key(2, "eu"); + byte[] thirdKey = key(3, "ap"); KvRecordBatch insertBatch = batch( - keyType, tableInfo.getRowType(), - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", ORIGINAL_PARTITION, "v1"}), + upsert(1, "us", ORIGINAL_PARTITION, "v1"), // The second record reuses the first record's staged state and must not // trigger another lake lookup for the same key. - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", ORIGINAL_PARTITION, "v1-updated"}), - Tuple2.of( - new Object[] {2, "eu"}, - new Object[] {2, "eu", ORIGINAL_PARTITION, "v2"})); + upsert(1, "us", ORIGINAL_PARTITION, "v1-updated"), + upsert(2, "eu", ORIGINAL_PARTITION, "v2"), + // The absent delete is a no-op. The following upsert must still see the + // resolved absence without another lake lookup. + delete(3, "ap"), + upsert(3, "ap", ORIGINAL_PARTITION, "v3")); long truncateCount = replicaManager.getServerMetricGroup().kvTruncateAsErrorCount().getCount(); - try { - historicalPartitionManager.processPut( - replica, - new PutKvDataForBucket(TABLE_BUCKET, insertBatch, ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + writeBatch(historicalPartitionManager, replica, ORIGINAL_PARTITION, insertBatch); - assertThat(lakeLookupManager.lookupCount).hasValue(2); + assertThat(lakeLookupManager.lookupCount).hasValue(3); assertThat(lakeLookupManager.lookupBatchCount).hasValue(1); assertThat(replicaManager.getServerMetricGroup().kvTruncateAsErrorCount().getCount()) .isEqualTo(truncateCount); @@ -187,8 +185,91 @@ void testResolvesMultipleLakeMissesWithoutPrewriteRollback() throws Exception { secondKey, tableInfo, row(2, "eu", ORIGINAL_PARTITION, "v2")); - } finally { - historicalPartitionManager.close(); + assertHistoricalValue( + kvTablet, + ORIGINAL_PARTITION, + thirdKey, + tableInfo, + row(3, "ap", ORIGINAL_PARTITION, "v3")); + } + } + + @Test + void testReusesLocalProbeResultsWhenEntriesDisappearBeforeApply() throws Exception { + TableInfo tableInfo = registerHistoricalTableAndBecomeLeader(); + Replica replica = replicaManager.getReplicaOrException(TABLE_BUCKET); + KvTablet kvTablet = replica.getKvTablet(); + assertThat(kvTablet).isNotNull(); + TestingHistoricalLakeLookupManager lakeLookupManager = + new TestingHistoricalLakeLookupManager(lookupConfiguration()); + + RowType rowType = tableInfo.getRowType(); + byte[] valueKey = key(1, "us"); + byte[] tombstoneKey = key(2, "eu"); + byte[] missingKey = key(3, "ap"); + + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + // Seed exactly the two local states that compaction can remove between phases. + kvTablet.getRocksDBKv() + .put( + HistoricalKvKeyEncoder.encode(ORIGINAL_PARTITION, valueKey), + ValueEncoder.forLayout(KvValueLayout.TAGGED) + .encodeValue( + new BinaryValue( + (short) tableInfo.getSchemaId(), + compactedRow( + rowType, + new Object[] { + 1, "us", ORIGINAL_PARTITION, "v1" + })), + 0L)); + kvTablet.getRocksDBKv() + .put( + HistoricalKvKeyEncoder.encode(ORIGINAL_PARTITION, tombstoneKey), + HistoricalKvTombstone.encode(1L)); + + // The third key forces lake I/O after all three keys have been probed. Remove the + // value and tombstone during that I/O to model compaction between probe and apply. + lakeLookupManager.setLookupHook( + () -> { + deleteHistoricalRocksDbKey(kvTablet, valueKey); + deleteHistoricalRocksDbKey(kvTablet, tombstoneKey); + }); + + LogAppendInfo appendInfo = + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch( + rowType, + upsert(1, "us", ORIGINAL_PARTITION, "v2"), + upsert(2, "eu", ORIGINAL_PARTITION, "recreated"), + upsert(3, "ap", ORIGINAL_PARTITION, "inserted"))); + + assertThat(appendInfo.numMessages()).isEqualTo(4); + assertThat(lakeLookupManager.lookupCount).hasValue(1); + assertThat(lakeLookupManager.lookupBatchCount).hasValue(1); + flushAndWait(kvTablet, Long.MAX_VALUE); + assertHistoricalValue( + kvTablet, + ORIGINAL_PARTITION, + valueKey, + tableInfo, + row(1, "us", ORIGINAL_PARTITION, "v2")); + assertHistoricalValue( + kvTablet, + ORIGINAL_PARTITION, + tombstoneKey, + tableInfo, + row(2, "eu", ORIGINAL_PARTITION, "recreated")); + assertHistoricalValue( + kvTablet, + ORIGINAL_PARTITION, + missingKey, + tableInfo, + row(3, "ap", ORIGINAL_PARTITION, "inserted")); } } @@ -200,31 +281,14 @@ void testWalFullRowUpsertDoesNotLookupLake() throws Exception { assertThat(kvTablet).isNotNull(); TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); - HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); - - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); - byte[] primaryKey = new CompactedKeyEncoder(keyType).encodeKey(row(1, "us")); + + byte[] primaryKey = key(1, "us"); KvRecordBatch insertBatch = - batch( - keyType, - tableInfo.getRowType(), - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", ORIGINAL_PARTITION, "v1"})); + batch(tableInfo.getRowType(), upsert(1, "us", ORIGINAL_PARTITION, "v1")); - try { - historicalPartitionManager.processPut( - replica, - new PutKvDataForBucket(TABLE_BUCKET, insertBatch, ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + writeBatch(historicalPartitionManager, replica, ORIGINAL_PARTITION, insertBatch); assertThat(lakeLookupManager.lookupCount).hasValue(0); assertThat(lakeLookupManager.lookupBatchCount).hasValue(0); @@ -235,8 +299,6 @@ void testWalFullRowUpsertDoesNotLookupLake() throws Exception { primaryKey, tableInfo, row(1, "us", ORIGINAL_PARTITION, "v1")); - } finally { - historicalPartitionManager.close(); } } @@ -249,36 +311,20 @@ void testHistoricalInsertUpdateAndDelete() throws Exception { assertThat(kvManager.getKv(TABLE_BUCKET)).contains(kvTablet); TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); - HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); RowType rowType = tableInfo.getRowType(); - byte[] primaryKey = new CompactedKeyEncoder(keyType).encodeKey(row(1, "us")); + byte[] primaryKey = key(1, "us"); - try { + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { // The first write misses both local state and lake, so it creates a local overlay. - KvRecordBatch insertBatch = - batch( - keyType, - rowType, - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", "20240107", "v1"})); + KvRecordBatch insertBatch = batch(rowType, upsert(1, "us", "20240107", "v1")); assertThat( - historicalPartitionManager - .processPut( + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, insertBatch, ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1) + ORIGINAL_PARTITION, + insertBatch) .lastOffset()) .isZero(); flushAndWait(kvTablet, Long.MAX_VALUE); @@ -293,19 +339,12 @@ void testHistoricalInsertUpdateAndDelete() throws Exception { // The same primary key in another original partition must use a separate state entry. KvRecordBatch anotherPartitionBatch = - batch( - keyType, - rowType, - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", "20240108", "another"})); - historicalPartitionManager.processPut( + batch(rowType, upsert(1, "us", "20240108", "another")); + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, anotherPartitionBatch, ANOTHER_ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); + ANOTHER_ORIGINAL_PARTITION, + anotherPartitionBatch); flushAndWait(kvTablet, Long.MAX_VALUE); assertHistoricalValue( kvTablet, @@ -322,13 +361,7 @@ void testHistoricalInsertUpdateAndDelete() throws Exception { assertThat(lakeLookupManager.lookupCount).hasValue(2); // Exercise the ReplicaManager entry point; the update should reuse the local overlay. - KvRecordBatch updateBatch = - batch( - keyType, - rowType, - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", "20240107", "v2"})); + KvRecordBatch updateBatch = batch(rowType, upsert(1, "us", "20240107", "v2")); CompletableFuture> updateResponse = new CompletableFuture<>(); assertThat(replica.tableMetrics().totalHistoricalPutKvRequests().getCount()).isZero(); @@ -386,14 +419,8 @@ void testHistoricalInsertUpdateAndDelete() throws Exception { assertThat(lookedUpValue.row.getString(3)).isEqualTo(BinaryString.fromString("v2")); // Keep a tombstone locally so a later lookup cannot resurrect the value from lake. - KvRecordBatch deleteBatch = - batch(keyType, rowType, Tuple2.of(new Object[] {1, "us"}, null)); - historicalPartitionManager.processPut( - replica, - new PutKvDataForBucket(TABLE_BUCKET, deleteBatch, ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); + KvRecordBatch deleteBatch = batch(rowType, delete(1, "us")); + writeBatch(historicalPartitionManager, replica, ORIGINAL_PARTITION, deleteBatch); flushAndWait(kvTablet, Long.MAX_VALUE); assertThat(kvTablet.lookupHistoricalLocal(ORIGINAL_PARTITION, primaryKey)) .isEqualTo(KvStateLookupResult.deleted()); @@ -425,8 +452,204 @@ void testHistoricalInsertUpdateAndDelete() throws Exception { replica.putRecordsToLeader( insertBatch, null, MergeMode.DEFAULT, 1)) .isInstanceOf(InvalidPartitionException.class); - } finally { - historicalPartitionManager.close(); + } + } + + @Test + void testHistoricalMutationsUseProducingWalOffsetsAsTags() throws Exception { + TableInfo tableInfo = registerHistoricalTableAndBecomeLeader(); + Replica replica = replicaManager.getReplicaOrException(TABLE_BUCKET); + KvTablet kvTablet = replica.getKvTablet(); + assertThat(kvTablet).isNotNull(); + RowType rowType = tableInfo.getRowType(); + byte[] primaryKey = key(1, "us"); + + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager( + new TestingHistoricalLakeLookupManager(lookupConfiguration()))) { + LogAppendInfo insertAppend = + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch(rowType, upsert(1, "us", ORIGINAL_PARTITION, "v1"))); + flushAndWait(kvTablet, Long.MAX_VALUE); + assertHistoricalValueTag( + kvTablet, ORIGINAL_PARTITION, primaryKey, insertAppend.lastOffset()); + + LogAppendInfo updateAppend = + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch(rowType, upsert(1, "us", ORIGINAL_PARTITION, "v2"))); + assertThat(updateAppend.numMessages()).isEqualTo(2); + flushAndWait(kvTablet, Long.MAX_VALUE); + // In FULL changelog mode, UPDATE_AFTER is the second WAL record and produces state. + assertHistoricalValueTag( + kvTablet, ORIGINAL_PARTITION, primaryKey, updateAppend.lastOffset()); + + LogAppendInfo deleteAppend = + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch(rowType, delete(1, "us"))); + flushAndWait(kvTablet, Long.MAX_VALUE); + byte[] tombstone = historicalRawValue(kvTablet, ORIGINAL_PARTITION, primaryKey); + assertThat(tombstone).hasSize(Long.BYTES); + assertHistoricalValueTag( + kvTablet, ORIGINAL_PARTITION, primaryKey, deleteAppend.lastOffset()); + } + } + + @Test + void testCompactionRespectsRetentionAndFallsBackToLake() throws Exception { + TableInfo tableInfo = registerHistoricalTableAndBecomeLeader(ChangelogImage.WAL); + Replica replica = replicaManager.getReplicaOrException(TABLE_BUCKET); + KvTablet kvTablet = replica.getKvTablet(); + assertThat(kvTablet).isNotNull(); + TestingHistoricalLakeLookupManager lakeLookupManager = + new TestingHistoricalLakeLookupManager(lookupConfiguration()); + byte[] coveredKey = key(1, "us"); + byte[] retainedKey = key(2, "eu"); + int retentionOffsetDistance = 60_000; + long cleanupOffset = retentionOffsetDistance + 1L; + + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch( + tableInfo.getRowType(), + upsert(1, "us", ORIGINAL_PARTITION, "covered"), + upsert(2, "eu", ORIGINAL_PARTITION, "retained"))); + // Repeated WAL-mode updates advance the log by one offset each without creating + // thousands of distinct local keys. This fills the hardcoded retention window. + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + genKvRecordBatch( + KEY_TYPE, + tableInfo.getRowType(), + Collections.nCopies( + retentionOffsetDistance, + upsert(3, "ap", ORIGINAL_PARTITION, "padding")))); + assertThat(replica.getLocalLogEndOffset()).isEqualTo(cleanupOffset + 1L); + flushAndWait(kvTablet, Long.MAX_VALUE); + assertHistoricalValueTag(kvTablet, ORIGINAL_PARTITION, coveredKey, 0L); + assertHistoricalValueTag(kvTablet, ORIGINAL_PARTITION, retainedKey, 1L); + + historicalPartitionManager.onLakeProgress(replica, 6L, retentionOffsetDistance); + + // Lake coverage alone does not remove values still within the retained offset range. + compactHistoricalKv(kvTablet); + assertHistoricalValueTag(kvTablet, ORIGINAL_PARTITION, coveredKey, 0L); + assertHistoricalValueTag(kvTablet, ORIGINAL_PARTITION, retainedKey, 1L); + + historicalPartitionManager.onLakeProgress(replica, 7L, cleanupOffset); + + assertThat(kvTablet.getHistoricalCleanupOffset()).isEqualTo(cleanupOffset); + compactHistoricalKv(kvTablet); + assertThat(historicalRawValue(kvTablet, ORIGINAL_PARTITION, coveredKey)).isNull(); + assertHistoricalValueTag(kvTablet, ORIGINAL_PARTITION, retainedKey, 1L); + + // A miss caused by cleanup must enter lake lookup only after the covering snapshot + // token has been published. + lakeLookupManager.putLakeValue( + ORIGINAL_PARTITION, + ValueEncoder.encodeValue( + (short) tableInfo.getSchemaId(), + compactedRow( + tableInfo.getRowType(), + new Object[] { + 1, "us", ORIGINAL_PARTITION, "covered-from-lake" + }))); + lakeLookupManager.setLookupHook( + () -> { + assertThat(lakeLookupManager.requiredSnapshotCount).hasValue(2); + assertThat(kvTablet.getHistoricalCleanupOffset()).isEqualTo(cleanupOffset); + }); + LookupResultForBucket fallbackResult = + historicalPartitionManager + .lookup( + replica, + new LookupDataForBucket( + TABLE_BUCKET, + Collections.singletonList(coveredKey), + ORIGINAL_PARTITION), + (lookupTimeNanos, lookupFileDownloaded) -> {}) + .get(10, TimeUnit.SECONDS); + assertThat(fallbackResult.failed()).isFalse(); + BinaryValue fallbackValue = + new ValueDecoder( + schemaGetter(tableInfo), + tableInfo.getTableConfig().getKvFormat(), + KvValueLayout.PLAIN) + .decodeValue(fallbackResult.lookupValues().get(0).toByteArray()); + assertThat(fallbackValue.row.getString(3)) + .isEqualTo(BinaryString.fromString("covered-from-lake")); + + historicalPartitionManager.onLakeProgress(replica, 8L, cleanupOffset + 1L); + compactHistoricalKv(kvTablet); + assertThat(historicalRawValue(kvTablet, ORIGINAL_PARTITION, retainedKey)).isNull(); + } + } + + @Test + void testLakeProgressPublishesSnapshotBeforeCleanupOffset() throws Exception { + TableInfo tableInfo = registerHistoricalTableAndBecomeLeader(ChangelogImage.WAL); + Replica replica = replicaManager.getReplicaOrException(TABLE_BUCKET); + KvTablet kvTablet = replica.getKvTablet(); + assertThat(kvTablet).isNotNull(); + TestingHistoricalLakeLookupManager lakeLookupManager = + new TestingHistoricalLakeLookupManager(lookupConfiguration()); + + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch( + tableInfo.getRowType(), + upsert(1, "us", ORIGINAL_PARTITION, "v1"), + upsert(2, "eu", ORIGINAL_PARTITION, "v2"))); + + // The covering snapshot must be published before entries become eligible for cleanup. + lakeLookupManager.setRequireSnapshotHook( + () -> assertThat(kvTablet.getHistoricalCleanupOffset()).isZero()); + historicalPartitionManager.onLakeProgress(replica, 7L, 1L); + assertThat(lakeLookupManager.requiredSnapshotCount).hasValue(1); + assertThat(kvTablet.getHistoricalCleanupOffset()).isOne(); + + // A new opaque snapshot token at the same offset is still published to the lookuper. + lakeLookupManager.setRequireSnapshotHook( + () -> assertThat(kvTablet.getHistoricalCleanupOffset()).isOne()); + historicalPartitionManager.onLakeProgress(replica, 8L, 1L); + assertThat(lakeLookupManager.requiredSnapshotCount).hasValue(2); + assertThat(kvTablet.getHistoricalCleanupOffset()).isOne(); + + // A regressing cleanup offset is ignored before it can publish an older required + // snapshot. + historicalPartitionManager.onLakeProgress(replica, 9L, 0L); + assertThat(lakeLookupManager.requiredSnapshotCount).hasValue(2); + assertThat(kvTablet.getHistoricalCleanupOffset()).isOne(); + + // The RPC entry point uses ReplicaManager's own historical manager. + CompletableFuture notifyFuture = + new CompletableFuture<>(); + replicaManager.notifyLakeTableOffset( + new NotifyLakeTableOffsetData( + 1, + Collections.singletonMap( + TABLE_BUCKET, new LakeBucketOffset(10L, null, 2L, null))), + notifyFuture::complete); + notifyFuture.get(10, TimeUnit.SECONDS); + assertThat(kvTablet.getHistoricalCleanupOffset()).isEqualTo(2L); } } @@ -439,20 +662,12 @@ void testUpdateAndDeleteFromLakeFallback() throws Exception { TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); - HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); RowType rowType = tableInfo.getRowType(); String updatePartition = "20240109"; String deletePartition = "20240110"; - byte[] updateKey = new CompactedKeyEncoder(keyType).encodeKey(row(1, "us")); - byte[] deleteKey = new CompactedKeyEncoder(keyType).encodeKey(row(2, "eu")); + byte[] updateKey = key(1, "us"); + byte[] deleteKey = key(2, "eu"); short schemaId = (short) tableInfo.getSchemaId(); // Seed lake-only values so the first local operations must use lake fallback. lakeLookupManager.putLakeValue( @@ -466,39 +681,35 @@ void testUpdateAndDeleteFromLakeFallback() throws Exception { schemaId, compactedRow(rowType, new Object[] {2, "eu", deletePartition, "lake-v1"}))); - try { - // Updating a lake-only value emits the before and after images. + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { + // Same-key records must see this batch's buffered values and deletes, even though + // the saved lake result still contains lake-v1. The repeated delete is a no-op. KvRecordBatch updateBatch = batch( - keyType, rowType, - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", updatePartition, "lake-v2"})); + upsert(1, "us", updatePartition, "lake-v2"), + delete(1, "us"), + delete(1, "us"), + upsert(1, "us", updatePartition, "recreated-v3"), + upsert(1, "us", updatePartition, "updated-v4")); assertThat( - historicalPartitionManager - .processPut( + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, updateBatch, updatePartition), - null, - MergeMode.DEFAULT, - 1) + updatePartition, + updateBatch) .numMessages()) - .isEqualTo(2); + .isEqualTo(6); // Deleting a lake-only value emits a delete and leaves a local tombstone. - KvRecordBatch deleteBatch = - batch(keyType, rowType, Tuple2.of(new Object[] {2, "eu"}, null)); + KvRecordBatch deleteBatch = batch(rowType, delete(2, "eu")); assertThat( - historicalPartitionManager - .processPut( + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, deleteBatch, deletePartition), - null, - MergeMode.DEFAULT, - 1) + deletePartition, + deleteBatch) .numMessages()) .isOne(); @@ -509,7 +720,7 @@ void testUpdateAndDeleteFromLakeFallback() throws Exception { updatePartition, updateKey, tableInfo, - row(1, "us", updatePartition, "lake-v2")); + row(1, "us", updatePartition, "updated-v4")); assertThat(kvTablet.lookupHistoricalLocal(deletePartition, deleteKey)) .isEqualTo(KvStateLookupResult.deleted()); assertThat(lakeLookupManager.lookupCount).hasValue(2); @@ -525,12 +736,22 @@ void testUpdateAndDeleteFromLakeFallback() throws Exception { Tuple2.of( ChangeType.UPDATE_AFTER, new Object[] {1, "us", updatePartition, "lake-v2"}), + Tuple2.of( + ChangeType.DELETE, + new Object[] {1, "us", updatePartition, "lake-v2"}), + Tuple2.of( + ChangeType.INSERT, + new Object[] {1, "us", updatePartition, "recreated-v3"}), + Tuple2.of( + ChangeType.UPDATE_BEFORE, + new Object[] {1, "us", updatePartition, "recreated-v3"}), + Tuple2.of( + ChangeType.UPDATE_AFTER, + new Object[] {1, "us", updatePartition, "updated-v4"}), Tuple2.of( ChangeType.DELETE, new Object[] {2, "eu", deletePartition, "lake-v1"})), schemaGetter(tableInfo)); - } finally { - historicalPartitionManager.close(); } } @@ -541,9 +762,7 @@ void testLeaderChangeDuringLakeLookupFencesHistoricalWrite() throws Exception { TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); + createHistoricalPartitionManager(lakeLookupManager); CountDownLatch lakeLookupStarted = new CountDownLatch(1); CountDownLatch finishLakeLookup = new CountDownLatch(1); lakeLookupManager.setLookupHook( @@ -552,17 +771,8 @@ void testLeaderChangeDuringLakeLookupFencesHistoricalWrite() throws Exception { await(finishLakeLookup); }); - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); KvRecordBatch insertBatch = - batch( - keyType, - tableInfo.getRowType(), - Tuple2.of( - new Object[] {1, "us"}, - new Object[] {1, "us", ORIGINAL_PARTITION, "v1"})); + batch(tableInfo.getRowType(), upsert(1, "us", ORIGINAL_PARTITION, "v1")); long logEndOffsetBeforeWrite = replica.getLocalLogEndOffset(); try { @@ -597,85 +807,50 @@ void testRecoversHistoricalOverlayFromLakeCommitOffset() throws Exception { Replica replica = replicaManager.getReplicaOrException(TABLE_BUCKET); TestingHistoricalLakeLookupManager lakeLookupManager = new TestingHistoricalLakeLookupManager(lookupConfiguration()); - HistoricalPartitionManager historicalPartitionManager = - new HistoricalPartitionManager( - new HistoricalPartitionTaskExecutor(lookupConfiguration()), - lakeLookupManager); - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); - CompactedKeyEncoder keyEncoder = new CompactedKeyEncoder(keyType); - byte[] tieredPrimaryKey = keyEncoder.encodeKey(row(1, "us")); - byte[] deletedPrimaryKey = keyEncoder.encodeKey(row(2, "eu")); + byte[] tieredPrimaryKey = key(1, "us"); + byte[] deletedPrimaryKey = key(2, "eu"); - try { + try (HistoricalPartitionManager historicalPartitionManager = + createHistoricalPartitionManager(lakeLookupManager)) { LogAppendInfo firstAppend = - historicalPartitionManager.processPut( + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, - batch( - keyType, - tableInfo.getRowType(), - Tuple2.of( - new Object[] {1, "us"}, - new Object[] { - 1, "us", ORIGINAL_PARTITION, "v1" - })), - ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); - historicalPartitionManager.processPut( - replica, - new PutKvDataForBucket( - TABLE_BUCKET, + ORIGINAL_PARTITION, batch( - keyType, tableInfo.getRowType(), - Tuple2.of( - new Object[] {1, "us"}, - new Object[] { - 1, "us", ANOTHER_ORIGINAL_PARTITION, "another" - })), - ANOTHER_ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); - historicalPartitionManager.processPut( - replica, - new PutKvDataForBucket( - TABLE_BUCKET, + upsert(1, "us", ORIGINAL_PARTITION, "v1"))); + LogAppendInfo anotherPartitionAppend = + writeBatch( + historicalPartitionManager, + replica, + ANOTHER_ORIGINAL_PARTITION, batch( - keyType, tableInfo.getRowType(), - Tuple2.of( - new Object[] {2, "eu"}, - new Object[] { - 2, "eu", ORIGINAL_PARTITION, "delete-me" - })), - ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); - historicalPartitionManager.processPut( + upsert(1, "us", ANOTHER_ORIGINAL_PARTITION, "another"))); + writeBatch( + historicalPartitionManager, replica, - new PutKvDataForBucket( - TABLE_BUCKET, - batch( - keyType, - tableInfo.getRowType(), - Tuple2.of(new Object[] {2, "eu"}, null)), - ORIGINAL_PARTITION), - null, - MergeMode.DEFAULT, - 1); + ORIGINAL_PARTITION, + batch( + tableInfo.getRowType(), + upsert(2, "eu", ORIGINAL_PARTITION, "delete-me"))); + LogAppendInfo deleteAppend = + writeBatch( + historicalPartitionManager, + replica, + ORIGINAL_PARTITION, + batch(tableInfo.getRowType(), delete(2, "eu"))); KvTablet kvTabletBeforeFollower = replica.getKvTablet(); assertThat(kvTabletBeforeFollower).isNotNull(); flushAndWait(kvTabletBeforeFollower, Long.MAX_VALUE); - assertThat(replica.getLogHighWatermark()).isEqualTo(replica.getLocalLogEndOffset()); + // The flush listener advances the high watermark after releasing the KV lock. + retry( + Duration.ofSeconds(10), + () -> + assertThat(replica.getLogHighWatermark()) + .isEqualTo(replica.getLocalLogEndOffset())); // Persist the exclusive end offset of the first write as the lake recovery point. The // replica has not received this offset locally, so becoming leader must load it before @@ -704,6 +879,10 @@ void testRecoversHistoricalOverlayFromLakeCommitOffset() throws Exception { KvTablet recoveredKvTablet = replica.getKvTablet(); assertThat(recoveredKvTablet).isNotNull(); assertThat(replica.getLakeLogEndOffset()).isEqualTo(lakeCommitOffset); + assertThat(recoveredKvTablet.getHistoricalCleanupOffset()).isEqualTo(lakeCommitOffset); + // Recovery retries can reuse the same tablet and restore offset. + assertThat(recoveredKvTablet.advanceHistoricalCleanupOffset(lakeCommitOffset)) + .isFalse(); assertThat(replica.getKvSnapshotManager()).isNull(); assertThat(recoveredKvTablet.getFlushedLogOffset()) .isEqualTo(replica.getLogHighWatermark()); @@ -723,8 +902,16 @@ void testRecoversHistoricalOverlayFromLakeCommitOffset() throws Exception { tieredPrimaryKey, tableInfo, row(1, "us", ANOTHER_ORIGINAL_PARTITION, "another")); - } finally { - historicalPartitionManager.close(); + assertHistoricalValueTag( + recoveredKvTablet, + ANOTHER_ORIGINAL_PARTITION, + tieredPrimaryKey, + anotherPartitionAppend.lastOffset()); + assertHistoricalValueTag( + recoveredKvTablet, + ORIGINAL_PARTITION, + deletedPrimaryKey, + deleteAppend.lastOffset()); } } @@ -739,11 +926,7 @@ void testHistoricalLookupThrottledWhenPermitsExhausted() throws Exception { new HistoricalPartitionTaskExecutor(lookupConfiguration(), executor), new TestingHistoricalLakeLookupManager(lookupConfiguration())); - RowType keyType = - DataTypes.ROW( - new DataField("id", DataTypes.INT()), - new DataField("region", DataTypes.STRING())); - byte[] primaryKey = new CompactedKeyEncoder(keyType).encodeKey(row(1, "us")); + byte[] primaryKey = key(1, "us"); LookupDataForBucket lookupData = new LookupDataForBucket( TABLE_BUCKET, Collections.singletonList(primaryKey), ORIGINAL_PARTITION); @@ -776,6 +959,39 @@ void testHistoricalLookupThrottledWhenPermitsExhausted() throws Exception { } } + private HistoricalPartitionManager createHistoricalPartitionManager( + TestingHistoricalLakeLookupManager lakeLookupManager) { + return new HistoricalPartitionManager( + new HistoricalPartitionTaskExecutor(lookupConfiguration()), lakeLookupManager); + } + + private static LogAppendInfo writeBatch( + HistoricalPartitionManager manager, + Replica replica, + String originalPartition, + KvRecordBatch records) + throws Exception { + return manager.processPut( + replica, + new PutKvDataForBucket(TABLE_BUCKET, records, originalPartition), + null, + MergeMode.DEFAULT, + 1); + } + + private static byte[] key(int id, String region) { + return new CompactedKeyEncoder(KEY_TYPE).encodeKey(row(id, region)); + } + + private static Tuple2 upsert( + int id, String region, String partition, String value) { + return Tuple2.of(new Object[] {id, region}, new Object[] {id, region, partition, value}); + } + + private static Tuple2 delete(int id, String region) { + return Tuple2.of(new Object[] {id, region}, null); + } + private LogRecords fetchLog(long fetchOffset) throws Exception { CompletableFuture> future = new CompletableFuture<>(); @@ -937,11 +1153,10 @@ private static void await(CountDownLatch latch) { } @SafeVarargs - private static KvRecordBatch batch( - RowType keyType, RowType rowType, Tuple2... keyAndValues) + private static KvRecordBatch batch(RowType rowType, Tuple2... keyAndValues) throws Exception { List> records = Arrays.asList(keyAndValues); - return genKvRecordBatch(keyType, rowType, records); + return genKvRecordBatch(KEY_TYPE, rowType, records); } private static void assertHistoricalValue( @@ -957,11 +1172,42 @@ private static void assertHistoricalValue( new ValueDecoder( schemaGetter(tableInfo), tableInfo.getTableConfig().getKvFormat(), - KvValueLayout.fromTableConfig(tableInfo.getTableConfig())) + KvValueLayout.TAGGED) .decodeValue(result.value()); assertThatRow(value.row).withSchema(tableInfo.getRowType()).isEqualTo(expectedRow); } + private static void assertHistoricalValueTag( + KvTablet kvTablet, String originalPartition, byte[] primaryKey, long expectedLogOffset) + throws IOException { + byte[] rawValue = historicalRawValue(kvTablet, originalPartition, primaryKey); + assertThat(rawValue).isNotNull(); + assertThat(KvValueLayout.TAGGED.readValueTag(MemorySegment.wrap(rawValue))) + .isEqualTo(expectedLogOffset); + } + + private static byte[] historicalRawValue( + KvTablet kvTablet, String originalPartition, byte[] primaryKey) throws IOException { + return kvTablet.getRocksDBKv() + .get(HistoricalKvKeyEncoder.encode(originalPartition, primaryKey)); + } + + private static void compactHistoricalKv(KvTablet kvTablet) throws Exception { + try (FlushOptions flushOptions = new FlushOptions().setWaitForFlush(true)) { + kvTablet.getRocksDBKv().getDb().flush(flushOptions); + kvTablet.getRocksDBKv().getDb().compactRange(); + } + } + + private static void deleteHistoricalRocksDbKey(KvTablet kvTablet, byte[] primaryKey) { + try { + kvTablet.getRocksDBKv() + .delete(HistoricalKvKeyEncoder.encode(ORIGINAL_PARTITION, primaryKey)); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + private static SchemaGetter schemaGetter(TableInfo tableInfo) { return new TestingSchemaGetter( new SchemaInfo(tableInfo.getSchema(), tableInfo.getSchemaId())); @@ -970,8 +1216,10 @@ private static SchemaGetter schemaGetter(TableInfo tableInfo) { private final class TestingHistoricalLakeLookupManager extends HistoricalLakeLookupManager { private final AtomicInteger lookupCount = new AtomicInteger(); private final AtomicInteger lookupBatchCount = new AtomicInteger(); + private final AtomicInteger requiredSnapshotCount = new AtomicInteger(); private final Map lakeValuesByPartition = new HashMap<>(); private volatile @Nullable Runnable lookupHook; + private volatile @Nullable Runnable requireSnapshotHook; private TestingHistoricalLakeLookupManager(Configuration configuration) { super( @@ -992,6 +1240,20 @@ private void setLookupHook(Runnable lookupHook) { this.lookupHook = lookupHook; } + private void setRequireSnapshotHook(Runnable requireSnapshotHook) { + this.requireSnapshotHook = requireSnapshotHook; + } + + @Override + void requireLakeSnapshot(long tableId, long snapshotId) { + super.requireLakeSnapshot(tableId, snapshotId); + requiredSnapshotCount.incrementAndGet(); + Runnable hook = requireSnapshotHook; + if (hook != null) { + hook.run(); + } + } + @Override List lookup( LookupDataForBucket lookupData,