This is an automated email from the ASF dual-hosted git repository.
KKcorps 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 7f024b6b623 Release the consumer semaphore only after the segment's
upsert metadata is removed on offload (#19444)
7f024b6b623 is described below
commit 7f024b6b6230f88f47526e7fb0f1479757bce6ca
Author: Kartik Khare <[email protected]>
AuthorDate: Tue Sep 8 14:11:04 2026 +0530
Release the consumer semaphore only after the segment's upsert metadata is
removed on offload (#19444)
RealtimeSegmentDataManager.doOffload() released the consumer semaphore
before _realtimeSegment.offload() removed the segment from the upsert and
dedup metadata managers. For partial-upsert tables in PROTECTED consistency
mode that removal reverts primary keys to their previous record locations,
so the next consuming segment of the partition could start replaying while
keys still pointed at the mutable segment being offloaded, and merge the
same update twice.
closeStreamConsumer() used to both close the stream consumer and release
the semaphore. Split the two: closeStreamConsumer() now only closes, and
still runs before the segment is offloaded so nothing can consume into an
offloaded segment. The build and download paths that let the next segment
start early call closeStreamConsumerAndReleaseSemaphore(), which keeps
their behavior under ALLOW_DURING_BUILD_ONLY and ALLOW_ALWAYS unchanged.
doOffload() releases the semaphore in a finally block after the metadata
removal, so a failure there cannot leave the partition unable to consume.
Co-authored-by: Kartik Khare <[email protected]>
---
.../realtime/RealtimeSegmentDataManager.java | 32 +++++-
.../realtime/RealtimeSegmentDataManagerTest.java | 111 +++++++++++++++++++++
2 files changed, 138 insertions(+), 5 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java
index 71a61b92d5b..4c511c00b2d 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java
@@ -1198,7 +1198,7 @@ public class RealtimeSegmentDataManager extends
SegmentDataManager {
protected SegmentBuildDescriptor buildSegmentInternal(boolean forCommit)
throws SegmentBuildFailureException {
if (_parallelSegmentConsumptionPolicy.isAllowedDuringBuild()) {
- closeStreamConsumer();
+ closeStreamConsumerAndReleaseSemaphore();
}
// Do not allow building segment when table data manager is already shut
down
if (_realtimeTableDataManager.isShutDown()) {
@@ -1445,15 +1445,25 @@ public class RealtimeSegmentDataManager extends
SegmentDataManager {
}
}
+ /// Closes the stream consumer so that no more data can be consumed into
this segment. Does NOT release the consumer
+ /// semaphore, see [#closeStreamConsumerAndReleaseSemaphore()] and
[#doOffload()].
private void closeStreamConsumer() {
if (_streamConsumerClosed.compareAndSet(false, true)) {
closePartitionGroupConsumer();
closePartitionMetadataProvider();
- releaseConsumerSemaphore();
_transformPipeline.reportStats();
}
}
+ /// Closes the stream consumer and releases the consumer semaphore so that
the next consuming segment of the
+ /// partition can start consuming in parallel with the build or download of
this segment. Only called when the
+ /// [ParallelSegmentConsumptionPolicy] allows it.
+ @VisibleForTesting
+ void closeStreamConsumerAndReleaseSemaphore() {
+ closeStreamConsumer();
+ releaseConsumerSemaphore();
+ }
+
private void closePartitionGroupConsumer() {
try {
_partitionGroupConsumer.close();
@@ -1722,7 +1732,7 @@ public class RealtimeSegmentDataManager extends
SegmentDataManager {
protected void downloadSegmentAndReplace(SegmentZKMetadata segmentZKMetadata)
throws Exception {
if (_parallelSegmentConsumptionPolicy.isAllowedDuringDownload()) {
- closeStreamConsumer();
+ closeStreamConsumerAndReleaseSemaphore();
}
_realtimeTableDataManager.downloadAndReplaceConsumingSegment(segmentZKMetadata);
}
@@ -1764,9 +1774,21 @@ public class RealtimeSegmentDataManager extends
SegmentDataManager {
} catch (Exception e) {
_segmentLogger.error("Caught exception while stopping the consumer
thread", e);
}
+ // Close the stream consumer first so that nothing can consume into the
segment once it is offloaded.
closeStreamConsumer();
- cleanupMetrics();
- _realtimeSegment.offload();
+ // Remove this segment's upsert/dedup metadata BEFORE releasing the
consumer semaphore. For partial upsert in
+ // PROTECTED consistency mode, offload() reverts the primary keys owned by
this consuming segment to their previous
+ // record locations. If the semaphore were released first, the next
consuming segment of the partition could start
+ // replaying while primary keys still point to this mutable segment, and
merge against the un-reverted state.
+ // When the parallel consumption policy allowed the next segment to start
during build or download, the semaphore
+ // was already released there and the release below is a no-op.
+ // The semaphore is released in a finally block so that a failure in
metadata removal cannot stall the partition.
+ try {
+ _realtimeSegment.offload();
+ } finally {
+ releaseConsumerSemaphore();
+ cleanupMetrics();
+ }
}
@Override
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java
index 9ef007b2e49..8843f289cee 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java
@@ -29,7 +29,9 @@ import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.function.BooleanSupplier;
@@ -48,6 +50,7 @@ import
org.apache.pinot.core.realtime.impl.fakestream.FakeStreamConsumerFactory;
import org.apache.pinot.core.realtime.impl.fakestream.FakeStreamMessageDecoder;
import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
import org.apache.pinot.segment.local.data.manager.TableDataManager;
+import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImpl;
import
org.apache.pinot.segment.local.realtime.impl.RealtimeSegmentStatsHistory;
import org.apache.pinot.segment.local.segment.creator.Fixtures;
import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
@@ -83,8 +86,11 @@ import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.anyString;
import static org.mockito.Mockito.atLeast;
+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.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -652,6 +658,96 @@ public class RealtimeSegmentDataManagerTest {
}
}
+ @Test
+ public void
testOffloadRemovesSegmentMetadataBeforeReleasingConsumerSemaphore()
+ throws Exception {
+ // Use a fresh coordinator. Other tests release the shared semaphore
without acquiring it, which inflates permits.
+ _partitionGroupIdToConsumerCoordinatorMap.remove(PARTITION_GROUP_ID);
+ try (FakeRealtimeSegmentDataManager segmentDataManager =
createFakeSegmentManager()) {
+ Semaphore semaphore =
_partitionGroupIdToConsumerCoordinatorMap.get(PARTITION_GROUP_ID).getSemaphore();
+ Assert.assertTrue(semaphore.tryAcquire());
+ segmentDataManager.getConsumerSemaphoreAcquired().set(true);
+
+ // MutableSegmentImpl.offload() removes the segment from the
upsert/dedup metadata managers. For partial upsert in
+ // PROTECTED mode that is where the primary keys are reverted to their
previous locations, so it has to run
+ // (1) after the stream consumer is closed, so nothing can consume into
the offloaded segment, and
+ // (2) while the consumer semaphore is still held, so the next consuming
segment cannot replay against
+ // un-reverted state.
+ AtomicInteger permitsWhenMetadataRemoved = new AtomicInteger(-1);
+ AtomicBoolean streamConsumerClosedWhenMetadataRemoved = new
AtomicBoolean(false);
+ MutableSegmentImpl realtimeSegment = spy((MutableSegmentImpl)
segmentDataManager.getSegment());
+ doAnswer(invocation -> {
+ permitsWhenMetadataRemoved.set(semaphore.availablePermits());
+
streamConsumerClosedWhenMetadataRemoved.set(segmentDataManager.isStreamConsumerClosed());
+ return null;
+ }).when(realtimeSegment).offload();
+ segmentDataManager.setRealtimeSegment(realtimeSegment);
+
+ segmentDataManager.offload();
+
+ verify(realtimeSegment).offload();
+ Assert.assertTrue(streamConsumerClosedWhenMetadataRemoved.get(),
+ "Stream consumer must be closed before the segment metadata is
removed");
+ Assert.assertEquals(permitsWhenMetadataRemoved.get(), 0,
+ "Segment metadata must be removed while the consumer semaphore is
still held");
+ Assert.assertEquals(semaphore.availablePermits(), 1, "Consumer semaphore
must be released after offload");
+
Assert.assertFalse(segmentDataManager.getConsumerSemaphoreAcquired().get());
+ }
+ }
+
+ @Test
+ public void
testOffloadAfterParallelConsumptionReleaseDoesNotReleaseSemaphoreTwice()
+ throws Exception {
+ _partitionGroupIdToConsumerCoordinatorMap.remove(PARTITION_GROUP_ID);
+ try (FakeRealtimeSegmentDataManager segmentDataManager =
createFakeSegmentManager()) {
+ Semaphore semaphore =
_partitionGroupIdToConsumerCoordinatorMap.get(PARTITION_GROUP_ID).getSemaphore();
+ Assert.assertTrue(semaphore.tryAcquire());
+ segmentDataManager.getConsumerSemaphoreAcquired().set(true);
+
+ // With ALLOW_DURING_BUILD_ONLY / ALLOW_ALWAYS, buildSegmentInternal()
and downloadSegmentAndReplace() let the
+ // next consuming segment start early by closing the consumer and
releasing the semaphore together.
+ segmentDataManager.closeStreamConsumerAndReleaseSemaphore();
+ Assert.assertTrue(segmentDataManager.isStreamConsumerClosed());
+ Assert.assertEquals(semaphore.availablePermits(), 1,
+ "Consumer semaphore must be released for parallel consumption");
+
Assert.assertFalse(segmentDataManager.getConsumerSemaphoreAcquired().get());
+
+ MutableSegmentImpl realtimeSegment = spy((MutableSegmentImpl)
segmentDataManager.getSegment());
+ doAnswer(invocation -> null).when(realtimeSegment).offload();
+ segmentDataManager.setRealtimeSegment(realtimeSegment);
+
+ segmentDataManager.offload();
+
+ verify(realtimeSegment).offload();
+ Assert.assertEquals(semaphore.availablePermits(), 1, "Offload must not
release the consumer semaphore twice");
+ }
+ }
+
+ @Test
+ public void testOffloadReleasesConsumerSemaphoreWhenMetadataRemovalFails()
+ throws Exception {
+ _partitionGroupIdToConsumerCoordinatorMap.remove(PARTITION_GROUP_ID);
+ try (FakeRealtimeSegmentDataManager segmentDataManager =
createFakeSegmentManager()) {
+ Semaphore semaphore =
_partitionGroupIdToConsumerCoordinatorMap.get(PARTITION_GROUP_ID).getSemaphore();
+ Assert.assertTrue(semaphore.tryAcquire());
+ segmentDataManager.getConsumerSemaphoreAcquired().set(true);
+
+ MutableSegmentImpl realtimeSegment = spy((MutableSegmentImpl)
segmentDataManager.getSegment());
+ doThrow(new RuntimeException("metadata removal
failed")).when(realtimeSegment).offload();
+ segmentDataManager.setRealtimeSegment(realtimeSegment);
+
+ try {
+ segmentDataManager.offload();
+ Assert.fail("Expected the metadata removal failure to propagate");
+ } catch (RuntimeException e) {
+ Assert.assertEquals(e.getMessage(), "metadata removal failed");
+ }
+ // A failed metadata removal must not leave the semaphore held, or the
partition can never consume again.
+ Assert.assertEquals(semaphore.availablePermits(), 1,
+ "Consumer semaphore must be released even when metadata removal
fails");
+ }
+ }
+
@Test
public void testOnlineTransitionSkipsLocalBuildOnCrcMismatch()
throws Exception {
@@ -1419,6 +1515,21 @@ public class RealtimeSegmentDataManagerTest {
return _tableDataManager;
}
+ /// Replaces the mutable segment so tests can observe or fail its
offload().
+ public void setRealtimeSegment(MutableSegmentImpl realtimeSegment)
+ throws Exception {
+ Field realtimeSegmentField =
RealtimeSegmentDataManager.class.getDeclaredField("_realtimeSegment");
+ realtimeSegmentField.setAccessible(true);
+ realtimeSegmentField.set(this, realtimeSegment);
+ }
+
+ public boolean isStreamConsumerClosed()
+ throws Exception {
+ Field streamConsumerClosedField =
RealtimeSegmentDataManager.class.getDeclaredField("_streamConsumerClosed");
+ streamConsumerClosedField.setAccessible(true);
+ return ((AtomicBoolean) streamConsumerClosedField.get(this)).get();
+ }
+
public String getStopReason() {
try {
return (String) _stopReason.get(this);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
