This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 4d93fce3399 Fix SIGSEGV when upsert reads a concurrently destroyed
segment (#19596)
4d93fce3399 is described below
commit 4d93fce3399c7796016089cdae5bd7196b7982fe
Author: Chaitanya Deepthi <[email protected]>
AuthorDate: Wed Oct 7 12:14:18 2026 -0700
Fix SIGSEGV when upsert reads a concurrently destroyed segment (#19596)
---
.../immutable/ImmutableSegmentImpl.java | 39 ++++++++++
.../upsert/BasePartitionUpsertMetadataManager.java | 15 ++++
...oncurrentMapPartitionUpsertMetadataManager.java | 73 ++++++++++++------
...nUpsertMetadataManagerForConsistentDeletes.java | 81 ++++++++++++-------
.../immutable/ImmutableSegmentImplTest.java | 90 +++++++++++++++++++++-
...ertMetadataManagerForConsistentDeletesTest.java | 75 ++++++++++++++++++
...rrentMapPartitionUpsertMetadataManagerTest.java | 82 ++++++++++++++++++++
7 files changed, 400 insertions(+), 55 deletions(-)
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
index 0cef98c9c97..65c06d5af59 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.segment.local.indexsegment.immutable;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
import java.io.File;
@@ -30,6 +31,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.segment.local.dedup.PartitionDedupMetadataManager;
@@ -98,6 +100,11 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
private ThreadSafeMutableRoaringBitmap _validDocIds;
private ThreadSafeMutableRoaringBitmap _queryableDocIds;
private volatile boolean _hasDeletedDocIds;
+ // Readers hold the read lock while reading columns through a cached
reference; destroy() takes the write lock to
+ // set _destroyed before closing the indexes
+ private final ReentrantReadWriteLock _destroyLock = new
ReentrantReadWriteLock();
+ // Monotonic: destroy() is the only writer and never clears it
+ private volatile boolean _destroyed;
public ImmutableSegmentImpl(
SegmentDirectory segmentDirectory,
@@ -374,6 +381,13 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
if (_partitionUpsertMetadataManager != null) {
_partitionUpsertMetadataManager.untrackSegmentForUpsertView(this);
}
+ // Wait for in-flight column reads (see tryAcquireReadLock) and make later
ones skip this segment
+ _destroyLock.writeLock().lock();
+ try {
+ _destroyed = true;
+ } finally {
+ _destroyLock.writeLock().unlock();
+ }
// StarTreeIndexContainer refers to other column index containers, so
close it firstly.
if (_starTreeIndexContainer != null) {
try {
@@ -404,6 +418,31 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
}
}
+ /// True once [#destroy()] has set the flag: the index buffers may be closed
and must not be read.
+ @VisibleForTesting
+ boolean isDestroyed() {
+ return _destroyed;
+ }
+
+ /// Blocks [#destroy()] while the caller reads this segment's columns
through a cached reference. Returns false,
+ /// without holding the lock, once the segment is destroyed. A true return
must be paired with [#releaseReadLock()].
+ public boolean tryAcquireReadLock() {
+ // Fast path: the flag is monotonic, so a true reading is definitive. A
false one is re-checked under the lock.
+ if (_destroyed) {
+ return false;
+ }
+ _destroyLock.readLock().lock();
+ if (_destroyed) {
+ _destroyLock.readLock().unlock();
+ return false;
+ }
+ return true;
+ }
+
+ public void releaseReadLock() {
+ _destroyLock.readLock().unlock();
+ }
+
@Nullable
@Override
public DataSource getDataSourceNullable(String column) {
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 f88a71434e4..f077c59aac7 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
@@ -1398,6 +1398,21 @@ public abstract class BasePartitionUpsertMetadataManager
implements PartitionUps
}
}
+ /// Blocks destroy() of an immutable segment while the caller reads its
columns through a cached reference. Returns
+ /// false once the segment is destroyed. Consuming segments always succeed:
their consumer keeps them alive
+ /// while the table runs, though not during shutdown, which destroys
segments without draining operations.
+ /// Liveness is keyed on the segment object on purpose: the TableDataManager
registers a replaced segment only after
+ /// replaceSegment(), so a name based lookup would report the live new
segment as gone for the whole replace.
+ protected static boolean tryAcquireSegmentReadLock(IndexSegment segment) {
+ return !(segment instanceof ImmutableSegmentImpl immutableSegment) ||
immutableSegment.tryAcquireReadLock();
+ }
+
+ protected static void releaseSegmentReadLock(IndexSegment segment) {
+ if (segment instanceof ImmutableSegmentImpl immutableSegment) {
+ immutableSegment.releaseReadLock();
+ }
+ }
+
@Override
public void trackNewlyAddedSegment(String segmentName) {
if (_newSegmentTrackingTimeMs > 0) {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java
index 5860901475b..cf6f11eef50 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java
@@ -240,28 +240,40 @@ public class ConcurrentMapPartitionUpsertMetadataManager
extends BasePartitionUp
}
// Revert to previous segment location
IndexSegment prevSegment = prevLocation.getSegment();
- ThreadSafeMutableRoaringBitmap prevValidDocIds =
prevSegment.getValidDocIds();
- if (prevValidDocIds != null) {
- try (UpsertUtils.RecordInfoReader recordInfoReader = new
UpsertUtils.RecordInfoReader(prevSegment,
- _primaryKeyColumns, _comparisonColumns,
_deleteRecordColumn)) {
- int prevDocId = prevLocation.getDocId();
- RecordInfo recordInfo =
recordInfoReader.getRecordInfo(prevDocId);
- replaceDocId(prevSegment, prevValidDocIds,
prevSegment.getQueryableDocIds(), segment, docId,
- prevDocId, recordInfo);
- return prevLocation;
- } catch (Exception e) {
- _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to revert to previous segment: {}, "
- + "removing key", segment.getSegmentName(),
prevSegment.getSegmentName(), e);
+ // Read lock: prevSegment cannot be destroyed while its columns
are read
+ if (!tryAcquireSegmentReadLock(prevSegment)) {
+ _logger.debug("Previous segment: {} is destroyed while
reconciling metadata for segment: {}, "
+ + "dropping primary key instead of reverting to it",
prevSegment.getSegmentName(),
+ segment.getSegmentName());
+ return null;
+ }
+ try {
+ ThreadSafeMutableRoaringBitmap prevValidDocIds =
prevSegment.getValidDocIds();
+ if (prevValidDocIds != null) {
+ try (UpsertUtils.RecordInfoReader recordInfoReader = new
UpsertUtils.RecordInfoReader(prevSegment,
+ _primaryKeyColumns, _comparisonColumns,
_deleteRecordColumn)) {
+ int prevDocId = prevLocation.getDocId();
+ RecordInfo recordInfo =
recordInfoReader.getRecordInfo(prevDocId);
+ replaceDocId(prevSegment, prevValidDocIds,
prevSegment.getQueryableDocIds(), segment, docId,
+ prevDocId, recordInfo);
+ return prevLocation;
+ } catch (Exception e) {
+ _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to revert to previous segment: "
+ + "{}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName(), e);
+ _serverMetrics.addMeteredTableValue(_tableNameWithType,
+ ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
+ return null;
+ }
+ } else {
+ // Should not happen
+ _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to find valid doc ids in previous "
+ + "segment: {}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName());
_serverMetrics.addMeteredTableValue(_tableNameWithType,
ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
return null;
}
- } else {
- // Should not happen
- _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to find valid doc ids in previous "
- + "segment: {}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName());
- _serverMetrics.addMeteredTableValue(_tableNameWithType,
ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
- return null;
+ } finally {
+ releaseSegmentReadLock(prevSegment);
}
} else if (recordLocation.getSegment() instanceof
ImmutableSegmentImpl) {
// The consuming segment's key is in a different immutable
segment
@@ -428,15 +440,26 @@ public class ConcurrentMapPartitionUpsertMetadataManager
extends BasePartitionUp
if (!recordInfo.isDeleteRecord()
&&
recordInfo.getComparisonValue().compareTo(recordLocation.getComparisonValue())
>= 0) {
IndexSegment currentSegment = recordLocation.getSegment();
- ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
int currentDocId = recordLocation.getDocId();
+ // queryableDocIds is heap backed and outlives destroy(), so a
deleted predecessor needs no lock at all.
+ ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
if (currentQueryableDocIds == null ||
currentQueryableDocIds.contains(currentDocId)) {
- try {
- _reusePreviousRow.init(currentSegment, currentDocId);
- _partialUpsertHandler.merge(_reusePreviousRow, record,
_reuseMergeResultHolder);
- } finally {
- _reuseMergeResultHolder.clear();
- _reusePreviousRow.clear();
+ // Read lock: currentSegment cannot be destroyed while LazyRow
reads its columns. A consuming segment
+ // needs no lock while the table runs: it is destroyed only
after replaceSegment()/removeSegment() has
+ // moved or dropped every location pointing at it, and those run
under the same per-key compute as this
+ // read. Shutdown is the exception, since removeSegment() skips
that cleanup once the manager is stopped.
+ if (tryAcquireSegmentReadLock(currentSegment)) {
+ try {
+ _reusePreviousRow.init(currentSegment, currentDocId);
+ _partialUpsertHandler.merge(_reusePreviousRow, record,
_reuseMergeResultHolder);
+ } finally {
+ _reuseMergeResultHolder.clear();
+ _reusePreviousRow.clear();
+ releaseSegmentReadLock(currentSegment);
+ }
+ } else {
+ _logger.debug("Current segment: {} is destroyed, storing
record without merging the previous row",
+ currentSegment.getSegmentName());
}
}
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java
index d64ee9e0b9f..afda18c35b5 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java
@@ -365,33 +365,45 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes
}
// Revert to previous segment location
IndexSegment prevSegment = prevLocation.getSegment();
- ThreadSafeMutableRoaringBitmap prevValidDocIds =
prevSegment.getValidDocIds();
- if (prevValidDocIds != null) {
- try (UpsertUtils.RecordInfoReader recordInfoReader = new
UpsertUtils.RecordInfoReader(prevSegment,
- _primaryKeyColumns, _comparisonColumns,
_deleteRecordColumn)) {
- int prevDocId = prevLocation.getDocId();
- RecordInfo recordInfo =
recordInfoReader.getRecordInfo(prevDocId);
- replaceDocId(prevSegment, prevValidDocIds,
prevSegment.getQueryableDocIds(), segment, docId,
- prevDocId, recordInfo);
- if (!uniquePrimaryKeys.add(pk)) {
- return prevLocation;
+ // Read lock: prevSegment cannot be destroyed while its columns
are read
+ if (!tryAcquireSegmentReadLock(prevSegment)) {
+ _logger.debug("Previous segment: {} is destroyed while
reconciling metadata for segment: {}, "
+ + "dropping primary key instead of reverting to it",
prevSegment.getSegmentName(),
+ segment.getSegmentName());
+ return null;
+ }
+ try {
+ ThreadSafeMutableRoaringBitmap prevValidDocIds =
prevSegment.getValidDocIds();
+ if (prevValidDocIds != null) {
+ try (UpsertUtils.RecordInfoReader recordInfoReader = new
UpsertUtils.RecordInfoReader(prevSegment,
+ _primaryKeyColumns, _comparisonColumns,
_deleteRecordColumn)) {
+ int prevDocId = prevLocation.getDocId();
+ RecordInfo recordInfo =
recordInfoReader.getRecordInfo(prevDocId);
+ replaceDocId(prevSegment, prevValidDocIds,
prevSegment.getQueryableDocIds(), segment, docId,
+ prevDocId, recordInfo);
+ if (!uniquePrimaryKeys.add(pk)) {
+ return prevLocation;
+ }
+ return new RecordLocation(prevLocation.getSegment(),
prevLocation.getDocId(),
+ prevLocation.getComparisonValue(),
+
RecordLocation.decrementSegmentCount(prevLocation.getDistinctSegmentCount()));
+ } catch (Exception e) {
+ _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to revert to previous segment: "
+ + "{}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName(), e);
+ _serverMetrics.addMeteredTableValue(_tableNameWithType,
+ ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
+ return null;
}
- return new RecordLocation(prevLocation.getSegment(),
prevLocation.getDocId(),
- prevLocation.getComparisonValue(),
-
RecordLocation.decrementSegmentCount(prevLocation.getDistinctSegmentCount()));
- } catch (Exception e) {
- _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to revert to previous segment: {}, "
- + "removing key", segment.getSegmentName(),
prevSegment.getSegmentName(), e);
+ } else {
+ // Should not happen
+ _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to find valid doc ids in previous "
+ + "segment: {}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName());
_serverMetrics.addMeteredTableValue(_tableNameWithType,
ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
return null;
}
- } else {
- // Should not happen
- _logger.error("UPSERT_METADATA_REVERT_FAILED: segment={}.
Failed to find valid doc ids in previous "
- + "segment: {}, removing key", segment.getSegmentName(),
prevSegment.getSegmentName());
- _serverMetrics.addMeteredTableValue(_tableNameWithType,
ServerMeter.UPSERT_METADATA_REVERT_FAILURES, 1);
- return null;
+ } finally {
+ releaseSegmentReadLock(prevSegment);
}
} else if (recordLocation.getSegment() instanceof
ImmutableSegmentImpl) {
// The consuming segment's key is in a different immutable
segment
@@ -558,15 +570,26 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes
if (!recordInfo.isDeleteRecord()
&&
recordInfo.getComparisonValue().compareTo(recordLocation.getComparisonValue())
>= 0) {
IndexSegment currentSegment = recordLocation.getSegment();
- ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
int currentDocId = recordLocation.getDocId();
+ // queryableDocIds is heap backed and outlives destroy(), so a
deleted predecessor needs no lock at all.
+ ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
if (currentQueryableDocIds == null ||
currentQueryableDocIds.contains(currentDocId)) {
- try {
- _reusePreviousRow.init(currentSegment, currentDocId);
- _partialUpsertHandler.merge(_reusePreviousRow, record,
_reuseMergeResultHolder);
- } finally {
- _reuseMergeResultHolder.clear();
- _reusePreviousRow.clear();
+ // Read lock: currentSegment cannot be destroyed while LazyRow
reads its columns. A consuming segment
+ // needs no lock while the table runs: it is destroyed only
after replaceSegment()/removeSegment() has
+ // moved or dropped every location pointing at it, and those run
under the same per-key compute as this
+ // read. Shutdown is the exception, since removeSegment() skips
that cleanup once the manager is stopped.
+ if (tryAcquireSegmentReadLock(currentSegment)) {
+ try {
+ _reusePreviousRow.init(currentSegment, currentDocId);
+ _partialUpsertHandler.merge(_reusePreviousRow, record,
_reuseMergeResultHolder);
+ } finally {
+ _reuseMergeResultHolder.clear();
+ _reusePreviousRow.clear();
+ releaseSegmentReadLock(currentSegment);
+ }
+ } else {
+ _logger.debug("Current segment: {} is destroyed, storing
record without merging the previous row",
+ currentSegment.getSegmentName());
}
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java
index db1ce1aafcb..0196ad963a7 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java
@@ -20,15 +20,26 @@ package
org.apache.pinot.segment.local.indexsegment.immutable;
import java.util.Map;
import java.util.TreeMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import org.apache.pinot.segment.spi.index.column.ColumnIndexContainer;
import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
import org.apache.pinot.segment.spi.store.SegmentDirectory;
import org.testng.annotations.Test;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
/// Tests for the post-registration lifecycle hook of [ImmutableSegmentImpl].
@@ -79,11 +90,88 @@ public class ImmutableSegmentImplTest {
verify(segmentDirectory, times(1)).onSegmentAdded();
}
+ @Test
+ public void testDestroyMarksSegmentDestroyedBeforeClosingIndexes()
+ throws Exception {
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ ColumnIndexContainer indexContainer = mock(ColumnIndexContainer.class);
+ ImmutableSegmentImpl segment = createSegment(segmentDirectory,
Map.of("col", indexContainer));
+ // Readers that see the flag skip the segment, so it must be set before
any index is closed
+ doAnswer(invocation -> {
+ assertTrue(segment.isDestroyed());
+ assertFalse(segment.tryAcquireReadLock());
+ return null;
+ }).when(indexContainer).close();
+
+ segment.destroy();
+
+ verify(indexContainer).close();
+ verify(segmentDirectory).close();
+ }
+
+ @Test
+ public void testDestroyWaitsForReadLock()
+ throws Exception {
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ ColumnIndexContainer indexContainer = mock(ColumnIndexContainer.class);
+ ImmutableSegmentImpl segment = createSegment(segmentDirectory,
Map.of("col", indexContainer));
+ assertTrue(segment.tryAcquireReadLock());
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<?> destroyer = executor.submit(segment::destroy);
+ try {
+ // A reader holds the read lock: destroy must not close anything
+ assertThrows(TimeoutException.class, () -> destroyer.get(500,
TimeUnit.MILLISECONDS));
+ assertFalse(segment.isDestroyed());
+ verify(indexContainer, never()).close();
+ } finally {
+ // destroy() blocks in a non-interruptible write lock, so
shutdownNow() cannot free it: a failed assertion
+ // above would strand the executor thread without this release.
+ segment.releaseReadLock();
+ }
+ destroyer.get(10, TimeUnit.SECONDS);
+ assertTrue(segment.isDestroyed());
+ assertFalse(segment.tryAcquireReadLock());
+ verify(indexContainer).close();
+ verify(segmentDirectory).close();
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testDestroyDoesNotWaitOnAnotherSegmentsReader()
+ throws Exception {
+ // The guard is per segment, so a reader pinning one segment must not hold
up destroy of an unrelated one. Under
+ // a lock shared across the partition this times out.
+ ImmutableSegmentImpl reading = createSegment(mock(SegmentDirectory.class));
+ SegmentDirectory destroyedDirectory = mock(SegmentDirectory.class);
+ ImmutableSegmentImpl destroyed = createSegment(destroyedDirectory);
+ assertTrue(reading.tryAcquireReadLock());
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ executor.submit(destroyed::destroy).get(10, TimeUnit.SECONDS);
+ assertTrue(destroyed.isDestroyed());
+ verify(destroyedDirectory).close();
+ assertFalse(reading.isDestroyed());
+ } finally {
+ reading.releaseReadLock();
+ executor.shutdownNow();
+ }
+ }
+
private static ImmutableSegmentImpl createSegment(SegmentDirectory
segmentDirectory) {
+ return createSegment(segmentDirectory, Map.of());
+ }
+
+ private static ImmutableSegmentImpl createSegment(SegmentDirectory
segmentDirectory,
+ Map<String, ColumnIndexContainer> columnIndexContainerMap) {
SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class);
when(segmentMetadata.getName()).thenReturn("seg");
// getColumnMetadataMap() is declared as a TreeMap, so an immutable
Map.of() will not do here.
when(segmentMetadata.getColumnMetadataMap()).thenReturn(new TreeMap<>());
- return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata,
Map.of(), null);
+ return new ImmutableSegmentImpl(segmentDirectory, segmentMetadata,
columnIndexContainerMap, null);
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java
index 5b320aec117..d75729bd048 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java
@@ -56,6 +56,7 @@ import org.apache.pinot.spi.config.table.UpsertConfig;
import org.apache.pinot.spi.data.DimensionFieldSpec;
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.ByteArray;
import org.apache.pinot.spi.utils.BytesUtils;
@@ -73,6 +74,8 @@ import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
import static org.mockito.Mockito.never;
@@ -101,6 +104,7 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest
ThreadSafeMutableRoaringBitmap validDocIds, @Nullable
ThreadSafeMutableRoaringBitmap queryableDocIds,
@Nullable List<PrimaryKey> primaryKeys) {
ImmutableSegmentImpl segment = mock(ImmutableSegmentImpl.class);
+ when(segment.tryAcquireReadLock()).thenReturn(true);
when(segment.getSegmentName()).thenReturn(getSegmentName(sequenceNumber));
when(segment.getValidDocIds()).thenReturn(validDocIds);
when(segment.getQueryableDocIds()).thenReturn(queryableDocIds);
@@ -270,6 +274,75 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest
FileUtils.forceDelete(INDEX_DIR);
}
+ @Test
+ public void testRevertDropsKeyWhenPreviousSegmentIsDestroyed()
+ throws IOException {
+ // dropOutOfOrderRecord turns on _previousKeyToRecordLocationMap tracking
+ UpsertContext upsertContext =
_contextBuilder.setDropOutOfOrderRecord(true).build();
+ ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes
upsertMetadataManager =
+ new
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes(REALTIME_TABLE_NAME,
0, upsertContext);
+
+ int[] primaryKeys = new int[]{10, 20};
+ ThreadSafeMutableRoaringBitmap validDocIds1 = new
ThreadSafeMutableRoaringBitmap();
+ ImmutableSegmentImpl segment1 =
+ mockImmutableSegment(1, validDocIds1, null,
List.of(makePrimaryKey(10), makePrimaryKey(20)));
+ upsertMetadataManager.addSegment(segment1, validDocIds1, null,
+ List.of(new RecordInfo(makePrimaryKey(10), 0, 100, false), new
RecordInfo(makePrimaryKey(20), 1, 200, false))
+ .iterator());
+
+ ThreadSafeMutableRoaringBitmap validDocIds2 = new
ThreadSafeMutableRoaringBitmap();
+ MutableSegment mutableSegment = mockMutableSegmentWithDataSource(2,
validDocIds2, null, primaryKeys);
+ upsertMetadataManager.addRecord(mutableSegment, new
RecordInfo(makePrimaryKey(10), 0, 150, false));
+ upsertMetadataManager.addRecord(mutableSegment, new
RecordInfo(makePrimaryKey(20), 1, 250, false));
+ assertEquals(upsertMetadataManager._previousKeyToRecordLocationMap.size(),
2);
+ // segment1 has no valid docs left, so removing it skips metadata cleanup
and leaves the previous locations dangling
+ assertTrue(validDocIds1.getMutableRoaringBitmap().isEmpty());
+
+ // Previous segment destroyed: nothing to revert to, the keys are dropped
without reading the segment
+ doReturn(false).when(segment1).tryAcquireReadLock();
+ upsertMetadataManager.revertAndRemoveSegment(mutableSegment,
+ List.of(Map.entry(0, makePrimaryKey(10)), Map.entry(1,
makePrimaryKey(20))).iterator());
+ assertTrue(upsertMetadataManager._primaryKeyToRecordLocationMap.isEmpty());
+
assertTrue(upsertMetadataManager._previousKeyToRecordLocationMap.isEmpty());
+ verify(segment1, never()).getDataSource(anyString());
+
+ upsertMetadataManager.stop();
+ upsertMetadataManager.close();
+ }
+
+ @Test
+ public void testPartialUpsertSkipsMergeWhenCurrentSegmentIsDestroyed()
+ throws IOException {
+ PartialUpsertHandler partialUpsertHandler =
mock(PartialUpsertHandler.class);
+ UpsertContext upsertContext =
+ _contextBuilder.setPartialUpsertHandlerSupplier(() ->
partialUpsertHandler).build();
+ ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes
upsertMetadataManager =
+ new
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes(REALTIME_TABLE_NAME,
0, upsertContext);
+
+ ThreadSafeMutableRoaringBitmap validDocIds1 = new
ThreadSafeMutableRoaringBitmap();
+ ImmutableSegmentImpl segment1 =
+ mockImmutableSegment(1, validDocIds1, null, List.of(makePrimaryKey(1),
makePrimaryKey(2)));
+ upsertMetadataManager.addSegment(segment1, validDocIds1, null,
+ List.of(new RecordInfo(makePrimaryKey(1), 0, 100, false), new
RecordInfo(makePrimaryKey(2), 1, 200, false))
+ .iterator());
+
+ // Current segment alive: the previous row is read from it and merged
+ GenericRow record1 = new GenericRow();
+ upsertMetadataManager.updateRecord(record1, new
RecordInfo(makePrimaryKey(1), 0, 150, false));
+ verify(partialUpsertHandler).merge(any(), same(record1), any());
+
+ // Current segment destroyed: the previous row must not be read, the
record is returned as is
+ doReturn(false).when(segment1).tryAcquireReadLock();
+ GenericRow record2 = new GenericRow();
+ assertSame(upsertMetadataManager.updateRecord(record2, new
RecordInfo(makePrimaryKey(2), 1, 250, false)),
+ record2);
+ // Still exactly one merge in total: none happened for the destroyed
segment
+ verify(partialUpsertHandler, times(1)).merge(any(), any(), any());
+
+ upsertMetadataManager.stop();
+ upsertMetadataManager.close();
+ }
+
@Test
public void testStartFinishOperation() {
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes
upsertMetadataManager =
@@ -1303,6 +1376,8 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest
MutableSegment segment = mockMutableSegmentWithDataSource(1,
validDocIds, null, new int[]{10});
ImmutableSegmentImpl previousSegment = mock(ImmutableSegmentImpl.class);
when(previousSegment.getSegmentName()).thenReturn(getSegmentName(0));
+ // A live segment: the revert reads its columns, so the destroy guard
must let it through
+ when(previousSegment.tryAcquireReadLock()).thenReturn(true);
ThreadSafeMutableRoaringBitmap previousValidDocIds = new
ThreadSafeMutableRoaringBitmap();
when(previousSegment.getValidDocIds()).thenReturn(failure.equals("bitmap") ?
null : previousValidDocIds);
PrimaryKey key = makePrimaryKey(10);
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java
index 1571dfc24e3..1b331cf03f7 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java
@@ -85,6 +85,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
@@ -943,6 +944,7 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerTest {
ThreadSafeMutableRoaringBitmap validDocIds, @Nullable
ThreadSafeMutableRoaringBitmap queryableDocIds,
List<PrimaryKey> primaryKeys, @Nullable int[] timestamps) {
ImmutableSegmentImpl segment = mock(ImmutableSegmentImpl.class);
+ when(segment.tryAcquireReadLock()).thenReturn(true);
when(segment.getSegmentName()).thenReturn(getSegmentName(sequenceNumber));
when(segment.getValidDocIds()).thenReturn(validDocIds);
when(segment.getQueryableDocIds()).thenReturn(queryableDocIds);
@@ -2501,6 +2503,8 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerTest {
MutableSegment segment = mockMutableSegmentWithDataSource(1,
validDocIds, null, new int[]{10});
ImmutableSegmentImpl previousSegment = mock(ImmutableSegmentImpl.class);
when(previousSegment.getSegmentName()).thenReturn(getSegmentName(0));
+ // A live segment: the revert reads its columns, so the destroy guard
must let it through
+ when(previousSegment.tryAcquireReadLock()).thenReturn(true);
when(previousSegment.getValidDocIds()).thenReturn(
failure.equals("bitmap") ? null : new
ThreadSafeMutableRoaringBitmap());
PrimaryKey key = makePrimaryKey(10);
@@ -2598,6 +2602,84 @@ public class
ConcurrentMapPartitionUpsertMetadataManagerTest {
upsertMetadataManager.close();
}
+ @Test
+ public void testRevertDropsKeyWhenPreviousSegmentIsDestroyed()
+ throws IOException {
+ // dropOutOfOrderRecord turns on _previousKeyToRecordLocationMap tracking
+ UpsertContext upsertContext =
_contextBuilder.setDropOutOfOrderRecord(true).build();
+ ConcurrentMapPartitionUpsertMetadataManager upsertMetadataManager =
+ new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME,
0, upsertContext);
+ Map<Object, RecordLocation> recordLocationMap =
upsertMetadataManager._primaryKeyToRecordLocationMap;
+
+ int[] primaryKeys = new int[]{10, 20, 30};
+ int[] timestamps = new int[]{100, 200, 300};
+ ThreadSafeMutableRoaringBitmap validDocIds1 = new
ThreadSafeMutableRoaringBitmap();
+ ImmutableSegmentImpl segment1 =
+ mockImmutableSegmentWithTimestamps(1, validDocIds1, null,
getPrimaryKeyList(3, primaryKeys), timestamps);
+ upsertMetadataManager.addSegment(segment1, validDocIds1, null,
+ getRecordInfoListWithIntegerComparison(3, primaryKeys, timestamps,
null).iterator());
+
+ ThreadSafeMutableRoaringBitmap validDocIds2 = new
ThreadSafeMutableRoaringBitmap();
+ MutableSegment mutableSegment = mockMutableSegmentWithDataSource(2,
validDocIds2, null, primaryKeys);
+ upsertMetadataManager.addRecord(mutableSegment, new
RecordInfo(makePrimaryKey(10), 0, 150, false));
+ upsertMetadataManager.addRecord(mutableSegment, new
RecordInfo(makePrimaryKey(20), 1, 250, false));
+ upsertMetadataManager.addRecord(mutableSegment, new
RecordInfo(makePrimaryKey(30), 2, 350, false));
+ assertEquals(upsertMetadataManager._previousKeyToRecordLocationMap.size(),
3);
+ // segment1 has no valid docs left, so removing it skips metadata cleanup
and leaves the previous locations dangling
+ assertTrue(validDocIds1.getMutableRoaringBitmap().isEmpty());
+
+ // Previous segment alive: the key is reverted to it
+ upsertMetadataManager.revertAndRemoveSegment(mutableSegment,
List.of(Map.entry(0, makePrimaryKey(10))).iterator());
+ checkRecordLocation(recordLocationMap, 10, segment1, 0, 100,
HashFunction.NONE);
+ assertEquals(validDocIds1.getMutableRoaringBitmap().toArray(), new
int[]{0});
+
+ // Previous segment destroyed: nothing to revert to, the keys are dropped
without reading the segment
+ doReturn(false).when(segment1).tryAcquireReadLock();
+ upsertMetadataManager.revertAndRemoveSegment(mutableSegment,
+ List.of(Map.entry(1, makePrimaryKey(20)), Map.entry(2,
makePrimaryKey(30))).iterator());
+ assertEquals(recordLocationMap.size(), 1);
+ checkRecordLocation(recordLocationMap, 10, segment1, 0, 100,
HashFunction.NONE);
+
assertTrue(upsertMetadataManager._previousKeyToRecordLocationMap.isEmpty());
+ assertEquals(validDocIds1.getMutableRoaringBitmap().toArray(), new
int[]{0});
+
+ upsertMetadataManager.stop();
+ upsertMetadataManager.close();
+ }
+
+ @Test
+ public void testPartialUpsertSkipsMergeWhenCurrentSegmentIsDestroyed()
+ throws IOException {
+ PartialUpsertHandler partialUpsertHandler =
mock(PartialUpsertHandler.class);
+ UpsertContext upsertContext =
+ _contextBuilder.setPartialUpsertHandlerSupplier(() ->
partialUpsertHandler).build();
+ ConcurrentMapPartitionUpsertMetadataManager upsertMetadataManager =
+ new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME,
0, upsertContext);
+
+ int[] primaryKeys = new int[]{1, 2};
+ int[] timestamps = new int[]{100, 200};
+ ThreadSafeMutableRoaringBitmap validDocIds1 = new
ThreadSafeMutableRoaringBitmap();
+ ImmutableSegmentImpl segment1 =
+ mockImmutableSegmentWithTimestamps(1, validDocIds1, null,
getPrimaryKeyList(2, primaryKeys), timestamps);
+ upsertMetadataManager.addSegment(segment1, validDocIds1, null,
+ getRecordInfoListWithIntegerComparison(2, primaryKeys, timestamps,
null).iterator());
+
+ // Current segment alive: the previous row is read from it and merged
+ GenericRow record1 = new GenericRow();
+ upsertMetadataManager.updateRecord(record1, new
RecordInfo(makePrimaryKey(1), 0, 150, false));
+ verify(partialUpsertHandler).merge(any(), same(record1), any());
+
+ // Current segment destroyed: the previous row must not be read, the
record is returned as is
+ doReturn(false).when(segment1).tryAcquireReadLock();
+ GenericRow record2 = new GenericRow();
+ assertSame(upsertMetadataManager.updateRecord(record2, new
RecordInfo(makePrimaryKey(2), 1, 250, false)),
+ record2);
+ // Still exactly one merge in total: none happened for the destroyed
segment
+ verify(partialUpsertHandler, times(1)).merge(any(), any(), any());
+
+ upsertMetadataManager.stop();
+ upsertMetadataManager.close();
+ }
+
@Test
public void testShouldRevertMetadataOnlyForConsumingSegmentsInProtectedMode()
throws IOException {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
