Jackie-Jiang commented on code in PR #19596:
URL: https://github.com/apache/pinot/pull/19596#discussion_r4201150082
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java:
##########
@@ -1332,6 +1332,20 @@ public void untrackSegmentForUpsertView(IndexSegment
segment) {
}
}
+ /// 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 are kept alive
by their consumer and always succeed.
+ /// 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();
Review Comment:
Non-blocking follow-up: the mutable-segment exemption still leaves a
use-after-free path during table shutdown. This also exists on the base commit;
it is not a regression introduced by this PR.
With pauseless partial upsert (`ALLOW_DURING_BUILD_ONLY`), predecessor A can
release its consumer semaphore and wait for the segment-build semaphore while
successor B reads an A row through `LazyRow`.
`RealtimeTableDataManager.doShutdown()` stops the upsert manager without
draining pending operations, then offloads/destroys segment managers
concurrently. A's cleanup joins A's consumer, but `removeSegment(A)` skips
metadata cleanup because the upsert manager is stopped. A can therefore free
its off-heap buffers while B's separate cleanup is still waiting for B's
in-flight merge. The immutable-add shutdown barrier does not cover A while it
is waiting to build.
Draining pending upsert operations before destroying segments, or extending
lifetime protection to mutable segments, would close this path. At minimum,
please qualify the unconditional lifecycle claim here and the matching comments
in both `doUpdateRecord()` implementations: stopping the owning consumer does
not keep a mutable predecessor alive for another consumer during shutdown.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java:
##########
@@ -404,6 +417,30 @@ public void destroy() {
}
}
+ /// True once [#destroy()] has set the flag: the index buffers may be closed
and must not be read.
+ public boolean isDestroyed() {
Review Comment:
Cleanup: `isDestroyed()` has no production callers; all callers are in
`ImmutableSegmentImplTest`, in the same package. We can make this
package-private with `@VisibleForTesting` instead of adding a public API.
Alternatively, the tests can assert the observable acquisition/close behavior
and avoid the accessor.
##########
pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java:
##########
@@ -1070,6 +1072,7 @@ private static ImmutableSegmentImpl
mockUploadedImmutableSegment(String suffix,
creationTimeMs = System.currentTimeMillis();
}
ImmutableSegmentImpl segment = mock(ImmutableSegmentImpl.class);
+ when(segment.tryAcquireReadLock()).thenReturn(true);
Review Comment:
Cleanup: this new stub is also unused by the uploaded-segment fixture's
current callers. `mockUploadedImmutableSegment()` is used only by
`verifyAddReplaceUploadedSegment1`, which adds/replaces/removes immutable
segments without entering partial-update or consuming-segment revert. The same
applies to the added stub in
`ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.mockUploadedImmutableSegment()`
(line 141). Both can be removed. The live-lock stubs in the general immutable
fixtures and existing revert-failure tests should remain.
##########
pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java:
##########
@@ -1465,6 +1468,7 @@ public void
testPreloadSegmentEmptyValidDocIdsWithDeleteColumn() {
new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME,
0, _contextBuilder.build());
ImmutableSegmentImpl segment = mock(ImmutableSegmentImpl.class);
+ when(segment.tryAcquireReadLock()).thenReturn(true);
Review Comment:
Cleanup: this new stub is unused, as is the corresponding stub at line 1500
in `testPreloadSegmentEmptyValidDocIdsWithoutDeleteColumn`. Both tests preload
an empty valid-doc snapshot; `doPreloadSegment()` enables the empty bitmaps and
returns without reaching either guarded operation. Both stubs can be removed.
##########
pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java:
##########
@@ -79,11 +90,84 @@ public void testFailedOnSegmentAddedIsNotRetried()
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);
+ // 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();
+
+ segment.releaseReadLock();
Review Comment:
Please put the reader-held assertions in an inner `try/finally` that
releases this read lock, before waiting for the destroyer to complete.
Currently, if an assertion or verification at lines 124–126 fails, execution
skips this release and only calls `executor.shutdownNow()`. The destroyer is
blocked in the non-interruptible write-lock `lock()` call, so shutdown cannot
release it and the executor thread remains stranded. This is test failure-path
cleanup, not a production-lock defect.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java:
##########
@@ -428,23 +440,40 @@ protected GenericRow doUpdateRecord(GenericRow record,
RecordInfo recordInfo) {
if (!recordInfo.isDeleteRecord()
&&
recordInfo.getComparisonValue().compareTo(recordLocation.getComparisonValue())
>= 0) {
IndexSegment currentSegment = recordLocation.getSegment();
- ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
int currentDocId = recordLocation.getDocId();
- if (currentQueryableDocIds == null ||
currentQueryableDocIds.contains(currentDocId)) {
+ // Read lock: currentSegment cannot be destroyed while LazyRow
reads its columns. A consuming segment needs
+ // no lock: 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.
+ if (tryAcquireSegmentReadLock(currentSegment)) {
try {
- _reusePreviousRow.init(currentSegment, currentDocId);
- _partialUpsertHandler.merge(_reusePreviousRow, record,
_reuseMergeResultHolder);
+ mergeWithPreviousRecord(currentSegment, currentDocId, record);
} finally {
- _reuseMergeResultHolder.clear();
- _reusePreviousRow.clear();
+ releaseSegmentReadLock(currentSegment);
}
+ } else {
+ _logger.debug("Current segment: {} is destroyed, storing record
without merging the previous row",
+ currentSegment.getSegmentName());
}
}
return recordLocation;
});
return record;
}
+ /// Merges the previous record into `record` unless the previous record is
marked deleted.
+ private void mergeWithPreviousRecord(IndexSegment currentSegment, int
currentDocId, GenericRow record) {
+ ThreadSafeMutableRoaringBitmap currentQueryableDocIds =
currentSegment.getQueryableDocIds();
+ if (currentQueryableDocIds == null ||
currentQueryableDocIds.contains(currentDocId)) {
Review Comment:
Optional simplification, also applicable to the consistent-deletes
implementation: this new helper has only one caller, so we can retain the
original merge block inline and wrap it with lock acquisition/release,
preserving both reusable-state cleanup and release in `finally`.
Also, the `queryableDocIds` check can stay before acquisition. The bitmap is
heap-backed and is not freed by segment destruction; a deleted predecessor
needs no column read. Moving the check inside this helper now acquires/releases
the lifetime lock even when no merge is needed. Keeping the original check
outside makes the diff smaller and avoids that unnecessary locking.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
