Supports Force Committing Segments in Batches

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

@@ -169,19 +169,29 @@ 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 (default = Integer.MAX_VALUE)")
+      @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();

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
   /**
@@ -189,6 +197,7 @@ public class PinotLLCRealtimeSegmentManager {
   private final AtomicInteger _numCompletingSegments = new AtomicInteger(0);
   private final ExecutorService _deepStoreUploadExecutor;
   private final Set<String> _deepStoreUploadExecutorPendingSegments;
+  private final ExecutorService _forceCommitExecutorService;
 
   private volatile boolean _isStopping = false;
 
@@ -213,6 +222,7 @@ public PinotLLCRealtimeSegmentManager(PinotHelixResourceManager helixResourceMan
         controllerConf.getDeepStoreRetryUploadParallelism()) : null;
     _deepStoreUploadExecutorPendingSegments =
         _isDeepStoreLLCSegmentUploadRetryEnabled ? ConcurrentHashMap.newKeySet() : null;
+    _forceCommitExecutorService = Executors.newFixedThreadPool(4);
   }
 
   public boolean isDeepStoreLLCSegmentUploadRetryEnabled() {
@@ -309,6 +319,8 @@ public void stop() {
         LOGGER.error("Failed to close fileUploadDownloadClient.");
       }
     }
+
+    _forceCommitExecutorService.shutdown();
   }
 
   /**
@@ -1848,15 +1860,122 @@ 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);
+
+    _forceCommitExecutorService.submit(() -> processBatchesSequentially(segmentBatchList, tableNameWithType));
+
     return targetConsumingSegments;
   }
 
+  private void processBatchesSequentially(List<Set<String>> segmentBatchList, String tableNameWithType) {
+    Set<String> prevBatch = null;
+    for (Set<String> segmentBatchToCommit: segmentBatchList) {
+      if (prevBatch != null) {
+        waitUntilPrevBatchIsComplete(tableNameWithType, prevBatch);
+      }
+      sendForceCommitMessageToServers(tableNameWithType, segmentBatchToCommit);
+      prevBatch = segmentBatchToCommit;
+    }
+  }
+
+  private void waitUntilPrevBatchIsComplete(String tableNameWithType, Set<String> 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.compute(instance, (key, value) -> {
+            if (value == null) {
+              value = new LinkedList<>();
+            }
+            value.add(segmentName);
+            return value;
+          });
+        }
+      }
+    }
+
+    return instanceToConsumingSegments;
+  }
+
   /**
    * Among all consuming segments, filter the ones that are in the given partitions or segments.
    */

pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java

@@ -19,6 +19,8 @@
 package org.apache.pinot.controller.helix.core.realtime;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import java.io.File;
@@ -31,9 +33,12 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
 import java.util.Random;
 import java.util.Set;
+import java.util.SortedSet;
 import java.util.TreeMap;
+import java.util.TreeSet;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -1248,6 +1253,112 @@ public void testGetPartitionIds()
     Assert.assertEquals(partitionIds.size(), 2);
   }
 
