Auto reset offset during ingestion lag

Author: Xin Gao

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

@@ -796,9 +796,10 @@ private String createNewSegmentMetadata(TableConfig tableConfig, IdealState idea
         //  segment. For tables with pauseless mode enabled, this size is unavailable at this step because the
         //  segment has not yet been built.
 
-        createNewSegmentZKMetadata(tableConfig, streamConfigs.get(0), newLLCSegment, newSegmentCreationTimeMs,
+        createNewSegmentZKMetadataWithOffsetAutoReset(tableConfig, streamConfigs.get(0), newLLCSegment,
+            newSegmentCreationTimeMs,
             committingSegmentDescriptor, committingSegmentZKMetadata, instancePartitions, partitionIds.size(),
-            numReplicas);
+            numReplicas, false);
         newConsumingSegmentName = newLLCSegment.getSegmentName();
         LOGGER.info("Created new segment metadata for segment: {} with status: {}.", newConsumingSegmentName,
             Status.IN_PROGRESS);
@@ -928,10 +929,9 @@ private boolean isPeerURL(String segmentLocation) {
   private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig streamConfig,
       LLCSegmentName newLLCSegmentName, long creationTimeMs, CommittingSegmentDescriptor committingSegmentDescriptor,
       @Nullable SegmentZKMetadata committingSegmentZKMetadata, InstancePartitions instancePartitions, int numPartitions,
-      int numReplicas) {
+      int numReplicas, String startOffset) {
     String realtimeTableName = tableConfig.getTableName();
     String segmentName = newLLCSegmentName.getSegmentName();
-    String startOffset = committingSegmentDescriptor.getNextOffset();
 
     LOGGER.info(
         "Creating segment ZK metadata for new CONSUMING segment: {} with start offset: {} and creation time: {}",
@@ -962,6 +962,72 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st
     persistSegmentZKMetadata(realtimeTableName, newSegmentZKMetadata, -1);
   }
 
+  /**
+   * Creates and persists segment ZK metadata for the new CONSUMING segment.
+   */
+  private void createNewSegmentZKMetadataWithOffsetAutoReset(TableConfig tableConfig, StreamConfig streamConfig,
+      LLCSegmentName newLLCSegmentName, long creationTimeMs, CommittingSegmentDescriptor committingSegmentDescriptor,
+      @Nullable SegmentZKMetadata committingSegmentZKMetadata, InstancePartitions instancePartitions, int numPartitions,
+      int numReplicas, boolean skipAutoReset) {
+    String oldStartOffset = committingSegmentDescriptor.getNextOffset();
+    String startOffset = skipAutoReset ? oldStartOffset : computeStartOffset(
+        oldStartOffset, streamConfig, newLLCSegmentName.getPartitionGroupId());
+    createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegmentName, creationTimeMs,
+        committingSegmentDescriptor, committingSegmentZKMetadata, instancePartitions, numPartitions, numReplicas,
+        startOffset);
+  }
+
+  private String computeStartOffset(
+      String nextOffset, StreamConfig streamConfig, int partitionId) {
+    long timeThreshold = streamConfig.getOffsetAutoResetTimeSecThreshold();
+    int offsetThreshold = streamConfig.getOffsetAutoResetOffsetThreshold();
+    if (timeThreshold <= 0 && offsetThreshold <= 0) {
+      return nextOffset;
+    }
+    String clientId =
+        PinotLLCRealtimeSegmentManager.class.getSimpleName() + "-" + streamConfig.getTableNameWithType() + "-"
+            + streamConfig.getTopicName();
+    StreamConsumerFactory consumerFactory = StreamConsumerFactoryProvider.create(streamConfig);
+    StreamPartitionMsgOffset offsetAtSLA;
+    StreamPartitionMsgOffset latestOffset;
+    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);
+    } 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);
+        return latestOffset.toString();
+      }
+      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);
+    }
+    return nextOffset;
+  }
+
   @Nullable
   private SegmentPartitionMetadata getPartitionMetadataFromTableConfig(TableConfig tableConfig, int partitionId,
       int numPartitionGroups) {
@@ -1585,8 +1651,10 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List<StreamConf
               CommittingSegmentDescriptor committingSegmentDescriptor =
                   new CommittingSegmentDescriptor(latestSegmentName,
                       (offsetFactory.create(latestSegmentZKMetadata.getEndOffset()).toString()), 0);
-              createNewSegmentZKMetadata(tableConfig, streamConfigs.get(0), newLLCSegmentName, currentTimeMs,
-                  committingSegmentDescriptor, latestSegmentZKMetadata, instancePartitions, numPartitions, numReplicas);
+              createNewSegmentZKMetadataWithOffsetAutoReset(tableConfig, streamConfigs.get(0), newLLCSegmentName,
+                  currentTimeMs,
+                  committingSegmentDescriptor, latestSegmentZKMetadata, instancePartitions, numPartitions, numReplicas,
+                  false);
               updateInstanceStatesForNewConsumingSegment(instanceStatesMap, latestSegmentName, newSegmentName,
                   segmentAssignment, instancePartitionsMap);
             } else { // partition group reached end of life
@@ -1711,8 +1779,8 @@ private void createNewConsumingSegment(TableConfig tableConfig, StreamConfig str
     LLCSegmentName newLLCSegmentName = getNextLLCSegmentName(latestLLCSegmentName, currentTimeMs);
     CommittingSegmentDescriptor committingSegmentDescriptor =
         new CommittingSegmentDescriptor(latestSegmentZKMetadata.getSegmentName(), startOffset.toString(), 0);
-    createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegmentName, currentTimeMs, committingSegmentDescriptor,
-        latestSegmentZKMetadata, instancePartitions, numPartitions, numReplicas);
+    createNewSegmentZKMetadataWithOffsetAutoReset(tableConfig, streamConfig, newLLCSegmentName, currentTimeMs,
+        committingSegmentDescriptor, latestSegmentZKMetadata, instancePartitions, numPartitions, numReplicas, false);
     String newSegmentName = newLLCSegmentName.getSegmentName();
     updateInstanceStatesForNewConsumingSegment(instanceStatesMap, null, newSegmentName, segmentAssignment,
         instancePartitionsMap);
@@ -1790,8 +1858,8 @@ private String setupNewPartitionGroup(TableConfig tableConfig, StreamConfig stre
     String newSegmentName = newLLCSegmentName.getSegmentName();
 
     CommittingSegmentDescriptor committingSegmentDescriptor = new CommittingSegmentDescriptor(null, startOffset, 0);
-    createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegmentName, creationTimeMs,
-        committingSegmentDescriptor, null, instancePartitions, numPartitions, numReplicas);
+    createNewSegmentZKMetadataWithOffsetAutoReset(tableConfig, streamConfig, newLLCSegmentName, creationTimeMs,
+        committingSegmentDescriptor, null, instancePartitions, numPartitions, numReplicas, true);
 
     return newSegmentName;
   }

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

