content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.sdk.transforms.splittabledofn;
  
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
  
  import java.math.BigDecimal;
  import java.math.MathContext;
  import org.apache.beam.sdk.io.range.OffsetRange;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Suppliers;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong;
  
  /**
   * An {@link OffsetRangeTracker} for tracking a growable offset range. {@code Long.MAX_VALUE} is
   * used as the end of the range to indicate infinity.
   *
   * <p>An offset range is considered growable when the end offset could grow (or change) during
   * execution time (e.g., Kafka topic partition offset, appended file, ...).
   *
   * <p>The growable range is marked as done by claiming {@code Long.MAX_VALUE}.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class GrowableOffsetRangeTracker extends OffsetRangeTracker {
    /**
     * Provides the estimated end offset of the range.
     *
     * <p>{@link #estimate} is called to give the end offset when {@link #trySplit} or {@link
     * #getProgress} is invoked. The end offset is exclusive for the range. The estimated end is not
     * required to monotonically increase as it will only be taken into consideration when the
     * estimated end offset is larger than the current position. Returning {@code Long.MAX_VALUE} as
     * the estimate implies the largest possible position for the range is {@code Long.MAX_VALUE - 1}.
     * Return {@code Long.MIN_VALUE} if an estimate can not be provided.
     *
     * <p>Providing a good estimate is important for an accurate progress signal and will impact
     * splitting decisions by the runner.
     *
     * <p>If {@link #estimate} is expensive to compute, consider wrapping the implementation with
     * {@link Suppliers#memoizeWithExpiration} or equivalent as an optimization.
     *
     * <p>TODO(https://github.com/apache/beam/issues/20227): Also consider using {@link
     * RangeEndEstimator} when the range is not ended with {@code Long.MAX_VALUE}.
     */
    @FunctionalInterface
    public interface RangeEndEstimator {
      long estimate();
    }
  
    private final RangeEndEstimator rangeEndEstimator;
  
    public GrowableOffsetRangeTracker(long start, RangeEndEstimator rangeEndEstimator) {
      super(new OffsetRange(start, Long.MAX_VALUE));
      this.rangeEndEstimator = checkNotNull(rangeEndEstimator);
    }
  
    // TODO(sjvanrossum): Use UnsignedLong instead of BigDecimal for splitting ranges
    @Override
    public SplitResult<OffsetRange> trySplit(double fractionOfRemainder) {
      // If current tracking range is no longer growable, split it as a normal range.
      if (range.getTo() != Long.MAX_VALUE || range.getTo() == range.getFrom()) {
        return super.trySplit(fractionOfRemainder);
      }
      // If current range has been done, there is no more space to split.
      if (lastAttemptedOffset != null && lastAttemptedOffset == Long.MAX_VALUE) {
        return null;
      }
      BigDecimal cur =
          (lastAttemptedOffset == null)
              ? BigDecimal.valueOf(range.getFrom()).subtract(BigDecimal.ONE, MathContext.DECIMAL128)
              : BigDecimal.valueOf(lastAttemptedOffset);
  
      // Fetch the estimated end offset. If the estimated end is smaller than the next offset, use
      // the next offset as end.
      BigDecimal estimateRangeEnd =
          BigDecimal.valueOf(rangeEndEstimator.estimate())
              .max(cur.add(BigDecimal.ONE, MathContext.DECIMAL128));
  
      // Convert to BigDecimal in computation to prevent overflow, which may result in loss of
      // precision.
      // split = cur + max(1, (estimateRangeEnd - cur) * fractionOfRemainder)
      BigDecimal splitPos =
          cur.add(
              estimateRangeEnd
                  .subtract(cur, MathContext.DECIMAL128)
                  .multiply(BigDecimal.valueOf(fractionOfRemainder), MathContext.DECIMAL128)
                  .max(BigDecimal.ONE),
              MathContext.DECIMAL128);
      long split = splitPos.longValue();
      if (split > estimateRangeEnd.longValue()) {
        return null;
      }
      OffsetRange res = new OffsetRange(split, range.getTo());
      this.range = new OffsetRange(range.getFrom(), split);
      return SplitResult.of(range, res);
    }
  
    @Override
    public Progress getProgress() {
      // If current tracking range is no longer growable, get progress as a normal range.
      if (range.getTo() != Long.MAX_VALUE || range.getTo() == range.getFrom()) {
        return super.getProgress();
      }
  
      final long completedEnd = lastAttemptedOffset == null ? range.getFrom() : lastAttemptedOffset;
      final long remainingEnd = Math.max(completedEnd, rangeEndEstimator.estimate());
  
      return Progress.from(
          UnsignedLong.fromLongBits(completedEnd - range.getFrom()).doubleValue(),
          UnsignedLong.fromLongBits(remainingEnd - completedEnd).doubleValue());
    }
  
    @Override
    public IsBounded isBounded() {
      // If current range has been done, the range should be bounded.
      if (lastAttemptedOffset != null && lastAttemptedOffset == Long.MAX_VALUE) {
        return IsBounded.BOUNDED;
      }
      return range.getTo() == Long.MAX_VALUE ? IsBounded.UNBOUNDED : IsBounded.BOUNDED;
    }
  }
  