<commit_before>/*******************************************************************************
 * Copyright (C) 2015 Google Inc.
 *
 * Licensed 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 com.google.cloud.dataflow.sdk.util.common.worker;

import com.google.cloud.dataflow.sdk.util.common.Counter;
import com.google.cloud.dataflow.sdk.util.common.CounterSet;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

import javax.annotation.concurrent.ThreadSafe;

/**
 * A StateSampler object may be used to obtain an approximate
 * breakdown of the time spent by an execution context in various
 * states, as a fraction of the total time.  The sampling is taken at
 * regular intervals, with adjustment for scheduling delay.
 */
@ThreadSafe
public class StateSampler implements AutoCloseable {

  /** Different kinds of states. */
  public enum StateKind {
    /** IO, user code, etc. */
    USER,
    /** Reading/writing from/to shuffle service, etc. */
    FRAMEWORK
  }

  public static final long DEFAULT_SAMPLING_PERIOD_MS = 200;

  private final String prefix;
  private final CounterSet.AddCounterMutator counterSetMutator;

  /** Array of counters indexed by their state. */
  private ArrayList<Counter<Long>> countersByState = new ArrayList<>();

  /** Map of state name to state. */
  private Map<String, Integer> statesByName = new HashMap<>();

  /** Map of state id to kind. */
  private Map<Integer, StateKind> kindsByState = new HashMap<>();

  /** The current state. */
  private volatile int currentState;

  /** Special value of {@code currentState} that means we do not sample. */
  private static final int DO_NOT_SAMPLE = -1;

  /**
   * A counter that increments with each state transition. May be used
   * to detect a context being stuck in a state for some amount of
   * time.
   */
  private volatile long stateTransitionCount;

  /**
   * The timestamp (in nanoseconds) corresponding to the last time the
   * state was sampled (and recorded).
   */
  private long stateTimestampNs = 0;

  /** Using a fixed number of timers for all StateSampler objects. */
  private static final int NUM_EXECUTOR_THREADS = 16;

  private static final ScheduledExecutorService executorService =
      Executors.newScheduledThreadPool(NUM_EXECUTOR_THREADS,
          new ThreadFactoryBuilder().setDaemon(true).build());

  private Random rand = new Random();

  private List<SamplingCallback> callbacks = new ArrayList<>();

  private ScheduledFuture<?> invocationTriggerFuture = null;

  private ScheduledFuture<?> invocationFuture = null;

  /**
   * Constructs a new {@link StateSampler} that can be used to obtain
   * an approximate breakdown of the time spent by an execution
   * context in various states, as a fraction of the total time.
   *
   * @param prefix the prefix of the counter names for the states
   * @param counterSetMutator the {@link CounterSet.AddCounterMutator}
   * used to create a counter for each distinct state
   * @param samplingPeriodMs the sampling period in milliseconds
   */
  public StateSampler(String prefix,
                      CounterSet.AddCounterMutator counterSetMutator,
                      final long samplingPeriodMs) {
    this.prefix = prefix;
    this.counterSetMutator = counterSetMutator;
    currentState = DO_NOT_SAMPLE;
    // Here "stratified sampling" is used, which makes sure that there's 1 uniformly chosen sampled
    // point in every bucket of samplingPeriodMs, to prevent pathological behavior in case some
    // states happen to occur at a similar period.
    // The current implementation uses a fixed-rate timer with a period samplingPeriodMs as a
    // trampoline to a one-shot random timer which fires with a random delay within
    // samplingPeriodMs.
    invocationTriggerFuture =
        executorService.scheduleAtFixedRate(
            new Runnable() {
              @Override
              public void run() {
                long delay = rand.nextInt((int) samplingPeriodMs);
                synchronized (StateSampler.this) {
                  if (invocationFuture != null) {
                    invocationFuture.cancel(false);
                  }
                  invocationFuture =
                      executorService.schedule(
                          new Runnable() {
                            @Override
                            public void run() {
                              StateSampler.this.run();
                            }
                          },
                          delay,
                          TimeUnit.MILLISECONDS);
                }
              }
            },
            0,
            samplingPeriodMs,
            TimeUnit.MILLISECONDS);
    stateTimestampNs = System.nanoTime();
  }

