public int computeThreshold(StreamConfig streamConfig, CommittingSegmentDescriptor committingSegmentDescriptor,
      @Nullable SegmentZKMetadata committingSegmentZKMetadata, String newSegmentName) {
    long desiredSegmentSizeBytes = streamConfig.getFlushThresholdSegmentSizeBytes();
    if (desiredSegmentSizeBytes <= 0) {
      desiredSegmentSizeBytes = StreamConfig.DEFAULT_FLUSH_THRESHOLD_SEGMENT_SIZE_BYTES;
    }
    long optimalSegmentSizeBytesMin = desiredSegmentSizeBytes / 2;
    double optimalSegmentSizeBytesMax = desiredSegmentSizeBytes * 1.5;

    if (committingSegmentZKMetadata == null) { // first segment of the partition, hence committing segment is null
      if (_latestSegmentRowsToSizeRatio > 0) { // new partition group added case
        long targetSegmentNumRows = (long) (desiredSegmentSizeBytes * _latestSegmentRowsToSizeRatio);
        targetSegmentNumRows = capNumRowsIfOverflow(targetSegmentNumRows);
        SegmentSizeBasedFlushThresholdUpdater.LOGGER.info(
            "Committing segment zk metadata is not available, using prev ratio {}, setting rows threshold for {} as {}",
            _latestSegmentRowsToSizeRatio, newSegmentName, targetSegmentNumRows);
        return (int) targetSegmentNumRows;
      } else {
        final int autotuneInitialRows = streamConfig.getFlushAutotuneInitialRows();
        SegmentSizeBasedFlushThresholdUpdater.LOGGER.info(
            "Committing segment zk metadata is not available, setting threshold for {} as {}", newSegmentName,
            autotuneInitialRows);
        return autotuneInitialRows;
      }
    }

    final long committingSegmentSizeBytes = committingSegmentDescriptor.getSegmentSizeBytes();
    if (committingSegmentSizeBytes <= 0 // repair segment case
        || SegmentCompletionProtocol.REASON_FORCE_COMMIT_MESSAGE_RECEIVED.equals(
        committingSegmentDescriptor.getStopReason())) {
      String reason = committingSegmentSizeBytes <= 0 //
          ? "Committing segment size is not available" //
          : "Committing segment is due to force-commit";
      final int targetNumRows = committingSegmentZKMetadata.getSizeThresholdToFlushSegment();
      SegmentSizeBasedFlushThresholdUpdater.LOGGER.info("{}, setting thresholds from previous segment for {} as {}",
          reason, newSegmentName, targetNumRows);
      return targetNumRows;
    }

    final long timeConsumed = _clock.millis() - committingSegmentZKMetadata.getCreationTime();
    final long numRowsConsumed = committingSegmentZKMetadata.getTotalDocs();
    final int numRowsThreshold = committingSegmentZKMetadata.getSizeThresholdToFlushSegment();
    SegmentSizeBasedFlushThresholdUpdater.LOGGER.info(
        "{}: Data from committing segment: Time {}  numRows {} threshold {} segmentSize(bytes) {}",
        newSegmentName, TimeUtils.convertMillisToPeriod(timeConsumed), numRowsConsumed, numRowsThreshold,
        committingSegmentSizeBytes);

    double currentRatio = (double) numRowsConsumed / committingSegmentSizeBytes;
    if (_latestSegmentRowsToSizeRatio > 0) {
      _latestSegmentRowsToSizeRatio =
          CURRENT_SEGMENT_RATIO_WEIGHT * currentRatio + PREVIOUS_SEGMENT_RATIO_WEIGHT * _latestSegmentRowsToSizeRatio;
    } else {
      _latestSegmentRowsToSizeRatio = currentRatio;
    }

    // If the number of rows consumed is less than what we set as target in metadata, then the segment hit time limit.
    // We can set the new target to be slightly higher than the actual number of rows consumed so that we can aim
    // to hit the row limit next time around.
    //
    // If the size of the committing segment is higher than the desired segment size, then the administrator has
    // set a lower segment size threshold. We should treat this case as if we have hit thw row limit and not the time
    // limit.
    //
    // TODO: add feature to adjust time threshold as well
    // If we set new threshold to be numRowsConsumed, we might keep oscillating back and forth between doubling limit
    // and time threshold being hit If we set new threshold to be committingSegmentZKMetadata
    // .getSizeThresholdToFlushSegment(),
    // we might end up using a lot more memory than required for the segment Using a minor bump strategy, until
    // we add feature to adjust time We will only slightly bump the threshold based on numRowsConsumed
    if (numRowsConsumed < numRowsThreshold && committingSegmentSizeBytes < desiredSegmentSizeBytes) {
      final long timeThresholdMillis = streamConfig.getFlushThresholdTimeMillis();
      long currentNumRows = numRowsConsumed;
      StringBuilder logStringBuilder = new StringBuilder().append("Time threshold reached. ");
      if (timeThresholdMillis < timeConsumed) {
        // The administrator has reduced the time threshold. Adjust the
        // number of rows to match the average consumption rate on the partition.
        currentNumRows = timeThresholdMillis * numRowsConsumed / timeConsumed;
        logStringBuilder.append(" Detected lower time threshold, adjusting numRowsConsumed to ").append(currentNumRows)
            .append(". ");
      }
      long targetSegmentNumRows = (long) (currentNumRows * ROWS_MULTIPLIER_WHEN_TIME_THRESHOLD_HIT);
      targetSegmentNumRows = capNumRowsIfOverflow(targetSegmentNumRows);
      logStringBuilder.append("Setting segment size for {} as {}");
      SegmentSizeBasedFlushThresholdUpdater.LOGGER.info(logStringBuilder.toString(),
          newSegmentName, targetSegmentNumRows);
      return (int) targetSegmentNumRows;
    }

    long targetSegmentNumRows;
    if (committingSegmentSizeBytes < optimalSegmentSizeBytesMin) {
      targetSegmentNumRows = numRowsConsumed + numRowsConsumed / 2;
    } else if (committingSegmentSizeBytes > optimalSegmentSizeBytesMax) {
      targetSegmentNumRows = numRowsConsumed / 2;
    } else {
      if (_latestSegmentRowsToSizeRatio > 0) {
        targetSegmentNumRows = (long) (desiredSegmentSizeBytes * _latestSegmentRowsToSizeRatio);
      } else {
        targetSegmentNumRows = (long) (desiredSegmentSizeBytes * currentRatio);
      }
    }
    targetSegmentNumRows = capNumRowsIfOverflow(targetSegmentNumRows);
    SegmentSizeBasedFlushThresholdUpdater.LOGGER.info(
        "Committing segment size {}, current ratio {}, setting threshold for {} as {}",
        committingSegmentSizeBytes, _latestSegmentRowsToSizeRatio, newSegmentName, targetSegmentNumRows);

    return (int) targetSegmentNumRows;
  }
