This is an automated email from the ASF dual-hosted git repository.
raghavyadav01 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 779210b9831 Add onSegmentAdded hook to IndexSegment and
SegmentDirectory (#19053)
779210b9831 is described below
commit 779210b98313d845f4edf4f1182c3a53d66f1981
Author: Jhow <[email protected]>
AuthorDate: Wed Jul 29 17:50:59 2026 -0700
Add onSegmentAdded hook to IndexSegment and SegmentDirectory (#19053)
* add onSegmentAdded hook
* address PR comments
* add unit test
---
.../core/data/manager/BaseTableDataManager.java | 33 +++++
.../data/manager/BaseTableDataManagerTest.java | 138 +++++++++++++++++++++
.../indexsegment/immutable/EmptyIndexSegment.java | 46 +++++++
.../immutable/ImmutableSegmentImpl.java | 19 +++
.../immutable/ImmutableSegmentLoader.java | 4 +-
.../immutable/EmptyIndexSegmentTest.java | 103 +++++++++++++++
.../immutable/ImmutableSegmentImplTest.java | 97 +++++++++++++++
.../org/apache/pinot/segment/spi/IndexSegment.java | 19 +++
.../pinot/segment/spi/store/SegmentDirectory.java | 17 +++
9 files changed, 475 insertions(+), 1 deletion(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java
index 8bb9393bb4e..a81b6f42652 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java
@@ -1248,9 +1248,42 @@ public abstract class BaseTableDataManager implements
TableDataManager {
oldSegmentDataManager = _segmentDataManagerMap.put(segmentName,
segmentDataManager);
}
_recentlyDeletedSegments.invalidate(segmentName);
+ // Fire the post-registration lifecycle hook now that the segment is
swapped into the serving set
+ fireOnSegmentAdded(segmentName, segmentDataManager);
return oldSegmentDataManager;
}
+ /**
+ * Fires the post-registration lifecycle hook on the segments exposed by a
newly registered segment data manager.
+ * <p>
+ * The hook runs after the segment is swapped into the serving set, so a
reference is held while it runs: without it
+ * a concurrent {@link #replaceSegment} / {@link #unregisterSegment} could
drop the reference count to 0 and destroy
+ * the segment (closing its {@link
org.apache.pinot.segment.spi.store.SegmentDirectory}) while the hook is still
+ * using it. A manager that is already destroyed cannot take a reference and
has nothing left to notify, so it is
+ * skipped.
+ * <p>
+ * Failures are contained here: the segment is already serving, so a failing
hook must not fail the registration or
+ * the enclosing Helix state transition. This also keeps a manager that
exposes a null segment (possible for a
+ * custom implementation, as the default {@link
SegmentDataManager#getReportableSegments()} wraps
+ * {@link SegmentDataManager#getSegment()}) from aborting registration.
+ */
+ private void fireOnSegmentAdded(String segmentName, SegmentDataManager
segmentDataManager) {
+ if (!segmentDataManager.increaseReferenceCount()) {
+ return;
+ }
+ try {
+ for (IndexSegment segment : segmentDataManager.getReportableSegments()) {
+ if (segment != null) {
+ segment.onSegmentAdded();
+ }
+ }
+ } catch (Exception e) {
+ _logger.warn("Caught exception while firing onSegmentAdded for segment:
{}", segmentName, e);
+ } finally {
+ releaseSegment(segmentDataManager);
+ }
+ }
+
/**
* De-registering a segment ensures that no query uses the given segment
until a segment with that name is
* re-registered. There may be scenarios where the broker thinks that a
segment is available even though it has
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java
index 248ad28ca35..229ffdcab60 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java
@@ -24,6 +24,7 @@ import java.io.FileInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -41,6 +42,7 @@ import
org.apache.pinot.common.utils.fetcher.SegmentFetcherFactory;
import org.apache.pinot.core.data.manager.offline.ImmutableSegmentDataManager;
import org.apache.pinot.core.data.manager.offline.OfflineTableDataManager;
import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
+import org.apache.pinot.segment.local.indexsegment.immutable.EmptyIndexSegment;
import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
@@ -55,6 +57,7 @@ import org.apache.pinot.segment.spi.V1Constants;
import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
import org.apache.pinot.segment.spi.creator.SegmentVersion;
import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.segment.spi.store.SegmentDirectory;
import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
import org.apache.pinot.spi.config.table.TableConfig;
@@ -78,8 +81,11 @@ import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
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.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.*;
@@ -943,6 +949,138 @@ public class BaseTableDataManagerTest {
}
}
+ /**
+ * registerSegment must fire the post-registration lifecycle hook on the
backing segment of a standalone manager.
+ * A standalone ImmutableSegmentDataManager exposes its segment via
getReportableSegments() (default), while
+ * getSegments() defaults to empty — so iterating getSegments() here would
silently skip the hook for the common
+ * single-segment case. This is a regression guard for that: it fails if
registerSegment iterates getSegments().
+ */
+ @Test
+ public void testRegisterSegmentFiresOnSegmentAddedForStandaloneSegment() {
+ BaseTableDataManager tableDataManager = createTableManager();
+ ImmutableSegment immutableSegment = mock(ImmutableSegment.class);
+ when(immutableSegment.getSegmentName()).thenReturn(SEGMENT_NAME);
+ tableDataManager.registerSegment(SEGMENT_NAME, new
ImmutableSegmentDataManager(immutableSegment));
+ verify(immutableSegment).onSegmentAdded();
+ }
+
+ /**
+ * registerSegment must fire the hook once per reportable segment for a
multi-segment manager
+ */
+ @Test
+ public void testRegisterSegmentFiresOnSegmentAddedForEachReportableSegment()
{
+ BaseTableDataManager tableDataManager = createTableManager();
+ ImmutableSegment segment1 = mock(ImmutableSegment.class);
+ ImmutableSegment segment2 = mock(ImmutableSegment.class);
+ SegmentDataManager multiSegmentManager = mock(SegmentDataManager.class);
+ when(multiSegmentManager.getSegmentName()).thenReturn(SEGMENT_NAME);
+
when(multiSegmentManager.getReportableSegments()).thenReturn(List.of(segment1,
segment2));
+ // registerSegment holds a reference while the hook runs; a mock returns
false unless stubbed.
+ when(multiSegmentManager.increaseReferenceCount()).thenReturn(true);
+ tableDataManager.registerSegment(SEGMENT_NAME, multiSegmentManager);
+ verify(segment1).onSegmentAdded();
+ verify(segment2).onSegmentAdded();
+ }
+
+ /**
+ * The hook fires after the segment is swapped into the serving set, so
registerSegment must hold a reference while
+ * it runs. A manager that has already been destroyed (reference count 0)
cannot be referenced and has nothing left
+ * to notify: firing the hook there would run it against a closed
SegmentDirectory.
+ */
+ @Test
+ public void testRegisterSegmentSkipsOnSegmentAddedForDestroyedSegment() {
+ BaseTableDataManager tableDataManager = createTableManager();
+ ImmutableSegment immutableSegment = mock(ImmutableSegment.class);
+ when(immutableSegment.getSegmentName()).thenReturn(SEGMENT_NAME);
+ ImmutableSegmentDataManager segmentDataManager = new
ImmutableSegmentDataManager(immutableSegment);
+ // Drop the initial reference, mirroring a concurrent
replaceSegment()/unregisterSegment() that destroyed it.
+ assertTrue(segmentDataManager.decreaseReferenceCount());
+
+ tableDataManager.registerSegment(SEGMENT_NAME, segmentDataManager);
+
+ verify(immutableSegment, never()).onSegmentAdded();
+ }
+
+ /**
+ * The hook must not leave a net reference behind: the segment stays
acquirable after registration, and is not
+ * destroyed by the reference registerSegment took while firing the hook.
+ */
+ @Test
+ public void testRegisterSegmentRestoresReferenceCountAfterFiringHook() {
+ BaseTableDataManager tableDataManager = createTableManager();
+ ImmutableSegment immutableSegment = mock(ImmutableSegment.class);
+ when(immutableSegment.getSegmentName()).thenReturn(SEGMENT_NAME);
+ ImmutableSegmentDataManager segmentDataManager = new
ImmutableSegmentDataManager(immutableSegment);
+
+ tableDataManager.registerSegment(SEGMENT_NAME, segmentDataManager);
+
+ verify(immutableSegment).onSegmentAdded();
+ assertEquals(segmentDataManager.getReferenceCount(), 1);
+ verify(immutableSegment, never()).destroy();
+ }
+
+ /**
+ * A failing hook must not fail the registration: the segment is already
serving by then, and propagating the
+ * failure would abort the enclosing Helix state transition for a segment
that is in fact up.
+ */
+ @Test
+ public void testRegisterSegmentSucceedsWhenOnSegmentAddedThrows() {
+ BaseTableDataManager tableDataManager = createTableManager();
+ ImmutableSegment immutableSegment = mock(ImmutableSegment.class);
+ when(immutableSegment.getSegmentName()).thenReturn(SEGMENT_NAME);
+ doThrow(new
RuntimeException("boom")).when(immutableSegment).onSegmentAdded();
+ ImmutableSegmentDataManager segmentDataManager = new
ImmutableSegmentDataManager(immutableSegment);
+
+ tableDataManager.registerSegment(SEGMENT_NAME, segmentDataManager);
+
+ assertSame(tableDataManager.getSegmentDataManager(SEGMENT_NAME),
segmentDataManager);
+ assertEquals(segmentDataManager.getReferenceCount(), 1);
+ }
+
+ /**
+ * The default getReportableSegments() wraps getSegment(), so a custom
manager exposing a null segment must not
+ * abort registration.
+ */
+ @Test
+ public void testRegisterSegmentToleratesNullReportableSegment() {
+ BaseTableDataManager tableDataManager = createTableManager();
+ SegmentDataManager segmentDataManager = mock(SegmentDataManager.class);
+ when(segmentDataManager.getSegmentName()).thenReturn(SEGMENT_NAME);
+ when(segmentDataManager.increaseReferenceCount()).thenReturn(true);
+
when(segmentDataManager.getReportableSegments()).thenReturn(Arrays.asList(null,
null));
+
+ tableDataManager.registerSegment(SEGMENT_NAME, segmentDataManager);
+
+ assertSame(tableDataManager.getSegmentDataManager(SEGMENT_NAME),
segmentDataManager);
+ }
+
+ /**
+ * An upsert replacement with a consistency mode other than NONE registers
the same new segment twice: first through
+ * a DuoSegmentDataManager (whose default getReportableSegments() returns
its primary, i.e. the new segment) and then
+ * directly. The hook must reach the underlying SegmentDirectory only once,
since implementations are not required to
+ * be idempotent — a second call would e.g. upload a duplicate marker file.
+ */
+ @Test
+ public void testConsistencyModeReplacementFiresOnSegmentAddedOnce()
+ throws Exception {
+ BaseTableDataManager tableDataManager = createTableManager();
+ SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class);
+ when(segmentMetadata.getName()).thenReturn(SEGMENT_NAME);
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ // A real ImmutableSegment (not a mock) so the at-most-once guard in the
implementation is exercised.
+ EmptyIndexSegment newSegment = new EmptyIndexSegment(segmentMetadata,
segmentDirectory);
+ ImmutableSegmentDataManager newSegmentManager = new
ImmutableSegmentDataManager(newSegment);
+ ImmutableSegment oldSegment = mock(ImmutableSegment.class);
+ when(oldSegment.getSegmentName()).thenReturn(SEGMENT_NAME);
+ ImmutableSegmentDataManager oldSegmentManager = new
ImmutableSegmentDataManager(oldSegment);
+
+ // Mirrors BaseTableDataManager.replaceUpsertSegment() for a non-NONE
consistency mode.
+ tableDataManager.registerSegment(SEGMENT_NAME, new
DuoSegmentDataManager(newSegmentManager, oldSegmentManager));
+ tableDataManager.registerSegment(SEGMENT_NAME, newSegmentManager);
+
+ verify(segmentDirectory, times(1)).onSegmentAdded();
+ }
+
protected BaseTableDataManager createTableManager() {
return createTableManager(createDefaultInstanceDataManagerConfig());
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java
index 5a19761cc53..cd85e1741d2 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegment.java
@@ -21,6 +21,7 @@ package org.apache.pinot.segment.local.indexsegment.immutable;
import com.google.common.base.Preconditions;
import java.util.List;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import org.apache.pinot.segment.local.segment.index.datasource.EmptyDataSource;
import org.apache.pinot.segment.spi.ColumnMetadata;
@@ -35,9 +36,12 @@ import
org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
import org.apache.pinot.segment.spi.index.reader.InvertedIndexReader;
import org.apache.pinot.segment.spi.index.reader.TextIndexReader;
import org.apache.pinot.segment.spi.index.startree.StarTreeV2;
+import org.apache.pinot.segment.spi.store.SegmentDirectory;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -45,10 +49,26 @@ import org.apache.pinot.spi.data.readers.GenericRow;
* Such an IndexSegment contains only the metadata, and no indexes
*/
public class EmptyIndexSegment implements ImmutableSegment {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(EmptyIndexSegment.class);
+
private final SegmentMetadataImpl _segmentMetadata;
+ // The directory this empty segment was loaded from, if any. An empty
(0-doc) segment holds no index buffers, but the
+ // directory may still own resources and post-registration work that a
non-empty segment gets via
+ // ImmutableSegmentImpl. Null when loaded without a directory.
+ @Nullable
+ private final SegmentDirectory _segmentDirectory;
+ // Guards the post-registration hook so it reaches the directory at most
once per segment instance, even when the
+ // same segment is registered more than once (e.g. an upsert replacement
with a consistency mode other than NONE
+ // registers the new segment through a DuoSegmentDataManager and then
directly).
+ private final AtomicBoolean _segmentAdded = new AtomicBoolean();
public EmptyIndexSegment(SegmentMetadataImpl segmentMetadata) {
+ this(segmentMetadata, null);
+ }
+
+ public EmptyIndexSegment(SegmentMetadataImpl segmentMetadata, @Nullable
SegmentDirectory segmentDirectory) {
_segmentMetadata = segmentMetadata;
+ _segmentDirectory = segmentDirectory;
}
@Override
@@ -71,12 +91,38 @@ public class EmptyIndexSegment implements ImmutableSegment {
return _segmentMetadata.getSchema().getPhysicalColumnNames();
}
+ @Override
+ public void onSegmentAdded() {
+ // Best-effort: this fires after the segment is already serving, so a
failure cannot roll back the registration.
+ if (_segmentDirectory == null) {
+ return;
+ }
+ if (!_segmentAdded.compareAndSet(false, true)) {
+ // Already notified for this segment instance; a repeated registration
must not notify the directory again.
+ return;
+ }
+ try {
+ _segmentDirectory.onSegmentAdded();
+ } catch (Exception e) {
+ LOGGER.warn("Caught exception in onSegmentAdded for empty segment: {}.
Continuing with error.", getSegmentName(),
+ e);
+ }
+ }
+
@Override
public void offload() {
}
@Override
public void destroy() {
+ if (_segmentDirectory != null) {
+ try {
+ _segmentDirectory.close();
+ } catch (Exception e) {
+ LOGGER.warn("Failed to close segment directory for empty segment: {}.
Continuing with error.",
+ getSegmentName(), e);
+ }
+ }
}
@Nullable
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 99706255793..0791af6dd66 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
@@ -29,6 +29,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.segment.local.dedup.PartitionDedupMetadataManager;
@@ -83,6 +84,10 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
private final StarTreeIndexContainer _starTreeIndexContainer;
private final TextIndexReader _multiColumnTextIndex;
private final Map<String, DataSource> _dataSources;
+ // Guards the post-registration hook so it reaches the directory at most
once per segment instance, even when the
+ // same segment is registered more than once (e.g. an upsert replacement
with a consistency mode other than NONE
+ // registers the new segment through a DuoSegmentDataManager and then
directly).
+ private final AtomicBoolean _segmentAdded = new AtomicBoolean();
// Dedupe
private PartitionDedupMetadataManager _partitionDedupMetadataManager;
@@ -323,6 +328,20 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
_segmentDirectory.release(fetchContext);
}
+ @Override
+ public void onSegmentAdded() {
+ if (!_segmentAdded.compareAndSet(false, true)) {
+ // Already notified for this segment instance; a repeated registration
must not notify the directory again.
+ return;
+ }
+ // Best-effort: this fires after the segment is already serving, so a
failure cannot roll back the registration.
+ try {
+ _segmentDirectory.onSegmentAdded();
+ } catch (Exception e) {
+ LOGGER.warn("Caught exception in onSegmentAdded for segment: {}.
Continuing with error.", getSegmentName(), e);
+ }
+ }
+
@Override
public void offload() {
if (_partitionUpsertMetadataManager != null) {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java
index 0fb0353fca5..ae0391a76f6 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentLoader.java
@@ -206,7 +206,9 @@ public class ImmutableSegmentLoader {
throws Exception {
SegmentMetadataImpl segmentMetadata =
segmentDirectory.getSegmentMetadata();
if (segmentMetadata.getTotalDocs() == 0) {
- return new EmptyIndexSegment(segmentMetadata);
+ // Hand the directory to the empty segment so it can run
post-registration work and own closing the directory,
+ // mirroring the non-empty ImmutableSegmentImpl path.
+ return new EmptyIndexSegment(segmentMetadata, segmentDirectory);
}
// Remove columns not in schema from the metadata
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java
new file mode 100644
index 00000000000..7621e071715
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/EmptyIndexSegmentTest.java
@@ -0,0 +1,103 @@
+/**
+ * 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.indexsegment.immutable;
+
+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.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+
+public class EmptyIndexSegmentTest {
+
+ /**
+ * An empty (0-doc) segment loaded from a directory must still run the
post-registration lifecycle hook and own
+ * closing the directory, mirroring the non-empty ImmutableSegmentImpl path.
+ */
+ @Test
+ public void testOnSegmentAddedDelegatesToDirectoryAndDestroyCloses()
+ throws Exception {
+ SegmentMetadataImpl metadata = mock(SegmentMetadataImpl.class);
+ when(metadata.getName()).thenReturn("seg");
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ EmptyIndexSegment segment = new EmptyIndexSegment(metadata,
segmentDirectory);
+
+ segment.onSegmentAdded();
+ verify(segmentDirectory).onSegmentAdded();
+
+ segment.destroy();
+ verify(segmentDirectory).close();
+ }
+
+ /**
+ * The hook fires after the segment is already serving, so a directory
failure must not propagate out of it.
+ */
+ @Test
+ public void testOnSegmentAddedIsBestEffort()
+ throws Exception {
+ SegmentMetadataImpl metadata = mock(SegmentMetadataImpl.class);
+ when(metadata.getName()).thenReturn("seg");
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ doThrow(new
RuntimeException("boom")).when(segmentDirectory).onSegmentAdded();
+ EmptyIndexSegment segment = new EmptyIndexSegment(metadata,
segmentDirectory);
+
+ // Must not throw.
+ segment.onSegmentAdded();
+ }
+
+ /**
+ * The hook must reach the directory at most once per segment instance: the
same segment can be registered more than
+ * once (e.g. an upsert replacement with a consistency mode other than NONE
registers it through a
+ * DuoSegmentDataManager and then directly), and implementations are not
required to be idempotent.
+ */
+ @Test
+ public void testOnSegmentAddedNotifiesDirectoryAtMostOnce()
+ throws Exception {
+ SegmentMetadataImpl metadata = mock(SegmentMetadataImpl.class);
+ when(metadata.getName()).thenReturn("seg");
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ EmptyIndexSegment segment = new EmptyIndexSegment(metadata,
segmentDirectory);
+
+ segment.onSegmentAdded();
+ segment.onSegmentAdded();
+ segment.onSegmentAdded();
+
+ verify(segmentDirectory, times(1)).onSegmentAdded();
+ }
+
+ /**
+ * An empty segment loaded without a directory (e.g. the local File path)
must treat the lifecycle callbacks as safe
+ * no-ops.
+ */
+ @Test
+ public void testMetadataOnlySegmentHasNoDirectorySideEffects() {
+ SegmentMetadataImpl metadata = mock(SegmentMetadataImpl.class);
+ when(metadata.getName()).thenReturn("seg");
+ EmptyIndexSegment segment = new EmptyIndexSegment(metadata);
+
+ // No directory: both must be safe no-ops.
+ segment.onSegmentAdded();
+ segment.destroy();
+ }
+}
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
new file mode 100644
index 00000000000..77f7d6cad31
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImplTest.java
@@ -0,0 +1,97 @@
+/**
+ * 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.indexsegment.immutable;
+
+import java.util.Map;
+import java.util.TreeMap;
+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.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+
+/**
+ * Tests for the post-registration lifecycle hook of {@link
ImmutableSegmentImpl}.
+ */
+public class ImmutableSegmentImplTest {
+
+ /**
+ * The hook must reach the directory at most once per segment instance: the
same segment can be registered more than
+ * once (e.g. an upsert replacement with a consistency mode other than NONE
registers it through a
+ * DuoSegmentDataManager and then directly), and implementations are not
required to be idempotent.
+ */
+ @Test
+ public void testOnSegmentAddedNotifiesDirectoryAtMostOnce()
+ throws Exception {
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ ImmutableSegmentImpl segment = createSegment(segmentDirectory);
+
+ segment.onSegmentAdded();
+ segment.onSegmentAdded();
+ segment.onSegmentAdded();
+
+ verify(segmentDirectory, times(1)).onSegmentAdded();
+ }
+
+ /**
+ * The hook fires after the segment is already serving, so a directory
failure must not propagate out of it.
+ */
+ @Test
+ public void testOnSegmentAddedIsBestEffort()
+ throws Exception {
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ doThrow(new
RuntimeException("boom")).when(segmentDirectory).onSegmentAdded();
+ ImmutableSegmentImpl segment = createSegment(segmentDirectory);
+
+ // Must not throw.
+ segment.onSegmentAdded();
+
+ verify(segmentDirectory).onSegmentAdded();
+ }
+
+ /**
+ * A failed attempt consumes the single notification: the directory was
already told, and retry semantics are the
+ * implementation's business, not the caller's.
+ */
+ @Test
+ public void testFailedOnSegmentAddedIsNotRetried()
+ throws Exception {
+ SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+ doThrow(new
RuntimeException("boom")).when(segmentDirectory).onSegmentAdded();
+ ImmutableSegmentImpl segment = createSegment(segmentDirectory);
+
+ segment.onSegmentAdded();
+ segment.onSegmentAdded();
+
+ verify(segmentDirectory, times(1)).onSegmentAdded();
+ }
+
+ private static ImmutableSegmentImpl createSegment(SegmentDirectory
segmentDirectory) {
+ 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);
+ }
+}
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
index b1d6b4592d4..d69d96701d3 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
@@ -165,6 +165,25 @@ public interface IndexSegment {
default void release(FetchContext fetchContext) {
}
+ /**
+ * Lifecycle callback invoked after this segment has been successfully added
to the serving set (i.e. registered and
+ * swapped in so queries can reach it). Default is a no-op; implementations
may override to run post-registration
+ * work. It is invoked only on the success path, and at most once per
segment instance even if the same segment is
+ * registered more than once.
+ * <p>
+ * Implementations must be best-effort: this fires after the segment is
already serving, so a failure here cannot
+ * roll back the registration and should be handled internally rather than
propagated.
+ * <p>
+ * Implementations must also return promptly. This runs inline on the
segment registration thread, which is the
+ * Helix state-transition thread for the ONLINE transition, so a blocking or
unbounded call here stalls the state
+ * transition itself; the caller applies no timeout. Any unbounded work (in
particular remote I/O) must be bounded
+ * by the implementation or performed asynchronously by it. An
implementation that goes asynchronous must not assume
+ * the segment is still alive once this method has returned, since the
caller only keeps the segment referenced for
+ * the duration of the call.
+ */
+ default void onSegmentAdded() {
+ }
+
/**
* Offloads the segment from the metadata management (e.g. upsert metadata),
but not releases the resources yet
* because there might be queries still accessing the segment.
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
index 3e0541c5292..18ffb6d748d 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
@@ -151,6 +151,23 @@ public abstract class SegmentDirectory implements
Closeable {
throws Exception {
}
+ /**
+ * Lifecycle callback invoked after the segment backed by this directory has
been successfully added to the serving
+ * set (registered and swapped in). Default is a no-op, invoked at most once
per segment instance backed by this
+ * directory.
+ * <p>
+ * Must be best-effort: it fires after the segment is already serving, so
failures cannot roll back the registration
+ * and should be handled internally.
+ * <p>
+ * Must also return promptly: this runs inline on the segment registration
thread, which is the Helix
+ * state-transition thread for the ONLINE transition, so blocking or
unbounded work here stalls the state transition.
+ * Bound any remote I/O (e.g. with a request timeout) or perform it
asynchronously, keeping in mind that the segment
+ * and this directory are only guaranteed to be alive for the duration of
this call.
+ */
+ public void onSegmentAdded()
+ throws Exception {
+ }
+
/**
* Get the storage tier where the segment directory is placed by server.
*
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
