Make segment completion protocol configurable from table config

apache · Dec 13, 2024 · 01e4b9c · 01e4b9c
1 parent 0c824ae
commit 01e4b9c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/...c/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java b/...c/main/java/org/apache/pinot/controller/helix/core/realtime/SegmentCompletionManager.java
@@ -32,6 +32,7 @@
 import org.apache.pinot.controller.helix.core.realtime.segment.CommittingSegmentDescriptor;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.stream.StreamConfig;
+import org.apache.pinot.spi.stream.StreamConfigProperties;
 import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
 import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
 import org.apache.pinot.spi.stream.StreamPartitionMsgOffsetFactory;
@@ -126,13 +127,21 @@ private SegmentCompletionFSM createFsm(LLCSegmentName llcSegmentName, String msg
     String segmentName = llcSegmentName.getSegmentName();
     SegmentZKMetadata segmentMetadata = _segmentManager.getSegmentZKMetadata(realtimeTableName, segmentName, null);
     Preconditions.checkState(segmentMetadata != null, "Failed to find ZK metadata for segment: %s", segmentName);
-    SegmentCompletionFSM fsm;
 
-    String factoryName = _segmentCompletionConfig.getDefaultFsmScheme();
+    TableConfig tableConfig = _segmentManager.getTableConfig(realtimeTableName);
+    String factoryName;
+    try {
+      Map<String, String> streamConfigMap = IngestionConfigUtils.getStreamConfigMap(tableConfig);
+      factoryName = streamConfigMap.get(StreamConfigProperties.SEGMENT_COMPLETION_FSM_SCHEME);
+    } catch (Exception e) {
+      factoryName = _segmentCompletionConfig.getDefaultFsmScheme();
+    }
+
     Preconditions.checkState(SegmentCompletionFSMFactory.isFactoryTypeSupported(factoryName),
         "No FSM registered for name: " + factoryName);
 
-    fsm = SegmentCompletionFSMFactory.createFSM(factoryName, this, _segmentManager, llcSegmentName, segmentMetadata);
+    SegmentCompletionFSM fsm =
+        SegmentCompletionFSMFactory.createFSM(factoryName, this, _segmentManager, llcSegmentName, segmentMetadata);
     fsm.transitionToInitialState(msgType);
 
     LOGGER.info("Created FSM {}", fsm);

diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamConfigProperties.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamConfigProperties.java
@@ -130,6 +130,11 @@ private StreamConfigProperties() {
    */
   public static final String SERVER_UPLOAD_TO_DEEPSTORE = "realtime.segment.serverUploadToDeepStore";
 
+  /**
+   * Config used to indicate which segment commit protocol implementation controller should use for this table
+   */
+  public static final String SEGMENT_COMPLETION_FSM_SCHEME = "segment.completion.fsm.scheme";
+
   /**
    * Helper method to create a stream specific property
    */