@Test
  public void testApplyMultiplierToAdjustedTotalDocsWhenTimeThresholdIsReached() {
    long currentTime = 1640216032391L;
    Clock clock = Clock.fixed(java.time.Instant.ofEpochMilli(currentTime), ZoneId.of("UTC"));

    SegmentFlushThresholdComputer computer = new SegmentFlushThresholdComputer(clock);

    StreamConfig streamConfig = mock(StreamConfig.class);
    when(streamConfig.getFlushThresholdSegmentSizeBytes()).thenReturn(300_0000L);
    when(streamConfig.getFlushThresholdTimeMillis()).thenReturn(MILLISECONDS.convert(1, TimeUnit.HOURS));

    CommittingSegmentDescriptor committingSegmentDescriptor = mock(CommittingSegmentDescriptor.class);
    when(committingSegmentDescriptor.getSegmentSizeBytes()).thenReturn(200_0000L);

    SegmentZKMetadata committingSegmentZKMetadata = mock(SegmentZKMetadata.class);
    when(committingSegmentZKMetadata.getTotalDocs()).thenReturn(30_000L);
    when(committingSegmentZKMetadata.getSizeThresholdToFlushSegment()).thenReturn(60_000);
    when(committingSegmentZKMetadata.getCreationTime()).thenReturn(
        currentTime - MILLISECONDS.convert(2, TimeUnit.HOURS));

    int threshold = computer.computeThreshold(streamConfig, committingSegmentDescriptor, committingSegmentZKMetadata,
        "events3__0__0__20211222T1646Z");

    // (totalDocs / 2) * 1.1
    // (30000 / 2) * 1.1
    // 15000 * 1.1
    assertEquals(threshold, 16_500);
  }