Supports Force Committing Segments in Batches

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotRealtimeTableResource.java

@@ -29,7 +29,6 @@
 import io.swagger.annotations.SecurityDefinition;
 import io.swagger.annotations.SwaggerDefinition;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
@@ -170,19 +169,28 @@ public Map<String, String> forceCommit(
       @ApiParam(value = "Comma separated list of partition group IDs to be committed") @QueryParam("partitions")
       String partitionGroupIds,
       @ApiParam(value = "Comma separated list of consuming segments to be committed") @QueryParam("segments")
-      String consumingSegments, @Context HttpHeaders headers) {
+      String consumingSegments,
+      @ApiParam(value = "Max number of consuming segments to commit at once") @QueryParam("batchSize")
+      Integer batchSize, @Context HttpHeaders headers) {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     if (partitionGroupIds != null && consumingSegments != null) {
       throw new ControllerApplicationException(LOGGER, "Cannot specify both partitions and segments to commit",
           Response.Status.BAD_REQUEST);
     }
+    if (batchSize == null) {
+      batchSize = Integer.MAX_VALUE;
+    } else if (batchSize <= 0) {
+      throw new ControllerApplicationException(LOGGER, "Batch size should be greater than zero",
+          Response.Status.BAD_REQUEST);
+    }
     long startTimeMs = System.currentTimeMillis();
     String tableNameWithType = TableNameBuilder.REALTIME.tableNameWithType(tableName);
     validateTable(tableNameWithType);
     Map<String, String> response = new HashMap<>();
     try {
       Set<String> consumingSegmentsForceCommitted =
-          _pinotLLCRealtimeSegmentManager.forceCommit(tableNameWithType, partitionGroupIds, consumingSegments);
+          _pinotLLCRealtimeSegmentManager.forceCommit(tableNameWithType, partitionGroupIds, consumingSegments,
+              batchSize);
       response.put("forceCommitStatus", "SUCCESS");
       try {
         String jobId = UUID.randomUUID().toString();
@@ -222,15 +230,26 @@ public JsonNode getForceCommitJobStatus(
     String tableNameWithType = controllerJobZKMetadata.get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE);
     Set<String> consumingSegmentCommitted = JsonUtils.stringToObject(
         controllerJobZKMetadata.get(CommonConstants.ControllerJob.CONSUMING_SEGMENTS_FORCE_COMMITTED_LIST), Set.class);
-    Set<String> onlineSegmentsForTable =
-        _pinotHelixResourceManager.getOnlineSegmentsFromIdealState(tableNameWithType, false);
 
-    Set<String> segmentsYetToBeCommitted = new HashSet<>();
-    consumingSegmentCommitted.forEach(segmentName -> {
-      if (!onlineSegmentsForTable.contains(segmentName)) {
-        segmentsYetToBeCommitted.add(segmentName);
-      }
-    });
+    Set<String> segmentsToCheck;
+    String segmentsPendingToBeComittedString =
+        controllerJobZKMetadata.get(CommonConstants.ControllerJob.CONSUMING_SEGMENTS_YET_TO_BE_COMMITTED_LIST);
+
+    if (segmentsPendingToBeComittedString != null) {
+      segmentsToCheck = JsonUtils.stringToObject(segmentsPendingToBeComittedString, Set.class);
+    } else {
+      segmentsToCheck = consumingSegmentCommitted;
+    }
+
+    Set<String> segmentsYetToBeCommitted =
+        _pinotLLCRealtimeSegmentManager.getSegmentsYetToBeCommitted(tableNameWithType, segmentsToCheck);
+
+    if (segmentsYetToBeCommitted.size() < segmentsToCheck.size()) {
+      controllerJobZKMetadata.put(CommonConstants.ControllerJob.CONSUMING_SEGMENTS_YET_TO_BE_COMMITTED_LIST,
+          JsonUtils.objectToString(segmentsYetToBeCommitted));
+      _pinotHelixResourceManager.addControllerJobToZK(forceCommitJobId, controllerJobZKMetadata,
+          ControllerJobType.FORCE_COMMIT, prev -> true);
+    }
 
     Map<String, Object> result = new HashMap<>(controllerJobZKMetadata);
     result.put("segmentsYetToBeCommitted", segmentsYetToBeCommitted);

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

@@ -31,8 +31,10 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
 import java.util.Random;
 import java.util.Set;
 import java.util.TreeSet;
@@ -117,7 +119,10 @@
 import org.apache.pinot.spi.utils.StringUtil;
 import org.apache.pinot.spi.utils.TimeUtils;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
+import org.apache.pinot.spi.utils.retry.RetriableOperationException;
 import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
 import org.apache.zookeeper.data.Stat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -152,6 +157,9 @@ public class PinotLLCRealtimeSegmentManager {
 
   // Max time to wait for all LLC segments to complete committing their metadata while stopping the controller.
   private static final long MAX_LLC_SEGMENT_METADATA_COMMIT_TIME_MILLIS = 30_000L;
+  private static final int FORCE_COMMIT_STATUS_CHECK_INTERVAL_MS = 15000;
+  private static final RetryPolicy DEFAULT_RETRY_POLICY =
+      RetryPolicies.fixedDelayRetryPolicy(10, FORCE_COMMIT_STATUS_CHECK_INTERVAL_MS);
 
   // TODO: make this configurable with default set to 10
   /**
@@ -1848,15 +1856,114 @@ private boolean isTmpAndCanDelete(String filePath, Set<String> downloadUrls, Pin
    * @return the set of consuming segments for which commit was initiated
    */
   public Set<String> forceCommit(String tableNameWithType, @Nullable String partitionGroupIdsToCommit,
-      @Nullable String segmentsToCommit) {
+      @Nullable String segmentsToCommit, int batchSize) {
     IdealState idealState = getIdealState(tableNameWithType);
     Set<String> allConsumingSegments = findConsumingSegments(idealState);
     Set<String> targetConsumingSegments = filterSegmentsToCommit(allConsumingSegments, partitionGroupIdsToCommit,
         segmentsToCommit);
-    sendForceCommitMessageToServers(tableNameWithType, targetConsumingSegments);
+
+    List<Set<String>> segmentBatchList = getSegmentBatchList(idealState, targetConsumingSegments, batchSize);
+    ExecutorService executorService = Executors.newFixedThreadPool(1);
+
+    try {
+      for (Set<String> segmentBatchToCommit : segmentBatchList) {
+        executorService.submit(() -> executeBatch(tableNameWithType, segmentBatchToCommit));
+      }
+    } finally {
+      executorService.shutdown();
+    }
+
     return targetConsumingSegments;
   }
 
+  private void executeBatch(String tableNameWithType, Set<String> segmentBatchToCommit) {
+    sendForceCommitMessageToServers(tableNameWithType, segmentBatchToCommit);
+
+    try {
+      Thread.sleep(FORCE_COMMIT_STATUS_CHECK_INTERVAL_MS);
+    } catch (InterruptedException ignored) {
+    }
+
+    int attemptCount = 0;
+    final Set<String>[] segmentsYetToBeCommitted = new Set[]{new HashSet<>()};
+    try {
+      attemptCount = DEFAULT_RETRY_POLICY.attempt(() -> {
+        segmentsYetToBeCommitted[0] = getSegmentsYetToBeCommitted(tableNameWithType, segmentBatchToCommit);
+        return segmentsYetToBeCommitted[0].isEmpty();
+      });
+    } catch (AttemptsExceededException | RetriableOperationException e) {
+      String errorMsg = String.format(
+          "Exception occurred while executing the forceCommit batch of segments: %s, attempt count: %d, "
+              + "segmentsYetToBeCommitted: %s",
+          segmentBatchToCommit, attemptCount, segmentsYetToBeCommitted[0]);
+      LOGGER.error(errorMsg, e);
+      throw new RuntimeException(e);
+    }
+  }
+
+  @VisibleForTesting
+  List<Set<String>> getSegmentBatchList(IdealState idealState, Set<String> targetConsumingSegments,
+      int batchSize) {
+    Map<String, Queue<String>> instanceToConsumingSegments =
+        getInstanceToConsumingSegments(idealState, targetConsumingSegments);
+
+    List<Set<String>> segmentBatchList = new ArrayList<>();
+    Set<String> currentBatch = new HashSet<>();
+    Set<String> segmentsAdded = new HashSet<>();
+    boolean segmentsRemaining = true;
+
+    while (segmentsRemaining) {
+      segmentsRemaining = false;
+      // pick segments in round-robin fashion to parallelize
+      // forceCommit across max servers
+      for (Queue<String> queue : instanceToConsumingSegments.values()) {
+        if (!queue.isEmpty()) {
+          segmentsRemaining = true;
+          String segmentName = queue.poll();
+          // there might be a segment replica hosted on
+          // another instance added before
+          if (segmentsAdded.contains(segmentName)) {
+            continue;
+          }
+          currentBatch.add(segmentName);
+          segmentsAdded.add(segmentName);
+          if (currentBatch.size() == batchSize) {
+            segmentBatchList.add(currentBatch);
+            currentBatch = new HashSet<>();
+          }
+        }
+      }
+    }
+
+    if (!currentBatch.isEmpty()) {
+      segmentBatchList.add(currentBatch);
+    }
+    return segmentBatchList;
+  }
+
+  @VisibleForTesting
+  Map<String, Queue<String>> getInstanceToConsumingSegments(IdealState idealState,
+      Set<String> targetConsumingSegments) {
+    Map<String, Queue<String>> instanceToConsumingSegments = new HashMap<>();
+
+    Map<String, Map<String, String>> segmentNameToInstanceToStateMap = idealState.getRecord().getMapFields();
+    for (String segmentName : segmentNameToInstanceToStateMap.keySet()) {
+      if (!targetConsumingSegments.contains(segmentName)) {
+        continue;
+      }
+      Map<String, String> instanceToStateMap = segmentNameToInstanceToStateMap.get(segmentName);
+      for (String instance : instanceToStateMap.keySet()) {
+        String state = instanceToStateMap.get(instance);
+        if (state.equals(SegmentStateModel.CONSUMING)) {
+          instanceToConsumingSegments.putIfAbsent(instance, new LinkedList<>());
+          instanceToConsumingSegments.get(instance).add(segmentName);
+        }
+      }
+    }
+
+    return instanceToConsumingSegments;
+  }
+
   /**
    * Among all consuming segments, filter the ones that are in the given partitions or segments.
    */
@@ -2009,4 +2116,20 @@ String moveSegmentFile(String rawTableName, String segmentName, String segmentLo
   URI createSegmentPath(String rawTableName, String segmentName) {
     return URIUtils.getUri(_controllerConf.getDataDir(), rawTableName, URIUtils.encode(segmentName));
   }
+
+  public Set<String> getSegmentsYetToBeCommitted(String tableNameWithType, Set<String> segmentsToCheck) {
+    Set<String> segmentsYetToBeCommitted = new HashSet<>();
+    for (String segmentName: segmentsToCheck) {
+      SegmentZKMetadata segmentZKMetadata =
+          _helixResourceManager.getSegmentZKMetadata(tableNameWithType, segmentName);
+      if (segmentZKMetadata == null) {
+        // Segment is deleted. No need to track this segment among segments yetToBeCommitted.
+        continue;
+      }
+      if (!(segmentZKMetadata.getStatus().equals(Status.DONE))) {
+        segmentsYetToBeCommitted.add(segmentName);
+      }
+    }
+    return segmentsYetToBeCommitted;
+  }
 }