Resolve comments

Author: Xin Gao

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

@@ -977,11 +977,16 @@ private void createNewSegmentZKMetadataWithOffsetAutoReset(TableConfig tableConf
         startOffset);
   }
 
-  private String computeStartOffset(
-      String nextOffset, StreamConfig streamConfig, int partitionId) {
+  private String computeStartOffset(String nextOffset, StreamConfig streamConfig, int partitionId) {
+    if (!streamConfig.isEnableOffsetAutoReset()) {
+      return nextOffset;
+    }
     long timeThreshold = streamConfig.getOffsetAutoResetTimeSecThreshold();
     int offsetThreshold = streamConfig.getOffsetAutoResetOffsetThreshold();
     if (timeThreshold <= 0 && offsetThreshold <= 0) {
+      LOGGER.warn("Invalid offset auto reset configuration for table: {}, topic: {}. "
+              + "timeThreshold: {}, offsetThreshold: {}",
+          streamConfig.getTableNameWithType(), streamConfig.getTopicName(), timeThreshold, offsetThreshold);
       return nextOffset;
     }
     String clientId =
@@ -990,40 +995,38 @@ private String computeStartOffset(
     StreamConsumerFactory consumerFactory = StreamConsumerFactoryProvider.create(streamConfig);
     StreamPartitionMsgOffset offsetAtSLA;
     StreamPartitionMsgOffset latestOffset;
-    try (StreamMetadataProvider metadataProvider = consumerFactory.createPartitionMetadataProvider(
-        clientId, partitionId)) {
+    try (StreamMetadataProvider metadataProvider = consumerFactory.createPartitionMetadataProvider(clientId,
+        partitionId)) {
       // Fetching timestamp from an offset is an expensive operation which requires reading the data,
       // while fetching offset from timestamp is lightweight and only needs to read metadata.
       // Hence, instead of checking if latestOffset's time - nextOffset's time < SLA, we would check
       // (CurrentTime - SLA)'s offset > nextOffset.
       // TODO: it is relying on System.currentTimeMillis() which might be affected by time drift. If we are able to
       // get nextOffset's time, we should instead check (nextOffset's time + SLA)'s offset < latestOffset
       latestOffset = metadataProvider.fetchStreamPartitionOffset(OffsetCriteria.LARGEST_OFFSET_CRITERIA, 5000);
-      LOGGER.info("Latest offset of topic {} and partition {} is {}",
-          streamConfig.getTopicName(), partitionId, latestOffset);
-      offsetAtSLA = metadataProvider.getOffsetAtTimestamp(
-          partitionId, System.currentTimeMillis() - timeThreshold * 1000);
-      LOGGER.info("Offset at SLA of topic {} and partition {} is {}",
-          streamConfig.getTopicName(), partitionId, offsetAtSLA);
+      LOGGER.info("Latest offset of topic {} and partition {} is {}", streamConfig.getTopicName(), partitionId,
+          latestOffset);
+      offsetAtSLA =
+          metadataProvider.getOffsetAtTimestamp(partitionId, System.currentTimeMillis() - timeThreshold * 1000);
+      LOGGER.info("Offset at SLA of topic {} and partition {} is {}", streamConfig.getTopicName(), partitionId,
+          offsetAtSLA);
     } catch (Exception e) {
       LOGGER.warn("Not able to fetch the offset metadata, skip auto resetting offsets", e);
       return nextOffset;
     }
     try {
-      if (timeThreshold > 0 && offsetAtSLA != null
-          && Long.valueOf(offsetAtSLA.toString()) > Long.valueOf(nextOffset)) {
-        LOGGER.info("Auto reset offset from {} to {} on partition {} because time threshold reached",
-            nextOffset, latestOffset, partitionId);
+      if (timeThreshold > 0 && offsetAtSLA != null && Long.valueOf(offsetAtSLA.toString()) > Long.valueOf(nextOffset)) {
+        LOGGER.info("Auto reset offset from {} to {} on partition {} because time threshold reached", nextOffset,
+            latestOffset, partitionId);
         return latestOffset.toString();
       }
-      if (offsetThreshold > 0
-          && Long.valueOf(latestOffset.toString()) - Long.valueOf(nextOffset) > offsetThreshold) {
+      if (offsetThreshold > 0 && Long.valueOf(latestOffset.toString()) - Long.valueOf(nextOffset) > offsetThreshold) {
         LOGGER.info("Auto reset offset from {} to {} on partition {} because number of offsets threshold reached",
             nextOffset, latestOffset, partitionId);
         return latestOffset.toString();
       }
     } catch (Exception e) {
-      LOGGER.warn("Not able to convert the offset to LONG type, skip auto resetting offsets", e);
+      LOGGER.warn("Not able to compare the offsets, skip auto resetting offsets", e);
     }
     return nextOffset;
   }

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

@@ -333,27 +333,25 @@ public void testCommitSegment() {
   }
 
   @Test
-  public void testCommitSegmentWithOffsetAutoReset() throws Exception {
+  public void testCommitSegmentWithOffsetAutoReset()
+      throws Exception {
     // Set up a new table with 2 replicas, 5 instances, 4 partition
     PinotHelixResourceManager mockHelixResourceManager = mock(PinotHelixResourceManager.class);
     FakePinotLLCRealtimeSegmentManager segmentManager =
         new FakePinotLLCRealtimeSegmentManager(mockHelixResourceManager);
     setUpNewTable(segmentManager, 2, 5, 4);
     Map<String, Map<String, String>> instanceStatesMap = segmentManager._idealState.getRecord().getMapFields();
     Map<String, String> streamConfigMap = IngestionConfigUtils.getStreamConfigMaps(segmentManager._tableConfig).get(0);
-    streamConfigMap.put(
-        StreamConfigProperties.OFFSET_AUTO_RESET_OFFSET_THRESHOLD_KEY, "100"
-    );
+    streamConfigMap.put(StreamConfigProperties.ENABLE_OFFSET_AUTO_RESET, String.valueOf(true));
+    streamConfigMap.put(StreamConfigProperties.OFFSET_AUTO_RESET_OFFSET_THRESHOLD_KEY, "100");
     segmentManager.makeTableConfig(streamConfigMap);
 
     StreamConsumerFactory mockConsumerFactory = mock(StreamConsumerFactory.class);
     StreamMetadataProvider mockMetadataProvider = mock(StreamMetadataProvider.class);
-    when(mockConsumerFactory.createPartitionMetadataProvider(anyString(), anyInt()))
-        .thenReturn(mockMetadataProvider);
-    when(mockMetadataProvider.fetchStreamPartitionOffset(eq(OffsetCriteria.LARGEST_OFFSET_CRITERIA), anyLong()))
-        .thenReturn(new LongMsgOffset(LATEST_OFFSET));
-    when(mockMetadataProvider.getOffsetAtTimestamp(eq(0), anyLong()))
-        .thenReturn(PARTITION_OFFSET);
+    when(mockConsumerFactory.createPartitionMetadataProvider(anyString(), anyInt())).thenReturn(mockMetadataProvider);
+    when(mockMetadataProvider.fetchStreamPartitionOffset(eq(OffsetCriteria.LARGEST_OFFSET_CRITERIA),
+        anyLong())).thenReturn(new LongMsgOffset(LATEST_OFFSET));
+    when(mockMetadataProvider.getOffsetAtTimestamp(eq(0), anyLong())).thenReturn(PARTITION_OFFSET);
 
     try (MockedStatic<StreamConsumerFactoryProvider> mockedStaticProvider = mockStatic(
         StreamConsumerFactoryProvider.class)) {

pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamConfig.java

@@ -75,6 +75,7 @@ public class StreamConfig {
 
   private final double _topicConsumptionRateLimit;
 
+  private final boolean _enableOffsetAutoReset;
   private final int _offsetAutoResetOffsetThreshold;
   private final long _offsetAutoResetTimeSecThreshold;
 
@@ -202,6 +203,7 @@ public StreamConfig(String tableNameWithType, Map<String, String> streamConfigMa
     String rate = streamConfigMap.get(StreamConfigProperties.TOPIC_CONSUMPTION_RATE_LIMIT);
     _topicConsumptionRateLimit = rate != null ? Double.parseDouble(rate) : CONSUMPTION_RATE_LIMIT_NOT_SPECIFIED;
 
+    _enableOffsetAutoReset = Boolean.parseBoolean(streamConfigMap.get(StreamConfigProperties.ENABLE_OFFSET_AUTO_RESET));
     String offsetThreshold = streamConfigMap.get(StreamConfigProperties.OFFSET_AUTO_RESET_OFFSET_THRESHOLD_KEY);
     _offsetAutoResetOffsetThreshold = offsetThreshold != null ? Integer.valueOf(offsetThreshold) : -1;
     String timeSecThreshold = streamConfigMap.get(StreamConfigProperties.OFFSET_AUTO_RESET_TIMESEC_THRESHOLD_KEY);
@@ -391,6 +393,10 @@ public Optional<Double> getTopicConsumptionRateLimit() {
         : Optional.of(_topicConsumptionRateLimit);
   }
 
+  public boolean isEnableOffsetAutoReset() {
+    return _enableOffsetAutoReset;
+  }
+
   public int getOffsetAutoResetOffsetThreshold() {
     return _offsetAutoResetOffsetThreshold;
   }
@@ -447,6 +453,7 @@ public boolean equals(Object o) {
         that._decoderClass) && Objects.equals(_decoderProperties, that._decoderProperties) && Objects.equals(_groupId,
         that._groupId) && Objects.equals(_streamConfigMap, that._streamConfigMap) && Objects.equals(_offsetCriteria,
         that._offsetCriteria) && Objects.equals(_flushThresholdVarianceFraction, that._flushThresholdVarianceFraction)
+        && _enableOffsetAutoReset == that._enableOffsetAutoReset
         && _offsetAutoResetOffsetThreshold == that._offsetAutoResetOffsetThreshold
         && _offsetAutoResetTimeSecThreshold == that._offsetAutoResetTimeSecThreshold;
   }
@@ -458,6 +465,6 @@ public int hashCode() {
         _flushThresholdSegmentRows, _flushThresholdTimeMillis, _flushThresholdSegmentSizeBytes,
         _flushAutotuneInitialRows, _groupId, _topicConsumptionRateLimit, _streamConfigMap, _offsetCriteria,
         _serverUploadToDeepStore, _flushThresholdVarianceFraction, _offsetAutoResetOffsetThreshold,
-        _offsetAutoResetTimeSecThreshold);
+        _enableOffsetAutoReset, _offsetAutoResetTimeSecThreshold);
   }
 }

pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamConfigProperties.java

@@ -143,6 +143,11 @@ private StreamConfigProperties() {
   public static final String PAUSELESS_SEGMENT_DOWNLOAD_TIMEOUT_SECONDS =
       "realtime.segment.pauseless.download.timeoutSeconds";
 
+  /**
+   * Config used to enable offset auto reset during segment commit.
+   */
+  public static final String ENABLE_OFFSET_AUTO_RESET = "realtime.segment.offsetAutoReset.enable";
+
   /**
    * During segment commit, the new segment startOffset would skip to the latest offset if thisValue is set as positive
    * and (latestStreamOffset - latestIngestedOffset > thisValue)