code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.buffered_stream_consumer;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import com.google.common.annotations.VisibleForTesting;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.protocol.models.AirbyteRecordMessage;
  import java.util.HashMap;
  import java.util.Map;
  
  /**
   * This class estimate the byte size of the record message. To reduce memory footprint, 1) it
   * assumes that a character is always four bytes, and 2) it only performs a sampling every N
   * records. The size of the samples are averaged together to protect the estimation against
   * outliers.
   */
  public class RecordSizeEstimator {
  
    // by default, perform one estimation for every 20 records
    private static final int DEFAULT_SAMPLE_BATCH_SIZE = 20;
  
    // latest estimated record message size for each stream
    private final Map<String, Long> streamRecordSizeEstimation;
    // number of record messages until next real sampling for each stream
    private final Map<String, Integer> streamSampleCountdown;
    // number of record messages
    private final int sampleBatchSize;
  
    /**
     * The estimator will perform a real calculation once per sample batch. The size of the batch is
     * determined by {@code sampleBatchSize}.
     */
    public RecordSizeEstimator(final int sampleBatchSize) {
      this.streamRecordSizeEstimation = new HashMap<>();
      this.streamSampleCountdown = new HashMap<>();
      this.sampleBatchSize = sampleBatchSize;
    }
  
    public RecordSizeEstimator() {
      this(DEFAULT_SAMPLE_BATCH_SIZE);
    }
  
    public long getEstimatedByteSize(final AirbyteRecordMessage recordMessage) {
      final String stream = recordMessage.getStream();
      final Integer countdown = streamSampleCountdown.get(stream);
  
      // this is a new stream; initialize its estimation
      if (countdown == null) {
        final long byteSize = getStringByteSize(recordMessage.getData());
        streamRecordSizeEstimation.put(stream, byteSize);
        streamSampleCountdown.put(stream, sampleBatchSize - 1);
        return byteSize;
      }
  
      // this stream needs update; compute a new estimation
      if (countdown <= 0) {
        final long prevMeanByteSize = streamRecordSizeEstimation.get(stream);
        final long currentByteSize = getStringByteSize(recordMessage.getData());
        final long newMeanByteSize = prevMeanByteSize / 2 + currentByteSize / 2;
        streamRecordSizeEstimation.put(stream, newMeanByteSize);
        streamSampleCountdown.put(stream, sampleBatchSize - 1);
        return newMeanByteSize;
      }
  
      // this stream does not need update; return current estimation
      streamSampleCountdown.put(stream, countdown - 1);
      return streamRecordSizeEstimation.get(stream);
    }
  
    @VisibleForTesting
    static long getStringByteSize(final JsonNode data) {
      // assume UTF-8 encoding, and each char is 4 bytes long
      return Jsons.serialize(data).length() * 4L;
    }
  
  }
  