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.runners.dataflow.worker.util.common.worker;
  
  import java.util.Random;
  import java.util.concurrent.ThreadLocalRandom;
  import org.apache.beam.runners.core.ElementByteSizeObservable;
  import org.apache.beam.runners.dataflow.worker.counters.Counter;
  import org.apache.beam.runners.dataflow.worker.counters.CounterBackedElementByteSizeObserver;
  import org.apache.beam.runners.dataflow.worker.counters.CounterFactory;
  import org.apache.beam.runners.dataflow.worker.counters.CounterFactory.CounterMean;
  import org.apache.beam.runners.dataflow.worker.counters.CounterName;
  import org.apache.beam.runners.dataflow.worker.counters.NameContext;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /** An {@link ElementCounter} that counts output objects, bytes, and mean bytes. */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class OutputObjectAndByteCounter implements ElementCounter {
  
    // Might be null, e.g., undeclared outputs will not have an
    // elementByteSizeObservable.
    private final ElementByteSizeObservable<Object> elementByteSizeObservable;
    private final CounterFactory counterFactory;
  
    // Lowest sampling probability: 0.001%.
    private static final int SAMPLING_TOKEN_UPPER_BOUND = 1000000;
    private static final int SAMPLING_CUTOFF = 10;
    private int samplingToken = 0;
  
    private Counter<Long, ?> objectCount = null;
    private Counter<Long, Long> byteCount = null;
    private Counter<Long, CounterMean<Long>> meanByteCount = null;
    private CounterBackedElementByteSizeObserver byteCountObserver = null;
    private CounterBackedElementByteSizeObserver meanByteCountObserver = null;
    private int samplingTokenUpperBound = SAMPLING_TOKEN_UPPER_BOUND;
  
    // A variant which does not take a {@link NameContext} for flume.
    @SuppressWarnings("unchecked")
    public OutputObjectAndByteCounter(
        ElementByteSizeObservable<?> elementByteSizeObservable, CounterFactory counterFactory) {
      this(elementByteSizeObservable, counterFactory, null);
    }
  
    @SuppressWarnings("unchecked")
    public OutputObjectAndByteCounter(
        ElementByteSizeObservable<?> elementByteSizeObservable,
        CounterFactory counterFactory,
        @Nullable NameContext nameContext) {
      this.elementByteSizeObservable = (ElementByteSizeObservable<Object>) elementByteSizeObservable;
      this.counterFactory = counterFactory;
    }
  
    /** Count output objects. */
    public OutputObjectAndByteCounter countObject(String objectCounterName) {
      objectCount = counterFactory.longSum(getCounterName(objectCounterName));
      return this;
    }
  
    /** Count output bytes. */
    public OutputObjectAndByteCounter countBytes(String bytesCounterName) {
      if (elementByteSizeObservable != null) {
        byteCount = counterFactory.longSum(getCounterName(bytesCounterName));
        byteCountObserver = new CounterBackedElementByteSizeObserver(byteCount);
      }
      return this;
    }
  
    /** Count output mean byte. */
    public OutputObjectAndByteCounter countMeanByte(String meanByteCounterName) {
      if (elementByteSizeObservable != null) {
        meanByteCount = counterFactory.longMean(getCounterName(meanByteCounterName));
        meanByteCountObserver = new CounterBackedElementByteSizeObserver(meanByteCount);
      }
      return this;
    }
  
    /** Sets the minimum sampling rate to the inverse of the given value. */
    public OutputObjectAndByteCounter setSamplingPeriod(int period) {
      this.samplingTokenUpperBound = period * SAMPLING_CUTOFF;
      return this;
    }
  
    public Counter<Long, ?> getObjectCount() {
      return objectCount;
    }
  
    public Counter<Long, Long> getByteCount() {
      return byteCount;
    }
  
    public Counter<Long, CounterMean<Long>> getMeanByteCount() {
      return meanByteCount;
    }
  
    @Override
    public void update(Object elem) throws Exception {
      // Increment object counter.
      if (objectCount != null) {
        objectCount.addValue(1L);
      }
  
      // Increment byte counter.
      if ((byteCountObserver != null || meanByteCountObserver != null)
          && (sampleElement() || elementByteSizeObservable.isRegisterByteSizeObserverCheap(elem))) {
        if (byteCountObserver != null) {
          byteCountObserver.setScalingFactor(
              Math.max(samplingToken, SAMPLING_CUTOFF) / (double) SAMPLING_CUTOFF);
          elementByteSizeObservable.registerByteSizeObserver(elem, byteCountObserver);
        }
        if (meanByteCountObserver != null) {
          elementByteSizeObservable.registerByteSizeObserver(elem, meanByteCountObserver);
        }
  
        if (byteCountObserver != null && !byteCountObserver.getIsLazy()) {
          byteCountObserver.advance();
        }
        if (meanByteCountObserver != null && !meanByteCountObserver.getIsLazy()) {
          meanByteCountObserver.advance();
        }
      }
    }
  
    @Override
    public void finishLazyUpdate(Object elem) {
      // Advance lazy ElementByteSizeObservers, if any.
      // Note that user's code is allowed to store the element of one
      // DoFn.processElement() call and access it later on. We are still
      // calling next() here, causing an update to byteCount. If user's
      // code really accesses more element's pieces later on, their byte
      // count would accrue against a future element. This is not ideal,
      // but still approximately correct.
      if (byteCountObserver != null && byteCountObserver.getIsLazy()) {
        byteCountObserver.advance();
      }
      if (meanByteCountObserver != null && meanByteCountObserver.getIsLazy()) {
        meanByteCountObserver.advance();
      }
    }
  
    protected boolean sampleElement() {
      // Sampling probability decreases as the element count is increasing.
      // We unconditionally sample the first samplingCutoff elements. For the
      // next samplingCutoff elements, the sampling probability drops from 100%
      // to 50%. The probability of sampling the Nth element is:
      // min(1, samplingCutoff / N), with an additional lower bound of
      // samplingCutoff / samplingTokenUpperBound. This algorithm may be refined
      // later.
      samplingToken = Math.min(samplingToken + 1, samplingTokenUpperBound);
      return getRandom().nextInt(samplingToken) < SAMPLING_CUTOFF;
    }
  
    @VisibleForTesting
    protected Random getRandom() {
      return ThreadLocalRandom.current();
    }
  
    private CounterName getCounterName(String name) {
      // TODO: use original name from the NameContext
      return CounterName.named(name);
    }
  }
  