code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.db.jdbc.streaming;
  
  import com.google.common.annotations.VisibleForTesting;
  import java.util.Optional;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * This estimator first uses the {@link InitialSizeEstimator} to calculate an initial fetch size by
   * sampling the first N rows consecutively, and then switches to {@link SamplingSizeEstimator} to
   * periodically adjust the fetch size by sampling every M rows.
   */
  public class TwoStageSizeEstimator implements FetchSizeEstimator {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(TwoStageSizeEstimator.class);
  
    private final int initialSampleSize;
    private BaseSizeEstimator delegate;
    private int counter = 0;
  
    public static TwoStageSizeEstimator getInstance() {
      return new TwoStageSizeEstimator();
    }
  
    private TwoStageSizeEstimator() {
      this.initialSampleSize = FetchSizeConstants.INITIAL_SAMPLE_SIZE;
      this.delegate = new InitialSizeEstimator(
          FetchSizeConstants.MIN_BUFFER_BYTE_SIZE,
          initialSampleSize,
          FetchSizeConstants.MIN_FETCH_SIZE,
          FetchSizeConstants.DEFAULT_FETCH_SIZE,
          FetchSizeConstants.MAX_FETCH_SIZE);
    }
  
    @Override
    public Optional<Integer> getFetchSize() {
      return delegate.getFetchSize();
    }
  
    @Override
    public void accept(final Object rowData) {
      if (counter <= initialSampleSize + 1) {
        counter++;
        // switch to SamplingSizeEstimator after the initial N rows
        if (delegate instanceof InitialSizeEstimator && counter > initialSampleSize) {
          delegate = new SamplingSizeEstimator(
              getTargetBufferByteSize(Runtime.getRuntime().maxMemory()),
              FetchSizeConstants.SAMPLE_FREQUENCY,
              delegate.getMaxRowByteSize(),
              FetchSizeConstants.MIN_FETCH_SIZE,
              FetchSizeConstants.DEFAULT_FETCH_SIZE,
              FetchSizeConstants.MAX_FETCH_SIZE);
        }
      }
  
      delegate.accept(rowData);
    }
  
    @VisibleForTesting
    static long getTargetBufferByteSize(final Long maxMemory) {
      if (maxMemory == null || maxMemory == Long.MAX_VALUE) {
        LOGGER.info("No max memory limit found, use min JDBC buffer size: {}", FetchSizeConstants.MIN_BUFFER_BYTE_SIZE);
        return FetchSizeConstants.MIN_BUFFER_BYTE_SIZE;
      }
      final long targetBufferByteSize = Math.round(maxMemory * FetchSizeConstants.TARGET_BUFFER_SIZE_RATIO);
      final long finalBufferByteSize = Math.min(FetchSizeConstants.MAX_BUFFER_BYTE_SIZE,
          Math.max(FetchSizeConstants.MIN_BUFFER_BYTE_SIZE, targetBufferByteSize));
      LOGGER.info("Max memory limit: {}, JDBC buffer size: {}", maxMemory, finalBufferByteSize);
      return finalBufferByteSize;
    }
  
    @VisibleForTesting
    BaseSizeEstimator getDelegate() {
      return delegate;
    }
  
  }
  