@@ -77,9 +77,14 @@
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.filesystem.PinotFSFactory;
 import org.apache.pinot.spi.stream.LongMsgOffset;
+import org.apache.pinot.spi.stream.OffsetCriteria;
 import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
 import org.apache.pinot.spi.stream.PartitionGroupMetadata;
 import org.apache.pinot.spi.stream.StreamConfig;
+import org.apache.pinot.spi.stream.StreamConfigProperties;
+import org.apache.pinot.spi.stream.StreamConsumerFactory;
+import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
+import org.apache.pinot.spi.stream.StreamMetadataProvider;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Helix;
 import org.apache.pinot.spi.utils.CommonConstants.Helix.Instance;
@@ -91,6 +96,7 @@
 import org.apache.pinot.util.TestUtils;
 import org.apache.zookeeper.data.Stat;
 import org.joda.time.Interval;
+import org.mockito.MockedStatic;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
@@ -122,6 +128,7 @@ public class PinotLLCRealtimeSegmentManagerTest {
   static final String CRC = Long.toString(RANDOM.nextLong() & 0xFFFFFFFFL);
   static final SegmentVersion SEGMENT_VERSION = RANDOM.nextBoolean() ? SegmentVersion.v1 : SegmentVersion.v3;
   static final int NUM_DOCS = RANDOM.nextInt(Integer.MAX_VALUE) + 1;
+  static final long LATEST_OFFSET = PARTITION_OFFSET.getOffset() * 2 + NUM_DOCS;
   static final int SEGMENT_SIZE_IN_BYTES = 100000000;
   @AfterClass
   public void tearDown()
@@ -325,6 +332,73 @@ public void testCommitSegment() {
     assertNull(consumingSegmentZKMetadata);
   }
 
+  @Test
+  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"
+    );
+    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);
+
+    try (MockedStatic<StreamConsumerFactoryProvider> mockedStaticProvider = mockStatic(
+        StreamConsumerFactoryProvider.class)) {
+
+      mockedStaticProvider.when(() -> StreamConsumerFactoryProvider.create(segmentManager._streamConfigs.get(0)))
+          .thenReturn(mockConsumerFactory);
+
+      // Commit a segment for partition group 0
+      String committingSegment = new LLCSegmentName(RAW_TABLE_NAME, 0, 0, CURRENT_TIME_MS).getSegmentName();
+      String endOffset = new LongMsgOffset(PARTITION_OFFSET.getOffset() + NUM_DOCS).toString();
+      CommittingSegmentDescriptor committingSegmentDescriptor =
+          new CommittingSegmentDescriptor(committingSegment, endOffset, SEGMENT_SIZE_IN_BYTES);
+      committingSegmentDescriptor.setSegmentMetadata(mockSegmentMetadata());
+      segmentManager.commitSegmentMetadata(REALTIME_TABLE_NAME, committingSegmentDescriptor);
+
+      // Verify instance states for committed segment and new consuming segment
+      Map<String, String> committedSegmentInstanceStateMap = instanceStatesMap.get(committingSegment);
+      assertNotNull(committedSegmentInstanceStateMap);
+      assertEquals(new HashSet<>(committedSegmentInstanceStateMap.values()),
+          Collections.singleton(SegmentStateModel.ONLINE));
+
+      String consumingSegment = new LLCSegmentName(RAW_TABLE_NAME, 0, 1, CURRENT_TIME_MS).getSegmentName();
+      Map<String, String> consumingSegmentInstanceStateMap = instanceStatesMap.get(consumingSegment);
+      assertNotNull(consumingSegmentInstanceStateMap);
+      assertEquals(new HashSet<>(consumingSegmentInstanceStateMap.values()),
+          Collections.singleton(SegmentStateModel.CONSUMING));
+
+      // Verify segment ZK metadata for committed segment and new consuming segment
+      SegmentZKMetadata committedSegmentZKMetadata = segmentManager._segmentZKMetadataMap.get(committingSegment);
+      assertEquals(committedSegmentZKMetadata.getStatus(), Status.DONE);
+      assertEquals(committedSegmentZKMetadata.getStartOffset(), PARTITION_OFFSET.toString());
+      assertEquals(committedSegmentZKMetadata.getEndOffset(), endOffset);
+      assertEquals(committedSegmentZKMetadata.getCreationTime(), CURRENT_TIME_MS);
+      assertEquals(committedSegmentZKMetadata.getCrc(), Long.parseLong(CRC));
+      assertEquals(committedSegmentZKMetadata.getIndexVersion(), SEGMENT_VERSION.name());
+      assertEquals(committedSegmentZKMetadata.getTotalDocs(), NUM_DOCS);
+      assertEquals(committedSegmentZKMetadata.getSizeInBytes(), SEGMENT_SIZE_IN_BYTES);
+
+      SegmentZKMetadata consumingSegmentZKMetadata = segmentManager._segmentZKMetadataMap.get(consumingSegment);
+      assertEquals(consumingSegmentZKMetadata.getStatus(), Status.IN_PROGRESS);
+      assertEquals(consumingSegmentZKMetadata.getStartOffset(), String.valueOf(LATEST_OFFSET));
+      assertEquals(committedSegmentZKMetadata.getCreationTime(), CURRENT_TIME_MS);
+    }
+  }
+
   /**
    * Test cases for the scenario where stream partitions increase, and the validation manager is attempting to create
    * segments for new partitions. This test assumes that all other factors remain the same (no error conditions or
@@ -1725,6 +1799,13 @@ void makeTableConfig() {
       _streamConfigs = IngestionConfigUtils.getStreamConfigs(_tableConfig);
     }
 
+    void makeTableConfig(Map<String, String> streamConfigMap) {
+      _tableConfig =
+          new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(_numReplicas)
+              .setStreamConfigs(streamConfigMap).build();
+      _streamConfigs = IngestionConfigUtils.getStreamConfigs(_tableConfig);
+    }
+
     void makeConsumingInstancePartitions() {
       List<String> instances = new ArrayList<>(_numInstances);
       for (int i = 0; i < _numInstances; i++) {

pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaStreamMetadataProvider.java

@@ -187,6 +187,12 @@ public List<TopicMetadata> getTopics() {
     }
   }
 
+  @Override
+  public StreamPartitionMsgOffset getOffsetAtTimestamp(int partitionId, long timestampMillis) {
+    return new LongMsgOffset(
+        _consumer.offsetsForTimes(Map.of(_topicPartition, timestampMillis)).get(_topicPartition).offset());
+  }
+
   public static class KafkaTopicMetadata implements TopicMetadata {
     private String _name;
 

pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java

@@ -199,6 +199,13 @@ public KafkaTopicMetadata setName(String name) {
       return this;
     }
   }
+
+  @Override
+  public StreamPartitionMsgOffset getOffsetAtTimestamp(int partitionId, long timestampMillis) {
+    return new LongMsgOffset(
+            _consumer.offsetsForTimes(Map.of(_topicPartition, timestampMillis)).get(_topicPartition).offset());
+  }
+
   @Override
   public void close()
       throws IOException {