Pauseless Consumption

pinot-common/src/main/java/org/apache/pinot/common/protocols/SegmentCompletionProtocol.java

@@ -132,6 +132,7 @@ public enum ControllerResponseStatus {
   public static final String PARAM_MEMORY_USED_BYTES = "memoryUsedBytes";
   public static final String PARAM_SEGMENT_SIZE_BYTES = "segmentSizeBytes";
   public static final String PARAM_REASON = "reason";
+  public static final String PARAM_PAUSELESS_CONSUMPTION_ENABLED = "pauselessConsumptionEnabled";
   // Sent by servers to request additional time to build
   public static final String PARAM_EXTRA_TIME_SEC = "extraTimeSec";
   // Sent by servers to indicate the number of rows read so far
@@ -199,6 +200,7 @@ public String getUrl(String hostPort, String protocol) {
       params.put(PARAM_SEGMENT_NAME, _params.getSegmentName());
       params.put(PARAM_OFFSET, String.valueOf(_params.getOffset()));
       params.put(PARAM_INSTANCE_ID, _params.getInstanceId());
+      params.put(PARAM_PAUSELESS_CONSUMPTION_ENABLED, String.valueOf(_params.getPauselessConsumptionEnabled()));
       if (_params.getReason() != null) {
         params.put(PARAM_REASON, _params.getReason());
       }
@@ -242,6 +244,7 @@ public static class Params {
       private long _memoryUsedBytes;
       private long _segmentSizeBytes;
       private String _streamPartitionMsgOffset;
+      private boolean _pauselessConsumptionEnabled;
 
       public Params() {
         _offset = -1L;
@@ -256,6 +259,7 @@ public Params() {
         _segmentSizeBytes = SEGMENT_SIZE_BYTES_DEFAULT;
         _streamPartitionMsgOffset = null;
         _reason = null;
+        _pauselessConsumptionEnabled = false;
       }
 
       public Params(Params params) {
@@ -271,6 +275,7 @@ public Params(Params params) {
         _segmentSizeBytes = params.getSegmentSizeBytes();
         _streamPartitionMsgOffset = params.getStreamPartitionMsgOffset();
         _reason = params.getReason();
+        _pauselessConsumptionEnabled = params.getPauselessConsumptionEnabled();
       }
 
       @Deprecated
@@ -334,6 +339,11 @@ public Params withStreamPartitionMsgOffset(String offset) {
         return this;
       }
 
+      public Params withPauselessConsumptionEnabled(boolean pauselessConsumptionEnabled) {
+        _pauselessConsumptionEnabled = pauselessConsumptionEnabled;
+        return this;
+      }
+
       public String getSegmentName() {
         return _segmentName;
       }
@@ -383,6 +393,10 @@ public String getStreamPartitionMsgOffset() {
         return _streamPartitionMsgOffset;
       }
 
+      public boolean getPauselessConsumptionEnabled() {
+        return _pauselessConsumptionEnabled;
+      }
+
       public String toString() {
         return "Offset: " + _offset + ",Segment name: " + _segmentName + ",Instance Id: " + _instanceId + ",Reason: "
             + _reason + ",NumRows: " + _numRows + ",BuildTimeMillis: " + _buildTimeMillis + ",WaitTimeMillis: "

pinot-common/src/test/java/org/apache/pinot/common/protocols/SegmentCompletionProtocolTest.java

@@ -112,11 +112,16 @@ public void testRequestURL()
     SegmentCompletionProtocol.SegmentCommitStartRequest segmentCommitStartRequest =
         new SegmentCompletionProtocol.SegmentCommitStartRequest(params);
     String url = segmentCommitStartRequest.getUrl("localhost:8080", "http");
+//    Assert.assertEquals(url,
+//    // CHECKSTYLE:OFF
+//        "http://localhost:8080/segmentCommitStart?extraTimeSec=3000&segmentSizeBytes=5000&reason=%7B%22type%22%3A%22ROW_LIMIT%22%2C%20%22value%22%3A1000%7D&buildTimeMillis=1000&streamPartitionMsgOffset=%7B%22shardId-000000000001%22%3A%2249615238429973311938200772279310862572716999467690098706%22%7D&instance=Server_localhost_8099&waitTimeMillis=2000&offset=-1&name=foo%25%25__0__0__12345Z&location=s3%3A%2F%2Fmy.bucket%2Fsegment&rowCount=6000&memoryUsedBytes=4000");
+//    // CHECKSTYLE:ON
+
+    // Adding the case for test checks
     Assert.assertEquals(url,
-    // CHECKSTYLE:OFF
-        "http://localhost:8080/segmentCommitStart?extraTimeSec=3000&segmentSizeBytes=5000&reason=%7B%22type%22%3A%22ROW_LIMIT%22%2C%20%22value%22%3A1000%7D&buildTimeMillis=1000&streamPartitionMsgOffset=%7B%22shardId-000000000001%22%3A%2249615238429973311938200772279310862572716999467690098706%22%7D&instance=Server_localhost_8099&waitTimeMillis=2000&offset=-1&name=foo%25%25__0__0__12345Z&location=s3%3A%2F%2Fmy.bucket%2Fsegment&rowCount=6000&memoryUsedBytes=4000");
+        // CHECKSTYLE:OFF
+        "http://localhost:8080/segmentCommitStart?segmentSizeBytes=5000&reason=%7B%22type%22%3A%22ROW_LIMIT%22%2C%20%22value%22%3A1000%7D&buildTimeMillis=1000&streamPartitionMsgOffset=%7B%22shardId-000000000001%22%3A%2249615238429973311938200772279310862572716999467690098706%22%7D&instance=Server_localhost_8099&offset=-1&pauselessConsumptionEnabled=false&memoryUsedBytes=4000&extraTimeSec=3000&waitTimeMillis=2000&name=foo%25%25__0__0__12345Z&location=s3%3A%2F%2Fmy.bucket%2Fsegment&rowCount=6000");
     // CHECKSTYLE:ON
-
     paramsMap = Arrays.stream(url.split("\\?")[1].split("&"))
         .collect(Collectors.toMap(e -> e.split("=")[0], e -> e.split("=")[1]));
     Assert.assertEquals(paramsMap.get(SegmentCompletionProtocol.PARAM_SEGMENT_NAME),

pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java

@@ -485,8 +485,7 @@ private void setUpPinotController() {
 
     // Helix resource manager must be started in order to create PinotLLCRealtimeSegmentManager
     LOGGER.info("Starting realtime segment manager");
-    _pinotLLCRealtimeSegmentManager =
-        new PinotLLCRealtimeSegmentManager(_helixResourceManager, _config, _controllerMetrics);
+    _pinotLLCRealtimeSegmentManager = createPinotLLCRealtimeSegmentManager();
     // TODO: Need to put this inside HelixResourceManager when HelixControllerLeadershipManager is removed.
     _helixResourceManager.registerPinotLLCRealtimeSegmentManager(_pinotLLCRealtimeSegmentManager);
     _segmentCompletionManager =
@@ -616,6 +615,10 @@ protected void configure() {
     _serviceStatusCallbackList.add(generateServiceStatusCallback(_helixParticipantManager));
   }
 
+  protected PinotLLCRealtimeSegmentManager createPinotLLCRealtimeSegmentManager() {
+    return new PinotLLCRealtimeSegmentManager(_helixResourceManager, _config, _controllerMetrics);
+  }
+
   /**
    * This method is used to fix table/schema names.
    * TODO: in the next release, maybe 2.0.0, we can remove this method. Meanwhile we can delete the orphan schemas
@@ -861,9 +864,7 @@ protected List<PeriodicTask> setupControllerPeriodicTasks() {
         new OfflineSegmentIntervalChecker(_config, _helixResourceManager, _leadControllerManager,
             new ValidationMetrics(_metricsRegistry), _controllerMetrics);
     periodicTasks.add(_offlineSegmentIntervalChecker);
-    _realtimeSegmentValidationManager =
-        new RealtimeSegmentValidationManager(_config, _helixResourceManager, _leadControllerManager,
-            _pinotLLCRealtimeSegmentManager, _validationMetrics, _controllerMetrics, _storageQuotaChecker);
+    _realtimeSegmentValidationManager = createRealtimeSegmentValidationManager();
     periodicTasks.add(_realtimeSegmentValidationManager);
     _brokerResourceValidationManager =
         new BrokerResourceValidationManager(_config, _helixResourceManager, _leadControllerManager, _controllerMetrics);
@@ -892,6 +893,11 @@ protected List<PeriodicTask> setupControllerPeriodicTasks() {
     return periodicTasks;
   }
 
+  protected RealtimeSegmentValidationManager createRealtimeSegmentValidationManager() {
+    return new RealtimeSegmentValidationManager(_config, _helixResourceManager, _leadControllerManager,
+        _pinotLLCRealtimeSegmentManager, _validationMetrics, _controllerMetrics, _storageQuotaChecker);
+  }
+
   @Override
   public void stop() {
     switch (_controllerMode) {

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

@@ -134,7 +134,8 @@ public String segmentConsumed(@QueryParam(SegmentCompletionProtocol.PARAM_INSTAN
       @QueryParam(SegmentCompletionProtocol.PARAM_STREAM_PARTITION_MSG_OFFSET) String streamPartitionMsgOffset,
       @QueryParam(SegmentCompletionProtocol.PARAM_REASON) String stopReason,
       @QueryParam(SegmentCompletionProtocol.PARAM_MEMORY_USED_BYTES) long memoryUsedBytes,
-      @QueryParam(SegmentCompletionProtocol.PARAM_ROW_COUNT) int numRows) {
+      @QueryParam(SegmentCompletionProtocol.PARAM_ROW_COUNT) int numRows,
+      @QueryParam(SegmentCompletionProtocol.PARAM_PAUSELESS_CONSUMPTION_ENABLED) boolean pauselessConsumptionEnabled) {
 
     if (instanceId == null || segmentName == null || (offset == -1 && streamPartitionMsgOffset == null)) {
       LOGGER.error("Invalid call: offset={}, segmentName={}, instanceId={}, streamPartitionMsgOffset={}", offset,
@@ -143,7 +144,8 @@ public String segmentConsumed(@QueryParam(SegmentCompletionProtocol.PARAM_INSTAN
     }
     SegmentCompletionProtocol.Request.Params requestParams = new SegmentCompletionProtocol.Request.Params();
     requestParams.withInstanceId(instanceId).withSegmentName(segmentName).withReason(stopReason)
-        .withMemoryUsedBytes(memoryUsedBytes).withNumRows(numRows);
+        .withMemoryUsedBytes(memoryUsedBytes).withNumRows(numRows)
+        .withPauselessConsumptionEnabled(pauselessConsumptionEnabled);
     extractOffsetFromParams(requestParams, streamPartitionMsgOffset, offset);
     LOGGER.info("Processing segmentConsumed:{}", requestParams.toString());
 
@@ -191,7 +193,8 @@ public String segmentCommitStart(@QueryParam(SegmentCompletionProtocol.PARAM_INS
       @QueryParam(SegmentCompletionProtocol.PARAM_BUILD_TIME_MILLIS) long buildTimeMillis,
       @QueryParam(SegmentCompletionProtocol.PARAM_WAIT_TIME_MILLIS) long waitTimeMillis,
       @QueryParam(SegmentCompletionProtocol.PARAM_ROW_COUNT) int numRows,
-      @QueryParam(SegmentCompletionProtocol.PARAM_SEGMENT_SIZE_BYTES) long segmentSizeBytes) {
+      @QueryParam(SegmentCompletionProtocol.PARAM_SEGMENT_SIZE_BYTES) long segmentSizeBytes,
+      @QueryParam(SegmentCompletionProtocol.PARAM_PAUSELESS_CONSUMPTION_ENABLED) boolean pauselessConsumptionEnabled) {
 
     if (instanceId == null || segmentName == null || (offset == -1 && streamPartitionMsgOffset == null)) {
       LOGGER.error("Invalid call: offset={}, segmentName={}, instanceId={}, streamPartitionMsgOffset={}", offset,
@@ -203,7 +206,7 @@ public String segmentCommitStart(@QueryParam(SegmentCompletionProtocol.PARAM_INS
     SegmentCompletionProtocol.Request.Params requestParams = new SegmentCompletionProtocol.Request.Params();
     requestParams.withInstanceId(instanceId).withSegmentName(segmentName).withMemoryUsedBytes(memoryUsedBytes)
         .withBuildTimeMillis(buildTimeMillis).withWaitTimeMillis(waitTimeMillis).withNumRows(numRows)
-        .withSegmentSizeBytes(segmentSizeBytes);
+        .withSegmentSizeBytes(segmentSizeBytes).withPauselessConsumptionEnabled(pauselessConsumptionEnabled);
     extractOffsetFromParams(requestParams, streamPartitionMsgOffset, offset);
 
     LOGGER.info("Processing segmentCommitStart:{}", requestParams.toString());
@@ -364,6 +367,7 @@ public String segmentCommitEndWithMetadata(@QueryParam(SegmentCompletionProtocol
       @QueryParam(SegmentCompletionProtocol.PARAM_ROW_COUNT) int numRows,
       @QueryParam(SegmentCompletionProtocol.PARAM_SEGMENT_SIZE_BYTES) long segmentSizeBytes,
       @QueryParam(SegmentCompletionProtocol.PARAM_REASON) String stopReason,
+      @QueryParam(SegmentCompletionProtocol.PARAM_PAUSELESS_CONSUMPTION_ENABLED) boolean pauselessConsumptionEnabled,
       FormDataMultiPart metadataFiles) {
     if (instanceId == null || segmentName == null || segmentLocation == null || metadataFiles == null || (offset == -1
         && streamPartitionMsgOffset == null)) {
@@ -377,7 +381,8 @@ public String segmentCommitEndWithMetadata(@QueryParam(SegmentCompletionProtocol
     SegmentCompletionProtocol.Request.Params requestParams = new SegmentCompletionProtocol.Request.Params();
     requestParams.withInstanceId(instanceId).withSegmentName(segmentName).withSegmentLocation(segmentLocation)
         .withSegmentSizeBytes(segmentSizeBytes).withBuildTimeMillis(buildTimeMillis).withWaitTimeMillis(waitTimeMillis)
-        .withNumRows(numRows).withMemoryUsedBytes(memoryUsedBytes).withReason(stopReason);
+        .withNumRows(numRows).withMemoryUsedBytes(memoryUsedBytes).withReason(stopReason)
+        .withPauselessConsumptionEnabled(pauselessConsumptionEnabled);
     extractOffsetFromParams(requestParams, streamPartitionMsgOffset, offset);
     LOGGER.info("Processing segmentCommitEndWithMetadata:{}", requestParams.toString());