Pauseless Consumption #3: Disaster Recovery with Reingestion

pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java

@@ -127,6 +127,7 @@ public static FileUploadType getDefaultUploadType() {
   private static final String FORCE_CLEANUP_PARAMETER = "&forceCleanup=";
 
   private static final String RETENTION_PARAMETER = "retention=";
+  public static final String REINGEST_SEGMENT_PATH = "/reingestSegment";
 
   private static final List<String> SUPPORTED_PROTOCOLS = Arrays.asList(HTTP, HTTPS);
 
@@ -1274,6 +1275,47 @@ public File downloadUntarFileStreamed(URI uri, File dest, AuthProvider authProvi
         httpHeaders, maxStreamRateInByte);
   }
 
+  /**
+   * Invokes the server's reIngestSegment API via a POST request with JSON payload,
+   * using Simple HTTP APIs.
+   *
+   * POST http://[serverURL]/reIngestSegment
+   * {
+   *   "tableNameWithType": [tableName],
+   *   "segmentName": [segmentName]
+   * }
+   */
+  public void triggerReIngestion(String serverHostPort, String tableNameWithType, String segmentName)
+      throws IOException, URISyntaxException, HttpErrorStatusException {
+    String scheme = HTTP;
+    if (serverHostPort.contains(HTTPS)) {
+      scheme = HTTPS;
+      serverHostPort = serverHostPort.replace(HTTPS + "://", "");
+    } else if (serverHostPort.contains(HTTP)) {
+      serverHostPort = serverHostPort.replace(HTTP + "://", "");
+    }
+
+    String serverHost = serverHostPort.split(":")[0];
+    String serverPort = serverHostPort.split(":")[1];
+
+    URI reIngestUri = getURI(scheme, serverHost, Integer.parseInt(serverPort), REINGEST_SEGMENT_PATH);
+
+    Map<String, Object> requestJson = new HashMap<>();
+    requestJson.put("tableNameWithType", tableNameWithType);
+    requestJson.put("segmentName", segmentName);
+
+    String jsonPayload = JsonUtils.objectToString(requestJson);
+    SimpleHttpResponse response =
+        HttpClient.wrapAndThrowHttpException(_httpClient.sendJsonPostRequest(reIngestUri, jsonPayload));
+
+    // Check that we got a 2xx response
+    int statusCode = response.getStatusCode();
+    if (statusCode / 100 != 2) {
+      throw new IOException(String.format("Failed POST to %s, HTTP %d: %s",
+          reIngestUri, statusCode, response.getResponse()));
+    }
+  }
+
   /**
    * Generate a param list with a table name attribute.
    *

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java

@@ -21,6 +21,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.BiMap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import java.io.IOException;
@@ -48,8 +49,11 @@
 import org.apache.helix.ClusterMessagingService;
 import org.apache.helix.Criteria;
 import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixDataAccessor;
 import org.apache.helix.HelixManager;
 import org.apache.helix.InstanceType;
+import org.apache.helix.PropertyKey;
+import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
@@ -442,6 +446,10 @@ public IdealState getIdealState(String realtimeTableName) {
     }
   }
 
+  public ExternalView getExternalView(String realtimeTableName) {
+    return _helixResourceManager.getTableExternalView(realtimeTableName);
+  }
+
   @VisibleForTesting
   void setIdealState(String realtimeTableName, IdealState idealState) {
     try {
@@ -2096,6 +2104,137 @@ URI createSegmentPath(String rawTableName, String segmentName) {
     return URIUtils.getUri(_controllerConf.getDataDir(), rawTableName, URIUtils.encode(segmentName));
   }
 
+  /**
+   * Re-ingests segments that are in ERROR state in EV but ONLINE in IS with no peer copy on any server. This method
+   * will call the server reIngestSegment API
+   * on one of the alive servers that are supposed to host that segment according to IdealState.
+   *
+   * API signature:
+   *   POST http://[serverURL]/reIngestSegment
+   *   Request body (JSON):
+   *   {
+   *     "tableNameWithType": [tableName],
+   *     "segmentName": [segmentName]
+   *   }
+   *
+   * @param tableNameWithType The table name with type, e.g. "myTable_REALTIME"
+   */
+  public void reIngestSegmentsWithErrorState(String tableNameWithType) {
+    // Step 1: Fetch the ExternalView and all segments
+    ExternalView externalView = getExternalView(tableNameWithType);
+    IdealState idealState = getIdealState(tableNameWithType);
+    Map<String, Map<String, String>> segmentToInstanceCurrentStateMap = externalView.getRecord().getMapFields();
+    Map<String, Map<String, String>> segmentToInstanceIdealStateMap = idealState.getRecord().getMapFields();
+
+    // find segments in ERROR state in externalView
+    List<String> segmentsInErrorState = new ArrayList<>();
+    for (Map.Entry<String, Map<String, String>> entry : segmentToInstanceCurrentStateMap.entrySet()) {
+      String segmentName = entry.getKey();
+      Map<String, String> instanceStateMap = entry.getValue();
+      boolean allReplicasInError = true;
+      for (String state : instanceStateMap.values()) {
+        if (!SegmentStateModel.ERROR.equals(state)) {
+          allReplicasInError = false;
+          break;
+        }
+      }
+      if (allReplicasInError) {
+        segmentsInErrorState.add(segmentName);
+      }
+    }
+
+    if (segmentsInErrorState.isEmpty()) {
+      LOGGER.info("No segments found in ERROR state for table {}", tableNameWithType);
+      return;
+    }
+
+    // filter out segments that are not ONLINE in IdealState
+    for (String segmentName : segmentsInErrorState) {
+      Map<String, String> instanceIdealStateMap = segmentToInstanceIdealStateMap.get(segmentName);
+      boolean isOnline = true;
+      for (String state : instanceIdealStateMap.values()) {
+        if (!SegmentStateModel.ONLINE.equals(state)) {
+          isOnline = false;
+          break;
+        }
+      }
+      if (!isOnline) {
+        segmentsInErrorState.remove(segmentName);
+      }
+    }
+
+    // Step 2: For each segment, check the ZK metadata for conditions
+    for (String segmentName : segmentsInErrorState) {
+      // Skip non-LLC segments or segments missing from the ideal state altogether
+      LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
+      if (llcSegmentName == null || !segmentToInstanceCurrentStateMap.containsKey(segmentName)) {
+        continue;
+      }
+
+      SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, segmentName);
+      // We only consider segments that are in COMMITTING which is indicated by having an endOffset
+      // but have a missing or placeholder download URL
+      if (segmentZKMetadata.getStatus() == Status.COMMITTING) {
+        Map<String, String> instanceStateMap = segmentToInstanceIdealStateMap.get(segmentName);
+
+        // Step 3: “No peer has that segment.” => Re-ingest from one server that is supposed to host it and is alive
+        LOGGER.info(
+            "Segment {} in table {} is COMMITTING with missing download URL and no peer copy. Triggering re-ingestion.",
+            segmentName, tableNameWithType);
+
+        // Find at least one server that should host this segment and is alive
+        String aliveServer = findAliveServerToReIngest(instanceStateMap.keySet());
+        if (aliveServer == null) {
+          LOGGER.warn("No alive server found to re-ingest segment {} in table {}", segmentName, tableNameWithType);
+          continue;
+        }
+
+        try {
+          _fileUploadDownloadClient.triggerReIngestion(aliveServer, tableNameWithType, segmentName);
+          LOGGER.info("Successfully triggered reIngestion for segment {} on server {}", segmentName, aliveServer);
+        } catch (Exception e) {
+          LOGGER.error("Failed to call reIngestSegment for segment {} on server {}", segmentName, aliveServer, e);
+        }
+      } else if (segmentZKMetadata.getStatus() == Status.UPLOADED) {
+        LOGGER.info(
+            "Segment {} in table {} is in ERROR state with download URL present. Resetting segment to ONLINE state.",
+            segmentName, tableNameWithType);
+        _helixResourceManager.resetSegment(tableNameWithType, segmentName, null);
+      }
+    }
+  }
+
+  /**
+   * Picks one 'alive' server among a set of servers that are supposed to host the segment,
+   * e.g. by checking if Helix says it is enabled or if it appears in the live instance list.
+   * This is a simple example; adapt to your environment’s definition of “alive.”
+   */
+  private String findAliveServerToReIngest(Set<String> candidateServers) {
+    // Get the current live instances from Helix
+    HelixDataAccessor helixDataAccessor = _helixManager.getHelixDataAccessor();
+    PropertyKey.Builder keyBuilder = helixDataAccessor.keyBuilder();
+    List<String> liveInstances = helixDataAccessor.getChildNames(keyBuilder.liveInstances());
+    try {
+      // This should ideally handle https scheme as well
+      BiMap<String, String> instanceToEndpointMap =
+          _helixResourceManager.getDataInstanceAdminEndpoints(candidateServers);
+
+      if (instanceToEndpointMap.isEmpty()) {
+        LOGGER.warn("No instance data admin endpoints found for servers: {}", candidateServers);
+        return null;
+      }
+
+      for (String server : candidateServers) {
+        if (liveInstances.contains(server)) {
+          return instanceToEndpointMap.get(server);
+        }
+      }
+    } catch (Exception e) {
+      LOGGER.warn("Failed to get Helix instance data admin endpoints for servers: {}", candidateServers, e);
+    }
+    return null;
+  }
+
   public Set<String> getSegmentsYetToBeCommitted(String tableNameWithType, Set<String> segmentsToCheck) {
     Set<String> segmentsYetToBeCommitted = new HashSet<>();
     for (String segmentName: segmentsToCheck) {

pinot-controller/src/main/java/org/apache/pinot/controller/validation/RealtimeSegmentValidationManager.java

@@ -28,6 +28,7 @@
 import org.apache.pinot.common.metrics.ControllerMeter;
 import org.apache.pinot.common.metrics.ControllerMetrics;
 import org.apache.pinot.common.metrics.ValidationMetrics;
+import org.apache.pinot.common.utils.PauselessConsumptionUtils;
 import org.apache.pinot.controller.ControllerConf;
 import org.apache.pinot.controller.LeadControllerManager;
 import org.apache.pinot.controller.api.resources.PauseStatusDetails;
@@ -176,6 +177,12 @@ private void runSegmentLevelValidation(TableConfig tableConfig) {
     // Update the total document count gauge
     _validationMetrics.updateTotalDocumentCountGauge(realtimeTableName, computeTotalDocumentCount(segmentsZKMetadata));
 
+    boolean isPauselessConsumptionEnabled = PauselessConsumptionUtils.isPauselessEnabled(tableConfig);
+
+    if (isPauselessConsumptionEnabled) {
+      _llcRealtimeSegmentManager.reIngestSegmentsWithErrorState(tableConfig.getTableName());
+    }
+
     // Check missing segments and upload them to the deep store
     if (_llcRealtimeSegmentManager.isDeepStoreLLCSegmentUploadRetryEnabled()) {
       _llcRealtimeSegmentManager.uploadToDeepStoreIfMissing(tableConfig, segmentsZKMetadata);

pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java

@@ -56,6 +56,7 @@
 import org.apache.pinot.common.metrics.ServerMeter;
 import org.apache.pinot.common.metrics.ServerMetrics;
 import org.apache.pinot.common.restlet.resources.SegmentErrorInfo;
+import org.apache.pinot.common.utils.PauselessConsumptionUtils;
 import org.apache.pinot.common.utils.TarCompressionUtils;
 import org.apache.pinot.common.utils.config.TierConfigUtils;
 import org.apache.pinot.common.utils.fetcher.SegmentFetcherFactory;
@@ -389,15 +390,21 @@ protected void replaceSegmentIfCrcMismatch(SegmentDataManager segmentDataManager
       IndexLoadingConfig indexLoadingConfig)
       throws Exception {
     String segmentName = segmentDataManager.getSegmentName();
-    Preconditions.checkState(segmentDataManager instanceof ImmutableSegmentDataManager,
-        "Cannot replace CONSUMING segment: %s in table: %s", segmentName, _tableNameWithType);
-    SegmentMetadata localMetadata = segmentDataManager.getSegment().getSegmentMetadata();
-    if (hasSameCRC(zkMetadata, localMetadata)) {
-      _logger.info("Segment: {} has CRC: {} same as before, not replacing it", segmentName, localMetadata.getCrc());
-      return;
+    TableConfig tableConfig = indexLoadingConfig.getTableConfig();
+    // For pauseless tables, we should replace the segment if download url is missing even if crc is same
+    // Without this the reingestion of ERROR segments in pauseless tables fails
+    // as the segment data manager is still an instance of RealtimeSegmentDataManager
+    if (!PauselessConsumptionUtils.isPauselessEnabled(tableConfig)) {
+      Preconditions.checkState(segmentDataManager instanceof ImmutableSegmentDataManager,
+          "Cannot replace CONSUMING segment: %s in table: %s", segmentName, _tableNameWithType);
+      SegmentMetadata localMetadata = segmentDataManager.getSegment().getSegmentMetadata();
+      if (hasSameCRC(zkMetadata, localMetadata)) {
+        _logger.info("Segment: {} has CRC: {} same as before, not replacing it", segmentName, localMetadata.getCrc());
+        return;
+      }
+      _logger.info("Replacing segment: {} because its CRC has changed from: {} to: {}", segmentName,
+          localMetadata.getCrc(), zkMetadata.getCrc());
     }
-    _logger.info("Replacing segment: {} because its CRC has changed from: {} to: {}", segmentName,
-        localMetadata.getCrc(), zkMetadata.getCrc());
     downloadAndLoadSegment(zkMetadata, indexLoadingConfig);
     _logger.info("Replaced segment: {} with new CRC: {}", segmentName, zkMetadata.getCrc());
   }

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java

@@ -1056,7 +1056,7 @@ AtomicBoolean getAcquiredConsumerSemaphore() {
   }
 
   @VisibleForTesting
-  SegmentBuildDescriptor buildSegmentInternal(boolean forCommit) {
+  protected SegmentBuildDescriptor buildSegmentInternal(boolean forCommit) {
     // for partial upsert tables, do not release _partitionGroupConsumerSemaphore proactively and rely on offload()
     // to release the semaphore. This ensures new consuming segment is not consuming until the segment replacement is
     // complete.
@@ -1243,7 +1243,7 @@ protected boolean buildSegmentAndReplace()
     return true;
   }
 
-  private void closeStreamConsumers() {
+  protected void closeStreamConsumers() {
     closePartitionGroupConsumer();
     closePartitionMetadataProvider();
     if (_acquiredConsumerSemaphore.compareAndSet(true, false)) {

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java

@@ -549,9 +549,8 @@ private void doAddConsumingSegment(String segmentName)
         _tableDedupMetadataManager != null ? _tableDedupMetadataManager.getOrCreatePartitionManager(partitionGroupId)
             : null;
     RealtimeSegmentDataManager realtimeSegmentDataManager =
-        new RealtimeSegmentDataManager(zkMetadata, tableConfig, this, _indexDir.getAbsolutePath(), indexLoadingConfig,
-            schema, llcSegmentName, semaphore, _serverMetrics, partitionUpsertMetadataManager,
-            partitionDedupMetadataManager, _isTableReadyToConsumeData);
+        createRealtimeSegmentDataManager(zkMetadata, tableConfig, indexLoadingConfig, schema, llcSegmentName, semaphore,
+            partitionUpsertMetadataManager, partitionDedupMetadataManager, _isTableReadyToConsumeData);
     registerSegment(segmentName, realtimeSegmentDataManager, partitionUpsertMetadataManager);
     if (partitionUpsertMetadataManager != null) {
       partitionUpsertMetadataManager.trackNewlyAddedSegment(segmentName);
@@ -638,6 +637,16 @@ private long getDownloadTimeOutMilliseconds(@Nullable TableConfig tableConfig) {
         }).orElse(DEFAULT_SEGMENT_DOWNLOAD_TIMEOUT_MS);
   }
 
+  protected RealtimeSegmentDataManager createRealtimeSegmentDataManager(SegmentZKMetadata zkMetadata,
+      TableConfig tableConfig, IndexLoadingConfig indexLoadingConfig, Schema schema, LLCSegmentName llcSegmentName,
+      Semaphore semaphore, PartitionUpsertMetadataManager partitionUpsertMetadataManager,
+      PartitionDedupMetadataManager partitionDedupMetadataManager, BooleanSupplier isTableReadyToConsumeData)
+      throws AttemptsExceededException, RetriableOperationException {
+    return new RealtimeSegmentDataManager(zkMetadata, tableConfig, this, _indexDir.getAbsolutePath(),
+        indexLoadingConfig, schema, llcSegmentName, semaphore, _serverMetrics, partitionUpsertMetadataManager,
+        partitionDedupMetadataManager, isTableReadyToConsumeData);
+  }
+
   /**
    * Sets the default time value in the schema as the segment creation time if it is invalid. Time column is used to
    * manage the segments, so its values have to be within the valid range.

pinot-core/src/main/java/org/apache/pinot/server/realtime/ServerSegmentCompletionProtocolHandler.java

@@ -100,6 +100,14 @@ public AuthProvider getAuthProvider() {
     return _authProvider;
   }
 
+  public String getProtocol() {
+    return _protocol;
+  }
+
+  public Integer getControllerHttpsPort() {
+    return _controllerHttpsPort;
+  }
+
   public SegmentCompletionProtocol.Response segmentCommitStart(SegmentCompletionProtocol.Request.Params params) {
     SegmentCompletionProtocol.SegmentCommitStartRequest request =
         new SegmentCompletionProtocol.SegmentCommitStartRequest(params);