This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new b38d073146a Fix duplicate upsert rows after restart when a snapshot is
missed with metadataTTL (#19832)
b38d073146a is described below
commit b38d073146aefe1e72f2966608503a453bd6bf4e
Author: Akshay Upadhyay <[email protected]>
AuthorDate: Fri Oct 16 21:43:51 2026 +0530
Fix duplicate upsert rows after restart when a snapshot is missed with
metadataTTL (#19832)
---
.../UpsertSnapshotRestartIntegrationTest.java | 350 +++++++++++++++++++++
.../upsert/BasePartitionUpsertMetadataManager.java | 61 +++-
.../local/upsert/UpsertSnapshotRestartTest.java | 345 ++++++++++++++++++++
3 files changed, 753 insertions(+), 3 deletions(-)
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertSnapshotRestartIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertSnapshotRestartIntegrationTest.java
new file mode 100644
index 00000000000..8ae14427f3c
--- /dev/null
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertSnapshotRestartIntegrationTest.java
@@ -0,0 +1,350 @@
+/**
+ * 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.pinot.integration.tests;
+
+import com.google.common.base.Joiner;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.locks.Lock;
+import javax.annotation.Nullable;
+import org.apache.commons.io.FileUtils;
+import org.apache.helix.model.IdealState;
+import org.apache.pinot.client.ResultSet;
+import org.apache.pinot.client.ResultSetGroup;
+import org.apache.pinot.common.utils.LLCSegmentName;
+import org.apache.pinot.core.data.manager.realtime.RealtimeTableDataManager;
+import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl;
+import
org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager.SnapshotPass;
+import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.UpsertConfig;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.spi.utils.Enablement;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.apache.pinot.util.TestUtils;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+
+/// Restarts the server of a full upsert table with metadata TTL, snapshot and
preload enabled, consuming from Kafka.
+/// Holding the segmentLock of the first segment while the next segment is
committed makes the snapshot pass skip it,
+/// like a segment reload holding the lock does, so its snapshot on disk is
stale or missing while it is out of TTL.
+/// The restarted server must still return exactly one row per primary key,
i.e. the segment must not be added as out
+/// of TTL with the stale snapshot, or with all docs valid when there is no
snapshot.
+public class UpsertSnapshotRestartIntegrationTest extends
BaseClusterIntegrationTestSet {
+ private static final String CSV_HEADER =
"playerId,name,game,score,timestampInEpoch,deleted";
+ private static final int NUM_KEYS = 20;
+ /// Must be smaller than the comparison-column gap between phase 1 (~1k) and
phase 2 (~200k) so seg0 is out of TTL
+ /// after phase 2.
+ private static final double TTL_MS = 100_000d;
+ private static final long PHASE1_BASE_TS = 1_000L;
+ private static final long PHASE2_BASE_TS = 200_000L;
+ private static final long WAIT_MS = 300_000L;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir);
+ startZk();
+ startController();
+ startBroker();
+ startServers(1);
+ startKafkaWithoutTopic();
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws Exception {
+ stopServer();
+ stopBroker();
+ stopController();
+ stopKafka();
+ stopZk();
+ FileUtils.deleteDirectory(_tempDir);
+ }
+
+ @Override
+ protected String getSchemaFileName() {
+ return "upsert_table_test.schema";
+ }
+
+ @Nullable
+ @Override
+ protected String getTimeColumnName() {
+ return "timestampInEpoch";
+ }
+
+ @Override
+ protected String getPartitionColumn() {
+ return "playerId";
+ }
+
+ @Override
+ protected int getNumKafkaPartitions() {
+ return 1;
+ }
+
+ @Override
+ protected int getRealtimeSegmentFlushSize() {
+ return NUM_KEYS;
+ }
+
+ @Override
+ protected void overrideServerConf(PinotConfiguration serverConf) {
+ serverConf.setProperty(Joiner.on(".")
+ .join(Server.INSTANCE_DATA_MANAGER_CONFIG_PREFIX,
Server.Upsert.CONFIG_PREFIX,
+ Server.Upsert.DEFAULT_ENABLE_SNAPSHOT), "true");
+ serverConf.setProperty(Joiner.on(".")
+ .join(Server.INSTANCE_DATA_MANAGER_CONFIG_PREFIX,
Server.Upsert.CONFIG_PREFIX,
+ Server.Upsert.DEFAULT_ENABLE_PRELOAD), "true");
+ }
+
+ @Test
+ public void testRestartWithoutLockContention()
+ throws Exception {
+ runScenario("noContention", false, false, false);
+ }
+
+ @Test
+ public void testRestartWithStaleSnapshot()
+ throws Exception {
+ runScenario("staleSnapshot", true, false, false);
+ }
+
+ @Test
+ public void testRestartWithMissingSnapshot()
+ throws Exception {
+ runScenario("missingSnapshot", true, true, false);
+ }
+
+ @Test
+ public void testReloadAfterRestart()
+ throws Exception {
+ runScenario("reload", true, false, true);
+ }
+
+ private List<String> rows(long baseTs) {
+ List<String> rows = new ArrayList<>();
+ for (int pk = 1; pk <= NUM_KEYS; pk++) {
+ rows.add(pk + ",name" + pk + ",game," + pk + "," + (baseTs + pk) +
",false");
+ }
+ return rows;
+ }
+
+ private RealtimeTableDataManager tdm(String realtimeTable) {
+ return (RealtimeTableDataManager)
_serverStarters.get(0).getServerInstance().getInstanceDataManager()
+ .getTableDataManager(realtimeTable);
+ }
+
+ private List<String> sortedLlcSegmentNames(String realtimeTable) {
+ IdealState idealState =
_helixResourceManager.getTableIdealState(realtimeTable);
+ if (idealState == null) {
+ return List.of();
+ }
+ List<String> names = new ArrayList<>();
+ for (String segmentName : idealState.getRecord().getMapFields().keySet()) {
+ if (LLCSegmentName.isLLCSegment(segmentName)) {
+ names.add(segmentName);
+ }
+ }
+ names.sort(Comparator.comparingInt(name ->
LLCSegmentName.of(name).getSequenceNumber()));
+ return names;
+ }
+
+ private File segmentDir(String realtimeTable, TableConfig tableConfig,
String segmentName) {
+ return tdm(realtimeTable).getSegmentDataDir(segmentName, null,
tableConfig);
+ }
+
+ private static File findSnapshot(File segmentDir) {
+ File v3 = new File(segmentDir, "v3");
+ File snap = new File(v3, V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME);
+ if (snap.exists()) {
+ return snap;
+ }
+ File[] found = new File[1];
+ FileUtils.listFiles(segmentDir, null, true).forEach(f -> {
+ if (f.getName().equals(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME)) {
+ found[0] = f;
+ }
+ });
+ return found[0];
+ }
+
+ private long countStar(String table)
+ throws Exception {
+ return getPinotConnection().execute("SELECT COUNT(*) FROM " +
table).getResultSet(0).getLong(0, 0);
+ }
+
+ private int duplicatePkCount(String table)
+ throws Exception {
+ ResultSet dups = getPinotConnection().execute(
+ "SELECT playerId, COUNT(*) FROM " + table + " GROUP BY playerId HAVING
COUNT(*) > 1 LIMIT 100")
+ .getResultSet(0);
+ return dups.getRowCount();
+ }
+
+ /// Returns the last finished snapshot pass of the only partition. A
snapshot pass runs when a consuming segment
+ /// starts, except the first one, which starts before any record is consumed.
+ @Nullable
+ private SnapshotPass lastSnapshotPass(String realtimeTable) {
+ return
tdm(realtimeTable).getTableUpsertMetadataManager().getOrCreatePartitionManager(0).getLastSnapshotPass();
+ }
+
+ /// Waits for a snapshot pass to finish after the given one, or for the
first one if null.
+ private SnapshotPass waitForNextSnapshotPass(String realtimeTable, @Nullable
SnapshotPass previousPass) {
+ TestUtils.waitForCondition(aVoid -> {
+ SnapshotPass pass = lastSnapshotPass(realtimeTable);
+ return pass != null && pass != previousPass;
+ }, 500L, WAIT_MS, "expected a snapshot pass to finish for " +
realtimeTable);
+ return lastSnapshotPass(realtimeTable);
+ }
+
+ /// Returns whether the validDocIds snapshot of the segment on disk matches
its validDocIds in memory.
+ private boolean isSnapshotUpToDate(String realtimeTable, String segmentName)
{
+ RealtimeTableDataManager tableDataManager = tdm(realtimeTable);
+ SegmentDataManager segmentDataManager =
tableDataManager.acquireSegment(segmentName);
+ assertNotNull(segmentDataManager, "segment " + segmentName + " not found");
+ try {
+ ImmutableSegmentImpl segment = (ImmutableSegmentImpl)
segmentDataManager.getSegment();
+ MutableRoaringBitmap snapshot =
segment.loadDocIdsFromSnapshot(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME);
+ return snapshot != null &&
snapshot.equals(segment.getValidDocIds().getMutableRoaringBitmap());
+ } finally {
+ tableDataManager.releaseSegment(segmentDataManager);
+ }
+ }
+
+ private String summary(String table)
+ throws Exception {
+ return "count(*)=" + countStar(table) + ", pks with >1 rows=" +
duplicatePkCount(table);
+ }
+
+ private void waitForRowCount(String table, long expected)
+ throws Exception {
+ TestUtils.waitForCondition(aVoid -> {
+ try {
+ return countStar(table) == expected;
+ } catch (Exception e) {
+ return false;
+ }
+ }, 500L, WAIT_MS, "expected " + expected + " rows in " + table);
+ }
+
+ private void waitForLlcSegments(String realtimeTable, int minSegments)
+ throws Exception {
+ TestUtils.waitForCondition(aVoid ->
sortedLlcSegmentNames(realtimeTable).size() >= minSegments, 500L, WAIT_MS,
+ "expected >=" + minSegments + " LLC segments in ideal state, have "
+ + sortedLlcSegmentNames(realtimeTable).size());
+ }
+
+ private void runScenario(String name, boolean holdLock, boolean
deleteSnapshotFile, boolean reloadAfterRestart)
+ throws Exception {
+ String table = "upsertSnapshotRestart_" + name;
+ String rtTable = TableNameBuilder.REALTIME.tableNameWithType(table);
+ String topic = table + "_topic";
+ createKafkaTopic(topic);
+ Schema schema = createSchema();
+ schema.setSchemaName(table);
+ addSchema(schema);
+ UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.FULL);
+ upsertConfig.setSnapshot(Enablement.ENABLE);
+ upsertConfig.setEnablePreload(true);
+ upsertConfig.setMetadataTTL(TTL_MS);
+ TableConfig tableConfig = createCSVUpsertTableConfig(table, topic,
getNumKafkaPartitions(),
+ getCSVDecoderProperties(",", CSV_HEADER), upsertConfig, "playerId");
+ addTableConfig(tableConfig);
+
+ pushCsvIntoKafka(rows(PHASE1_BASE_TS), topic, 0);
+ waitForRowCount(table, NUM_KEYS);
+ waitForLlcSegments(rtTable, 2);
+ String seg0Name = sortedLlcSegmentNames(rtTable).get(0);
+ // The first snapshot pass runs when the second consuming segment starts
+ SnapshotPass firstPass = waitForNextSnapshotPass(rtTable, null);
+
+ Lock lock = holdLock ? tdm(rtTable).getSegmentLock(seg0Name) : null;
+ if (lock != null) {
+ lock.lock();
+ }
+ try {
+ pushCsvIntoKafka(rows(PHASE2_BASE_TS), topic, 0);
+ // The next snapshot pass runs when the third consuming segment starts,
i.e. after the rows of phase 2 replaced
+ // all the rows of seg0. Hold the lock until the pass finishes, so that
it skips seg0.
+ waitForNextSnapshotPass(rtTable, firstPass);
+ } finally {
+ if (lock != null) {
+ lock.unlock();
+ }
+ }
+ if (deleteSnapshotFile) {
+ File snapshot = findSnapshot(segmentDir(rtTable, tableConfig, seg0Name));
+ if (snapshot != null) {
+ FileUtils.forceDelete(snapshot);
+ }
+ }
+ // Make sure the restart below exercises a stale or missing snapshot of
seg0 when holding the lock, and an
+ // up-to-date one otherwise
+ assertEquals(isSnapshotUpToDate(rtTable, seg0Name), !holdLock,
+ name + ": unexpected state of the validDocIds snapshot of " + seg0Name
+ " before restart");
+ assertEquals(countStar(table), (long) NUM_KEYS, "pre-restart count");
+
+ _serverStarters.set(0, restartServer(_serverStarters.get(0)));
+ // A segment is routed to only after its state transition, which adds it
to the upsert metadata, so the result is
+ // final as soon as all the segments are queried without exception
+ TestUtils.waitForCondition(aVoid -> {
+ try {
+ ResultSetGroup resultSetGroup = getPinotConnection().execute("SELECT
COUNT(*) FROM " + table);
+ return resultSetGroup.getExceptions().isEmpty() &&
resultSetGroup.getExecutionStats().getNumSegmentsQueried()
+ >= 3;
+ } catch (Exception e) {
+ return false;
+ }
+ }, 1000L, WAIT_MS, "segments back online after restart");
+ String afterRestart = summary(table);
+
+ assertEquals(countStar(table), (long) NUM_KEYS,
+ name + ": expected one row per primary key after restart, got " +
afterRestart);
+ assertEquals(duplicatePkCount(table), 0,
+ name + ": found duplicate primary keys after restart, got " +
afterRestart);
+
+ if (reloadAfterRestart) {
+ String jobId = reloadTableAndValidateResponse(table, TableType.REALTIME,
false);
+ TestUtils.waitForCondition(aVoid -> {
+ try {
+ return isReloadJobCompleted(jobId);
+ } catch (Exception e) {
+ return false;
+ }
+ }, 1000L, WAIT_MS, "reload job finished");
+ String afterReload = summary(table);
+ assertEquals(countStar(table), (long) NUM_KEYS, "reload should keep one
row per PK, got " + afterReload);
+ assertEquals(duplicatePkCount(table), 0, "duplicate PKs after reload,
got " + afterReload);
+ }
+ dropRealtimeTable(table);
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java
index 3a2c00aeab5..894925faa48 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java
@@ -131,6 +131,9 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
// Used to maintain the largestSeenComparisonValue to avoid handling
out-of-ttl segments/records.
// If upsertTTL enabled, we will keep track of largestSeenComparisonValue to
compute expired segments.
protected final AtomicDouble _largestSeenComparisonValue;
+ // The TTL watermark last persisted to disk. A segment out of TTL against it
has an up-to-date snapshot on disk, see
+ // getWatermarkToPersist().
+ private volatile double _persistedWatermark;
// The following variables are always accessed within synchronized block
private boolean _stopped;
@@ -202,6 +205,7 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
_largestSeenComparisonValue = new AtomicDouble(TTL_WATERMARK_NOT_SET);
WatermarkUtils.deleteWatermark(getWatermarkFile());
}
+ _persistedWatermark = _largestSeenComparisonValue.get();
}
@Override
@@ -346,6 +350,14 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
&& maxComparisonValue < _largestSeenComparisonValue.get() -
_metadataTTL;
}
+ /// Returns whether the segment can be handled as out of TTL when it's added
or preloaded, i.e. with its validDocIds
+ /// snapshot taken as is. Besides the watermark in memory, the segment must
also be out of TTL against the persisted
+ /// watermark, which is capped to keep any segment with a stale or missing
snapshot within TTL. The watermark in
+ /// memory alone is not enough, e.g. on server restart it's bumped by the
newer segments loaded earlier.
+ protected boolean canSkipSegmentOutOfTTL(double maxComparisonValue) {
+ return isOutOfMetadataTTL(maxComparisonValue) && maxComparisonValue <
_persistedWatermark - _metadataTTL;
+ }
+
protected boolean isOutOfMetadataTTL(IndexSegment segment) {
if (_metadataTTL > 0) {
double maxComparisonValue = getMaxComparisonValue(segment);
@@ -390,7 +402,7 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
// watermark stays unchanged - the segment's rows are not in the map, so
recording their max would misrepresent
// what we have observed and let the sweep expire other keys against a
phantom watermark.
double maxComparisonValue = isTTLEnabled() ?
getMaxComparisonValue(segment) : TTL_WATERMARK_NOT_SET;
- if (isTTLEnabled() && isOutOfMetadataTTL(maxComparisonValue) &&
skipAddSegmentOutOfTTL(segment)) {
+ if (isTTLEnabled() && canSkipSegmentOutOfTTL(maxComparisonValue) &&
skipAddSegmentOutOfTTL(segment)) {
updateLargestSeenComparisonValue(maxComparisonValue);
return;
}
@@ -474,7 +486,7 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
}
// Bump watermark after rows are added; see doAddSegment.
double maxComparisonValue = isTTLEnabled() ?
getMaxComparisonValue(segment) : TTL_WATERMARK_NOT_SET;
- if (isTTLEnabled() && isOutOfMetadataTTL(maxComparisonValue)
+ if (isTTLEnabled() && canSkipSegmentOutOfTTL(maxComparisonValue)
&& skipPreloadSegmentOutOfTTL(segment, validDocIds)) {
updateLargestSeenComparisonValue(maxComparisonValue);
return;
@@ -1167,7 +1179,9 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
// updated validDocIds bitmaps. If the TTL watermark is persisted first,
segments out of TTL may get loaded with
// stale bitmaps or even no bitmap snapshots to use.
if (isTTLEnabled()) {
- WatermarkUtils.persistWatermark(_largestSeenComparisonValue.get(),
getWatermarkFile());
+ double watermark = getWatermarkToPersist();
+ WatermarkUtils.persistWatermark(watermark, getWatermarkFile());
+ _persistedWatermark = watermark;
}
updateSnapshotMetrics(numImmutableSegments, numPrimaryKeysInSnapshot,
numQueryableDocIdsInSnapshot,
numTrackedSegments, numConsumingSegments, numUnchangedSegments);
@@ -1177,6 +1191,38 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
numConsumingSegments, System.currentTimeMillis() - startTimeMs);
}
+ /// Returns the TTL watermark to persist after taking snapshots. A segment
still tracked as updated after the snapshot
+ /// pass, e.g. skipped on segmentLock contention or added while the pass
ran, has a stale snapshot or no snapshot on
+ /// disk. If such segment were out of TTL after a restart, it would be added
with the stale snapshot, or with all
+ /// docs treated as valid when there is no snapshot, which shows up as
duplicate rows on this replica only. So cap
+ /// the persisted watermark to keep such segments within TTL, and get them
added with the regular flow after restart.
+ /// The watermark in memory is not changed, and primary keys are still
expired against it, so a segment that keeps
+ /// missing its snapshot doesn't keep the expired primary keys in memory.
+ private double getWatermarkToPersist() {
+ double watermark = _largestSeenComparisonValue.get();
+ if (_metadataTTL <= 0) {
+ return watermark;
+ }
+ double cappedWatermark = watermark;
+ // Unlike taking the snapshots above, reading the max comparison value
doesn't need the segmentLock or the check
+ // against _trackedSegments: it comes from the segment metadata, which can
still be read after the segment is
+ // destroyed, and a segment removed in the meantime can only cap the
watermark lower than needed.
+ for (ImmutableSegmentImpl segment : _updatedSegmentsSinceLastSnapshot) {
+ double maxComparisonValue = getMaxComparisonValue(segment);
+ double maxWatermark = maxComparisonValue + _metadataTTL;
+ // Make sure the segment is not out of TTL as checked by
isOutOfMetadataTTL(), regardless of rounding
+ while (maxComparisonValue < maxWatermark - _metadataTTL) {
+ maxWatermark = Math.nextDown(maxWatermark);
+ }
+ cappedWatermark = Math.min(cappedWatermark, maxWatermark);
+ }
+ if (cappedWatermark < watermark) {
+ _logger.warn("Capped TTL watermark to persist from: {} to: {} because
some segments missed taking snapshot",
+ watermark, cappedWatermark);
+ }
+ return cappedWatermark;
+ }
+
private void updateSnapshotMetrics(int numImmutableSegments, long
numPrimaryKeysInSnapshot,
long numQueryableDocIdsInSnapshot, int numTrackedSegments, int
numConsumingSegments, int numUnchangedSegments) {
_serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId,
@@ -1214,14 +1260,23 @@ public abstract class
BasePartitionUpsertMetadataManager implements PartitionUps
return new File(_tableIndexDir, V1Constants.TTL_WATERMARK_TABLE_PARTITION
+ _partitionId);
}
+ /// Returns the TTL watermark last persisted to disk, as computed by
getWatermarkToPersist(). The validDocIds
+ /// snapshots on disk are up to date for all the segments out of TTL against
it, while they can be stale for the
+ /// segments out of TTL only against the watermark in memory.
+ protected double getPersistedWatermark() {
+ return _persistedWatermark;
+ }
+
@VisibleForTesting
double getWatermark() {
return _largestSeenComparisonValue.get();
}
+ /// Sets the watermark as if it was loaded from the persisted watermark file.
@VisibleForTesting
void setWatermark(double watermark) {
_largestSeenComparisonValue.set(watermark);
+ _persistedWatermark = watermark;
}
@Override
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertSnapshotRestartTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertSnapshotRestartTest.java
new file mode 100644
index 00000000000..dc1723d5f58
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertSnapshotRestartTest.java
@@ -0,0 +1,345 @@
+/**
+ * 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.pinot.segment.local.upsert;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.segment.local.data.manager.TableDataManager;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.local.utils.WatermarkUtils;
+import org.apache.pinot.segment.spi.MutableSegment;
+import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import
org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.PrimaryKey;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+
+/// Restarts a full upsert partition with metadata TTL and snapshots enabled,
using real segments, snapshot files and
+/// TTL watermark file. Only the [TableDataManager] is mocked, to make
segmentLock contended for chosen segments.
+///
+/// segA (ts=100) has pk 1, 2, 3 and the newer segB (ts=290) has pk 1, 2, so
the expected valid docs are segA={2} and
+/// segB={0,1}. With TTL=30, segA is out of TTL against the largest seen
comparison value 290. When the snapshot pass
+/// misses segA, its snapshot on disk is stale or missing, and segA must not
be treated as out of TTL after a restart,
+/// otherwise it is added with the stale snapshot or with all docs valid, i.e.
with duplicate primary keys.
+public class UpsertSnapshotRestartTest {
+ private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(),
"UpsertSnapshotRestartTest");
+ private static final File SEGMENTS_DIR = new File(TEMP_DIR, "segments");
+ private static final String RAW_TABLE_NAME = "testTable";
+ private static final String REALTIME_TABLE_NAME = "testTable_REALTIME";
+ private static final Schema SCHEMA = new
Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME)
+ .addSingleValueDimension("pk", FieldSpec.DataType.INT).addMetric("ts",
FieldSpec.DataType.INT)
+ .setPrimaryKeyColumns(List.of("pk")).build();
+ private static final TableConfig TABLE_CONFIG =
+ new
TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).build();
+ private static final double METADATA_TTL = 30;
+
+ private final Set<String> _contendedSegments = ConcurrentHashMap.newKeySet();
+ private TableDataManager _tableDataManager;
+
+ @BeforeMethod
+ public void setUp()
+ throws Exception {
+ FileUtils.deleteQuietly(TEMP_DIR);
+ FileUtils.forceMkdir(TEMP_DIR);
+ ServerMetrics.register(mock(ServerMetrics.class));
+ _contendedSegments.clear();
+ _tableDataManager = mock(TableDataManager.class);
+ when(_tableDataManager.getTableDataDir()).thenReturn(TEMP_DIR);
+ ConcurrentHashMap<String, ReentrantLock> segmentLocks = new
ConcurrentHashMap<>();
+ when(_tableDataManager.getSegmentLock(anyString())).thenAnswer(
+ invocation -> segmentLocks.computeIfAbsent(invocation.getArgument(0),
this::newSegmentLock));
+ buildSegment("segA", new int[]{1, 2, 3}, 100);
+ buildSegment("segB", new int[]{1, 2}, 290);
+ }
+
+ @AfterMethod
+ public void tearDown() {
+ FileUtils.deleteQuietly(TEMP_DIR);
+ }
+
+ @DataProvider
+ public static Object[][] restartModes() {
+ // Whether to preload the segments with snapshot on restart, and whether
to add segB before segA
+ return new Object[][]{{false, false}, {false, true}, {true, false}, {true,
true}};
+ }
+
+ @Test(dataProvider = "restartModes")
+ public void testRestartWithoutMissedSnapshot(boolean preload, boolean
addSegBFirst)
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ ImmutableSegmentImpl segB = addSegment(manager, "segB");
+ manager.takeSnapshot();
+ assertEquals(persistedWatermark(manager), 290.0);
+ close(manager, segA, segB);
+
+ assertEquals(restart(preload, addSegBFirst), "{2} | {0,1}");
+ }
+
+ /// segA never got a snapshot, e.g. its segmentLock was held by a reload
whenever the snapshot was taken.
+ @Test(dataProvider = "restartModes")
+ public void testRestartWithMissingSnapshot(boolean preload, boolean
addSegBFirst)
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ ImmutableSegmentImpl segB = addSegment(manager, "segB");
+ _contendedSegments.add("segA");
+ manager.takeSnapshot();
+ assertFalse(hasSnapshot("segA"));
+ assertTrue(hasSnapshot("segB"));
+ // Capped to keep segA within TTL: 100 >= 130 - 30
+ assertEquals(persistedWatermark(manager), 130.0);
+ close(manager, segA, segB);
+
+ assertEquals(restart(preload, addSegBFirst), "{2} | {0,1}");
+ }
+
+ /// segA got a snapshot while all its docs were valid, then missed the
snapshot after segB invalidated 2 of them.
+ @Test(dataProvider = "restartModes")
+ public void testRestartWithStaleSnapshot(boolean preload, boolean
addSegBFirst)
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ manager.takeSnapshot();
+ assertEquals(loadSnapshot("segA"), "{0,1,2}");
+ ImmutableSegmentImpl segB = addSegment(manager, "segB");
+ _contendedSegments.add("segA");
+ manager.takeSnapshot();
+ assertEquals(loadSnapshot("segA"), "{0,1,2}");
+ // segB is not snapshotted either as the loop for segments without
snapshot is skipped
+ assertFalse(hasSnapshot("segB"));
+ assertEquals(persistedWatermark(manager), 130.0);
+ close(manager, segA, segB);
+
+ assertEquals(restart(preload, addSegBFirst), "{2} | {0,1}");
+ }
+
+ @Test
+ public void testWatermarkAdvancesOnceMissedSnapshotIsTaken()
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ ImmutableSegmentImpl segB = addSegment(manager, "segB");
+ _contendedSegments.add("segA");
+ manager.takeSnapshot();
+ assertEquals(persistedWatermark(manager), 130.0);
+ // The watermark in memory is not capped, so the expired primary keys are
still removed as before
+ assertEquals(manager.getWatermark(), 290.0);
+
+ _contendedSegments.clear();
+ manager.takeSnapshot();
+ assertEquals(loadSnapshot("segA"), "{2}");
+ assertEquals(persistedWatermark(manager), 290.0);
+ close(manager, segA, segB);
+ }
+
+ /// The just committed segment can miss its snapshot because its segmentLock
is still held by the thread replacing the
+ /// consuming segment. It is within TTL, so it should not hold back the
persisted watermark.
+ @Test
+ public void testMissedSnapshotWithinTtlDoesNotCapWatermark()
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ ImmutableSegmentImpl segB = addSegment(manager, "segB");
+ _contendedSegments.add("segB");
+ manager.takeSnapshot();
+ assertTrue(hasSnapshot("segA"));
+ assertFalse(hasSnapshot("segB"));
+ assertEquals(persistedWatermark(manager), 290.0);
+ close(manager, segA, segB);
+
+ for (Object[] restartMode : restartModes()) {
+ assertEquals(restart((boolean) restartMode[0], (boolean)
restartMode[1]), "{2} | {0,1}");
+ }
+ }
+
+ /// Stopping the metadata manager does not take a snapshot. Rows consumed
into the consuming segment are not
+ /// committed, and are consumed again after a restart, so the snapshots must
not include the doc invalidated by them.
+ /// Otherwise pk 1 would be invisible after a restart until the consuming
segment catches up again.
+ @Test
+ public void testStopDoesNotSnapshotUncommittedUpdates()
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ ImmutableSegmentImpl segA = addSegment(manager, "segA");
+ manager.takeSnapshot();
+ MutableSegment consumingSegment = mock(MutableSegment.class);
+ when(consumingSegment.getSegmentName()).thenReturn("segC");
+ when(consumingSegment.getValidDocIds()).thenReturn(new
ThreadSafeMutableRoaringBitmap());
+ manager.addRecord(consumingSegment, new RecordInfo(new PrimaryKey(new
Object[]{1}), 0, 200, false));
+ assertEquals(segA.getValidDocIds().getMutableRoaringBitmap().toString(),
"{1,2}");
+ manager.stop();
+ manager.close();
+ segA.destroy();
+ assertEquals(loadSnapshot("segA"), "{0,1,2}");
+
+ manager = newManager();
+ segA = load("segA");
+ manager.preloadSegment(segA);
+ assertEquals(segA.getValidDocIds().getMutableRoaringBitmap().toString(),
"{0,1,2}");
+ manager.stop();
+ manager.close();
+ segA.destroy();
+ }
+
+ private ReentrantLock newSegmentLock(String segmentName) {
+ return new ReentrantLock() {
+ @Override
+ public boolean tryLock() {
+ return !_contendedSegments.contains(segmentName) && super.tryLock();
+ }
+ };
+ }
+
+ private ConcurrentMapPartitionUpsertMetadataManager newManager() {
+ UpsertContext context = new
UpsertContext.Builder().setTableConfig(TABLE_CONFIG).setSchema(SCHEMA)
+
.setTableDataManager(_tableDataManager).setPrimaryKeyColumns(List.of("pk"))
+
.setComparisonColumns(List.of("ts")).setEnableSnapshot(true).setMetadataTTL(METADATA_TTL)
+ .setTableIndexDir(TEMP_DIR).build();
+ ConcurrentMapPartitionUpsertMetadataManager manager =
+ new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME,
0, context);
+ manager._gotFirstConsumingSegment = true;
+ return manager;
+ }
+
+ private static void buildSegment(String segmentName, int[] primaryKeys, int
timestamp)
+ throws Exception {
+ List<GenericRow> rows = new ArrayList<>();
+ for (int primaryKey : primaryKeys) {
+ GenericRow row = new GenericRow();
+ row.putValue("pk", primaryKey);
+ row.putValue("ts", timestamp);
+ rows.add(row);
+ }
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG,
SCHEMA);
+ config.setOutDir(SEGMENTS_DIR.getAbsolutePath());
+ config.setSegmentName(segmentName);
+ config.setTableName(RAW_TABLE_NAME);
+ SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
+ driver.init(config, new GenericRowRecordReader(rows));
+ driver.build();
+ }
+
+ private static ImmutableSegmentImpl load(String segmentName)
+ throws Exception {
+ return (ImmutableSegmentImpl) ImmutableSegmentLoader.load(new
File(SEGMENTS_DIR, segmentName), ReadMode.mmap);
+ }
+
+ private static ImmutableSegmentImpl
addSegment(ConcurrentMapPartitionUpsertMetadataManager manager,
+ String segmentName)
+ throws Exception {
+ ImmutableSegmentImpl segment = load(segmentName);
+ manager.addSegment(segment);
+ return segment;
+ }
+
+ private static boolean hasSnapshot(String segmentName)
+ throws Exception {
+ ImmutableSegmentImpl segment = load(segmentName);
+ try {
+ return
segment.loadDocIdsFromSnapshot(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME) !=
null;
+ } finally {
+ segment.destroy();
+ }
+ }
+
+ private static String loadSnapshot(String segmentName)
+ throws Exception {
+ ImmutableSegmentImpl segment = load(segmentName);
+ try {
+ return
String.valueOf(segment.loadDocIdsFromSnapshot(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME));
+ } finally {
+ segment.destroy();
+ }
+ }
+
+ private static double
persistedWatermark(ConcurrentMapPartitionUpsertMetadataManager manager) {
+ double watermark =
WatermarkUtils.loadWatermark(manager.getWatermarkFile(),
BasePartitionUpsertMetadataManager
+ .TTL_WATERMARK_NOT_SET);
+ assertEquals(manager.getPersistedWatermark(), watermark, "persisted
watermark in memory and on disk differ");
+ return watermark;
+ }
+
+ private static void close(ConcurrentMapPartitionUpsertMetadataManager
manager, ImmutableSegmentImpl... segments)
+ throws Exception {
+ manager.stop();
+ manager.close();
+ for (ImmutableSegmentImpl segment : segments) {
+ segment.destroy();
+ }
+ }
+
+ /// Loads the segments into a new metadata manager as a server restart does:
preloads the segments with snapshot
+ /// first if preload is enabled, then adds the remaining segments. Returns
the valid docs as "segA | segB", where
+ /// "ALL" means no validDocIds, i.e. all docs are returned by queries.
+ private String restart(boolean preload, boolean addSegBFirst)
+ throws Exception {
+ ConcurrentMapPartitionUpsertMetadataManager manager = newManager();
+ // The persisted watermark is loaded from the watermark file on restart
+ persistedWatermark(manager);
+ ImmutableSegmentImpl segA = load("segA");
+ ImmutableSegmentImpl segB = load("segB");
+ List<ImmutableSegmentImpl> segments = addSegBFirst ? List.of(segB, segA) :
List.of(segA, segB);
+ List<ImmutableSegmentImpl> segmentsToAdd = new ArrayList<>();
+ for (ImmutableSegmentImpl segment : segments) {
+ if (preload &&
segment.loadDocIdsFromSnapshot(V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME) !=
null) {
+ manager.preloadSegment(segment);
+ } else {
+ segmentsToAdd.add(segment);
+ }
+ }
+ for (ImmutableSegmentImpl segment : segmentsToAdd) {
+ manager.addSegment(segment);
+ }
+ String validDocIds = validDocIds(segA) + " | " + validDocIds(segB);
+ close(manager, segA, segB);
+ return validDocIds;
+ }
+
+ private static String validDocIds(ImmutableSegmentImpl segment) {
+ ThreadSafeMutableRoaringBitmap validDocIds = segment.getValidDocIds();
+ return validDocIds != null ?
validDocIds.getMutableRoaringBitmap().toString() : "ALL";
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
