private void setConsumeEndTime(SegmentZKMetadata segmentZKMetadata, long now) {
    long maxConsumeTimeMillis = _streamConfig.getFlushThresholdTimeMillis();
    _consumeEndTime = segmentZKMetadata.getCreationTime() + maxConsumeTimeMillis;

    // When we restart a server, the consuming segments retain their creationTime (derived from segment
    // metadata), but a couple of corner cases can happen:
    // (1) The server was down for a very long time, and the consuming segment is not yet completed.
    // (2) The consuming segment was just about to be completed, but the server went down.
    // In either of these two cases, if a different replica could not complete the segment, it is possible
    // that we get a value for _consumeEndTime that is in the very near future, or even in the past. In such
    // cases, we let some minimum consumption happen before we attempt to complete the segment (unless, of course
    // the max consumption time has been configured to be less than the minimum time we use in this class).
    long minConsumeTimeMillis =
        Math.min(maxConsumeTimeMillis, TimeUnit.MILLISECONDS.convert(MINIMUM_CONSUME_TIME_MINUTES, TimeUnit.MINUTES));
    if (_consumeEndTime - now < minConsumeTimeMillis) {
      _consumeEndTime = now + minConsumeTimeMillis;
    }
  }
@Test
  public void testEndCriteriaChecking()
      throws Exception {
    // test reaching max row limit
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.INITIAL_CONSUMING);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager.setNumRowsIndexed(Fixtures.MAX_ROWS_IN_SEGMENT - 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager.setNumRowsIndexed(Fixtures.MAX_ROWS_IN_SEGMENT);
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
      Assert.assertEquals(segmentDataManager.getStopReason(), SegmentCompletionProtocol.REASON_ROW_LIMIT);
    }
    // test reaching max time limit
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.INITIAL_CONSUMING);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      // We should still get false because there is no messages fetched
      segmentDataManager._timeSupplier.add(Fixtures.MAX_TIME_FOR_SEGMENT_CLOSE_MS + 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      // Once there are messages fetched, and the time exceeds the extended hour, we should get true
      setHasMessagesFetched(segmentDataManager, true);
      segmentDataManager._timeSupplier.add(TimeUnit.HOURS.toMillis(1));
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
      Assert.assertEquals(segmentDataManager.getStopReason(), SegmentCompletionProtocol.REASON_TIME_LIMIT);
    }
    // In catching up state, test reaching final offset
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CATCHING_UP);
      final long finalOffset = START_OFFSET_VALUE + 100;
      segmentDataManager.setFinalOffset(finalOffset);
      segmentDataManager.setCurrentOffset(finalOffset - 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager.setCurrentOffset(finalOffset);
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
    }
    // In catching up state, test reaching final offset ignoring time
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._timeSupplier.add(Fixtures.MAX_TIME_FOR_SEGMENT_CLOSE_MS);
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CATCHING_UP);
      final long finalOffset = START_OFFSET_VALUE + 100;
      segmentDataManager.setFinalOffset(finalOffset);
      segmentDataManager.setCurrentOffset(finalOffset - 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager.setCurrentOffset(finalOffset);
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
    }
    // When we go from consuming to online state, time and final offset matter.
    // Case 1. We have reached final offset.
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._timeSupplier.add(1);
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CONSUMING_TO_ONLINE);
      segmentDataManager.setConsumeEndTime(segmentDataManager._timeSupplier.get() + 10);
      final long finalOffset = START_OFFSET_VALUE + 100;
      segmentDataManager.setFinalOffset(finalOffset);
      segmentDataManager.setCurrentOffset(finalOffset - 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager.setCurrentOffset(finalOffset);
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
    }
    // Case 2. We have reached time limit.
    try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) {
      segmentDataManager._state.set(segmentDataManager, RealtimeSegmentDataManager.State.CONSUMING_TO_ONLINE);
      final long endTime = segmentDataManager._timeSupplier.get() + 10;
      segmentDataManager.setConsumeEndTime(endTime);
      final long finalOffset = START_OFFSET_VALUE + 100;
      segmentDataManager.setFinalOffset(finalOffset);
      segmentDataManager.setCurrentOffset(finalOffset - 1);
      segmentDataManager._timeSupplier.set(endTime - 1);
      Assert.assertFalse(segmentDataManager.invokeEndCriteriaReached());
      segmentDataManager._timeSupplier.set(endTime);
      Assert.assertTrue(segmentDataManager.invokeEndCriteriaReached());
    }
  }