  /**
   * Constructs a new {@link StateSampler} that can be used to obtain
   * an approximate breakdown of the time spent by an execution
   * context in various states, as a fraction of the total time.
   *
   * @param prefix the prefix of the counter names for the states
   * @param counterSetMutator the {@link CounterSet.AddCounterMutator}<commit_msg>Fixed a race in StateSampler.<commit_after>/*******************************************************************************
 * Copyright (C) 2015 Google Inc.
 *
 * Licensed 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 com.google.cloud.dataflow.sdk.util.common.worker;

import com.google.cloud.dataflow.sdk.util.common.Counter;
import com.google.cloud.dataflow.sdk.util.common.CounterSet;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

import javax.annotation.concurrent.ThreadSafe;

/**
 * A StateSampler object may be used to obtain an approximate
 * breakdown of the time spent by an execution context in various
 * states, as a fraction of the total time.  The sampling is taken at
 * regular intervals, with adjustment for scheduling delay.
 */
@ThreadSafe
public class StateSampler implements AutoCloseable {

  /** Different kinds of states. */
  public enum StateKind {
    /** IO, user code, etc. */
    USER,
    /** Reading/writing from/to shuffle service, etc. */
    FRAMEWORK
  }

  public static final long DEFAULT_SAMPLING_PERIOD_MS = 200;

  private final String prefix;
  private final CounterSet.AddCounterMutator counterSetMutator;

  /** Array of counters indexed by their state. */
  private ArrayList<Counter<Long>> countersByState = new ArrayList<>();

  /** Map of state name to state. */
  private Map<String, Integer> statesByName = new HashMap<>();

  /** Map of state id to kind. */
  private Map<Integer, StateKind> kindsByState = new HashMap<>();

  /** The current state. */
  private volatile int currentState;

  /** Special value of {@code currentState} that means we do not sample. */
  private static final int DO_NOT_SAMPLE = -1;

  /**
   * A counter that increments with each state transition. May be used
   * to detect a context being stuck in a state for some amount of
   * time.
   */
  private volatile long stateTransitionCount;

  /**
   * The timestamp (in nanoseconds) corresponding to the last time the
   * state was sampled (and recorded).
   */
  private long stateTimestampNs = 0;

  /** Using a fixed number of timers for all StateSampler objects. */
  private static final int NUM_EXECUTOR_THREADS = 16;

  private static final ScheduledExecutorService executorService =
      Executors.newScheduledThreadPool(NUM_EXECUTOR_THREADS,
          new ThreadFactoryBuilder().setDaemon(true).build());

  private Random rand = new Random();

  private List<SamplingCallback> callbacks = new ArrayList<>();

  private ScheduledFuture<?> invocationTriggerFuture = null;

  private ScheduledFuture<?> invocationFuture = null;

  /**
   * Constructs a new {@link StateSampler} that can be used to obtain
   * an approximate breakdown of the time spent by an execution
   * context in various states, as a fraction of the total time.
   *
   * @param prefix the prefix of the counter names for the states
   * @param counterSetMutator the {@link CounterSet.AddCounterMutator}
   * used to create a counter for each distinct state
   * @param samplingPeriodMs the sampling period in milliseconds
   */
  public StateSampler(String prefix,
                      CounterSet.AddCounterMutator counterSetMutator,
                      final long samplingPeriodMs) {
    this.prefix = prefix;
    this.counterSetMutator = counterSetMutator;
    currentState = DO_NOT_SAMPLE;
    // Here "stratified sampling" is used, which makes sure that there's 1 uniformly chosen sampled
    // point in every bucket of samplingPeriodMs, to prevent pathological behavior in case some
    // states happen to occur at a similar period.
    // The current implementation uses a fixed-rate timer with a period samplingPeriodMs as a
    // trampoline to a one-shot random timer which fires with a random delay within
    // samplingPeriodMs.
    stateTimestampNs = System.nanoTime();
    invocationTriggerFuture =
        executorService.scheduleAtFixedRate(
            new Runnable() {
              @Override
              public void run() {
                long delay = rand.nextInt((int) samplingPeriodMs);
                synchronized (StateSampler.this) {
                  if (invocationFuture != null) {
                    invocationFuture.cancel(false);
                  }
                  invocationFuture =
                      executorService.schedule(
                          new Runnable() {
                            @Override
                            public void run() {
                              StateSampler.this.run();
                            }
                          },
                          delay,
                          TimeUnit.MILLISECONDS);
                }
              }
            },
            0,
            samplingPeriodMs,
            TimeUnit.MILLISECONDS);
  }

  /**
   * Constructs a new {@link StateSampler} that can be used to obtain
   * an approximate breakdown of the time spent by an execution
   * context in various states, as a fraction of the total time.
   *
   * @param prefix the prefix of the counter names for the states
   * @param counterSetMutator the {@link CounterSet.AddCounterMutator}
{'fraction_non_alphanumeric': 0.09936592691473385, 'fraction_numerical': 0.002836642749874854, 'mean_word_length': 2.788558786346397, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3144700', 'n_tokens_mistral': 3035, 'n_tokens_neox': 2897, 'n_words': 1129}