+  @Test
+  public void testGetInstanceToConsumingSegments() {
+    PinotHelixResourceManager mockHelixResourceManager = mock(PinotHelixResourceManager.class);
+    FakePinotLLCRealtimeSegmentManager realtimeSegmentManager =
+        new FakePinotLLCRealtimeSegmentManager(mockHelixResourceManager);
+    IdealState idealState = mock(IdealState.class);
+    Map<String, Map<String, String>> map = ImmutableMap.of(
+        "seg0", ImmutableMap.of("i1", "CONSUMING", "i4", "ONLINE"),
+        "seg1", ImmutableMap.of("i2", "CONSUMING"),
+        "seg2", ImmutableMap.of("i3", "CONSUMING", "i2", "OFFLINE"),
+        "seg3", ImmutableMap.of("i4", "CONSUMING", "i2", "CONSUMING", "i3", "CONSUMING"),
+        "seg4", ImmutableMap.of("i5", "CONSUMING", "i1", "CONSUMING", "i3", "CONSUMING")
+    );
+
+    ZNRecord znRecord = mock(ZNRecord.class);
+    when(znRecord.getMapFields()).thenReturn(map);
+    when(idealState.getRecord()).thenReturn(znRecord);
+    Set<String> targetConsumingSegment = new HashSet<>(map.keySet());
+
+    Map<String, Queue<String>> instanceToConsumingSegments =
+        realtimeSegmentManager.getInstanceToConsumingSegments(idealState, targetConsumingSegment);
+    List<String> instanceList = ImmutableList.of("i1", "i2", "i3", "i4", "i5");
+
+    StringBuilder expectedSegNames = new StringBuilder();
+    for (String instanceName : instanceList) {
+      SortedSet<String> sortedSegNames = new TreeSet<>(instanceToConsumingSegments.get(instanceName));
+      expectedSegNames.append(sortedSegNames);
+    }
+    assert expectedSegNames.toString()
+        .equals("[seg0, seg4][seg1, seg3][seg2, seg3, seg4][seg3][seg4]");
+  }
+
+  @Test
+  public void getSegmentBatchList() {
+    PinotHelixResourceManager mockHelixResourceManager = mock(PinotHelixResourceManager.class);
+    FakePinotLLCRealtimeSegmentManager realtimeSegmentManager =
+        new FakePinotLLCRealtimeSegmentManager(mockHelixResourceManager);
+    IdealState idealState = mock(IdealState.class);
+
+    Map<String, Map<String, String>> map = ImmutableMap.of(
+        "seg0", ImmutableMap.of("i1", "CONSUMING", "i4", "ONLINE"),
+        "seg1", ImmutableMap.of("i2", "CONSUMING"),
+        "seg2", ImmutableMap.of("i3", "CONSUMING", "i2", "OFFLINE"),
+        "seg3", ImmutableMap.of("i4", "CONSUMING", "i2", "CONSUMING", "i3", "CONSUMING"),
+        "seg4", ImmutableMap.of("i5", "CONSUMING", "i1", "CONSUMING", "i3", "CONSUMING"),
+        "seg5", ImmutableMap.of("i6", "CONSUMING", "i1", "CONSUMING", "i3", "CONSUMING"),
+        "seg6", ImmutableMap.of("i7", "CONSUMING", "i1", "CONSUMING", "i3", "CONSUMING")
+    );
+
+    ZNRecord znRecord = mock(ZNRecord.class);
+    when(znRecord.getMapFields()).thenReturn(map);
+    when(idealState.getRecord()).thenReturn(znRecord);
+    Set<String> targetConsumingSegment = new HashSet<>(map.keySet());
+
+    List<Set<String>> segmentBatchList =
+        realtimeSegmentManager.getSegmentBatchList(idealState, targetConsumingSegment, 2);
+
+    assert segmentBatchList.size() == 4;
+    Set<String> segmentsAdded = new HashSet<>();
+
+    for (Set<String> segmentBatch : segmentBatchList) {
+      assert segmentBatch.size() <= 2;
+      for (String segmentName : segmentBatch) {
+        assert !segmentsAdded.contains(segmentName);
+        segmentsAdded.add(segmentName);
+      }
+    }
+
+    Random random = new Random();
+    int numOfServers = 1 + random.nextInt(20);
+    int numOfSegments = Math.max(numOfServers, random.nextInt(500));
+    int rf = Math.min(numOfServers, random.nextInt(7));
+    int batchSize = random.nextInt(100);
+
+    map = new HashMap<>();
+    for (int segmentIndex = 0; segmentIndex < numOfSegments; segmentIndex++) {
+      String segmentName = "seg_" + segmentIndex;
+      Map<String, String> instanceToStateMap = new HashMap<>();
+      for (int rfIndex = 0; rfIndex < rf; rfIndex++) {
+        instanceToStateMap.put("i_" + random.nextInt(numOfServers), "CONSUMING");
+      }
+      map.put(segmentName, instanceToStateMap);
+    }
+
+    znRecord = mock(ZNRecord.class);
+    when(znRecord.getMapFields()).thenReturn(map);
+    when(idealState.getRecord()).thenReturn(znRecord);
+    targetConsumingSegment = new HashSet<>(map.keySet());
+
+    segmentBatchList = realtimeSegmentManager.getSegmentBatchList(idealState, targetConsumingSegment, batchSize);
+    int numBatchesExpected = (targetConsumingSegment.size() + batchSize - 1) / batchSize;
+
+    assert numBatchesExpected == segmentBatchList.size();
+
+    segmentsAdded = new HashSet<>();
+    for (Set<String> batch : segmentBatchList) {
+      assert batch.size() <= batchSize;
+      for (String segmentName : batch) {
+        assert !segmentsAdded.contains(segmentName);
+        segmentsAdded.add(segmentName);
+      }
+    }
+
+    assert segmentsAdded.equals(targetConsumingSegment);
+  }
+
   @Test
   public void getSegmentsYetToBeCommitted() {
     PinotHelixResourceManager mockHelixResourceManager = mock(PinotHelixResourceManager.class);

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java

@@ -428,10 +428,22 @@ public void testForceCommit()
       throws Exception {
     Set<String> consumingSegments = getConsumingSegmentsFromIdealState(getTableName() + "_REALTIME");
     String jobId = forceCommit(getTableName());
+    testForceCommitInternal(jobId, consumingSegments, 60000L);
+  }
+
+  @Test
+  public void testForceCommitInBatches()
+      throws Exception {
+    Set<String> consumingSegments = getConsumingSegmentsFromIdealState(getTableName() + "_REALTIME");
+    String jobId = forceCommit(getTableName(), 1);
+    testForceCommitInternal(jobId, consumingSegments, 180000L);
+  }
+
+  private void testForceCommitInternal(String jobId, Set<String> consumingSegments, long timeoutMs) {
     Map<String, String> jobMetadata =
         _helixResourceManager.getControllerJobZKMetadata(jobId, ControllerJobType.FORCE_COMMIT);
     assert jobMetadata != null;
-    assert jobMetadata.get("segmentsForceCommitted") != null;
+    assert jobMetadata.get(CommonConstants.ControllerJob.CONSUMING_SEGMENTS_FORCE_COMMITTED_LIST) != null;
 
     TestUtils.waitForCondition(aVoid -> {
       try {
@@ -446,7 +458,7 @@ public void testForceCommit()
       } catch (Exception e) {
         return false;
       }
-    }, 60000L, "Error verifying force commit operation on table!");
+    }, timeoutMs, "Error verifying force commit operation on table!");
   }
 
   public Set<String> getConsumingSegmentsFromIdealState(String tableNameWithType) {
@@ -492,6 +504,13 @@ private String forceCommit(String tableName)
     return JsonUtils.stringToJsonNode(response).get("forceCommitJobId").asText();
   }
 
+  private String forceCommit(String tableName, int batchSize)
+      throws Exception {
+    String response =
+        sendPostRequest(_controllerRequestURLBuilder.forTableForceCommit(tableName) + "?batchSize=" + batchSize, null);
+    return JsonUtils.stringToJsonNode(response).get("forceCommitJobId").asText();
+  }
+
   @Test
   @Override
   public void testHardcodedServerPartitionedSqlQueries()