<commit_before>/*
 * 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;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;

import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.SideInputInfo;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState;
import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext;
import org.apache.beam.runners.dataflow.worker.counters.CounterFactory;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.metrics.MetricsContainer;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.FluentIterable;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** {@link DataflowExecutionContext} for use in streaming mode. */
public class StreamingModeExecutionContext extends DataflowExecutionContext<StepContext> {

  private static final Logger LOG = LoggerFactory.getLogger(StreamingModeExecutionContext.class);

  /**
   * The current user-facing key for this execution context.
   *
   * <p>This field is set to a new key upon each call to {@link #start}, which corresponds to one
   * input bundle. For non-keyed bundles, the key will be {@code null}.
   *
   * <p>This key should not be mistaken for the sharding key of the computation, which is always
   * present.
   */
  @Nullable private Object key = null;

  private final String computationId;
  private final Map<TupleTag<?>, Map<BoundedWindow, Object>> sideInputCache;

  // Per-key cache of active Reader objects in use by this process.
  private final ImmutableMap<String, String> stateNameMap;
  private final WindmillStateCache.ForComputation stateCache;

  private Windmill.WorkItem work;
  private StateFetcher stateFetcher;
  private Windmill.WorkItemCommitRequest.Builder outputBuilder;
  private UnboundedSource.UnboundedReader<?> activeReader;
  private volatile long backlogBytes;
  private final ReaderCache readerCache;

  public StreamingModeExecutionContext(
      CounterFactory counterFactory,
      String computationId,
      ReaderCache readerCache,
      Map<String, String> stateNameMap,
      WindmillStateCache.ForComputation stateCache,
      MetricsContainerRegistry<StreamingStepMetricsContainer> metricsContainerRegistry,
      DataflowExecutionStateTracker executionStateTracker,
      StreamingModeExecutionStateRegistry executionStateRegistry,
      long sinkByteLimit) {
    super(
        counterFactory,
        metricsContainerRegistry,
        executionStateTracker,
        executionStateRegistry,
        sinkByteLimit);
    this.computationId = computationId;
    this.readerCache = readerCache;
    this.sideInputCache = new HashMap<>();
    this.stateNameMap = ImmutableMap.copyOf(stateNameMap);
    this.stateCache = stateCache;
    this.backlogBytes = UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN;
  }

  /**
   * Execution states in Streaming are shared between multiple map-task executors. Thus this class
   * needs to be thread safe for multiple writers. A single stage could have have multiple executors
   * running concurrently.
   */
  public static class StreamingModeExecutionState
      extends DataflowOperationContext.DataflowExecutionState {

    // AtomicLong is used because this value is written in two places:
    // 1. The sampling thread calls takeSample to increment the time spent in this state
    // 2. The reporting thread calls extractUpdate which reads the current sum *AND* sets it to 0.
    private final AtomicLong totalMillisInState = new AtomicLong();

    // The worker that created this state.  Used to report lulls back to the worker.
    private final StreamingDataflowWorker worker;

    public StreamingModeExecutionState(
        NameContext nameContext,
        String stateName,
        MetricsContainer metricsContainer,
        ProfileScope profileScope,
        StreamingDataflowWorker worker) {
      // TODO: Take in the requesting step name and side input index for streaming.
      super(nameContext, stateName, null, null, metricsContainer, profileScope);
      this.worker = worker;
    }

    /*
     * Report the lull to the StreamingDataflowWorker that is stuck in addition to logging the
     * lull.
     */
    @Override
    public void reportLull(Thread trackedThread, long millis) {
      super.reportLull(trackedThread, millis);
      // Also report the failure to the list of pending failures to report on the worker thread
      // so that the failure gets communicated to the StreamingDataflowWorker.
      String errorMessage = getLullMessage(trackedThread, Duration.millis(millis));
      worker.addFailure(errorMessage);
    }

    /**
     * Take sample is only called by the ExecutionStateSampler thread. It is the only place that
     * increments totalMillisInState, however the reporting thread periodically calls extractUpdate
     * which will read the sum and reset it to 0, so totalMillisInState does have multiple writers.
     */
    @Override
    public void takeSample(long millisSinceLastSample) {
      totalMillisInState.addAndGet(millisSinceLastSample);
    }

    /**
     * Extract updates in the form of a {@link CounterUpdate}.
     *
     * <p>Non-final updates are extracted periodically and report the physical value as a delta.
     * This requires setting the totalMillisInState back to 0.
     *
     * <p>Final updates should never be requested from a Streaming job since the work unit never
     * completes.
     */
    @Nullable
    @Override
    public CounterUpdate extractUpdate(boolean isFinalUpdate) {
      // Streaming reports deltas, so isFinalUpdate doesn't matter, and should never be true.
      long sum = totalMillisInState.getAndSet(0);<commit_msg>No longer reporting Lulls as errors in the worker.<commit_after>/*
 * 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;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;

import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.SideInputInfo;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState;
import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StepContext;
import org.apache.beam.runners.dataflow.worker.counters.CounterFactory;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.metrics.MetricsContainer;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.FluentIterable;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** {@link DataflowExecutionContext} for use in streaming mode. */
public class StreamingModeExecutionContext extends DataflowExecutionContext<StepContext> {

  private static final Logger LOG = LoggerFactory.getLogger(StreamingModeExecutionContext.class);

  /**
   * The current user-facing key for this execution context.
   *
   * <p>This field is set to a new key upon each call to {@link #start}, which corresponds to one
   * input bundle. For non-keyed bundles, the key will be {@code null}.
   *
   * <p>This key should not be mistaken for the sharding key of the computation, which is always
   * present.
   */
  @Nullable private Object key = null;

  private final String computationId;
  private final Map<TupleTag<?>, Map<BoundedWindow, Object>> sideInputCache;

  // Per-key cache of active Reader objects in use by this process.
  private final ImmutableMap<String, String> stateNameMap;
  private final WindmillStateCache.ForComputation stateCache;

  private Windmill.WorkItem work;
  private StateFetcher stateFetcher;
  private Windmill.WorkItemCommitRequest.Builder outputBuilder;
  private UnboundedSource.UnboundedReader<?> activeReader;
  private volatile long backlogBytes;
  private final ReaderCache readerCache;

  public StreamingModeExecutionContext(
      CounterFactory counterFactory,
      String computationId,
      ReaderCache readerCache,
      Map<String, String> stateNameMap,
      WindmillStateCache.ForComputation stateCache,
      MetricsContainerRegistry<StreamingStepMetricsContainer> metricsContainerRegistry,
      DataflowExecutionStateTracker executionStateTracker,
      StreamingModeExecutionStateRegistry executionStateRegistry,
      long sinkByteLimit) {
    super(
        counterFactory,
        metricsContainerRegistry,
        executionStateTracker,
        executionStateRegistry,
        sinkByteLimit);
    this.computationId = computationId;
    this.readerCache = readerCache;
    this.sideInputCache = new HashMap<>();
    this.stateNameMap = ImmutableMap.copyOf(stateNameMap);
    this.stateCache = stateCache;
    this.backlogBytes = UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN;
  }

  /**
   * Execution states in Streaming are shared between multiple map-task executors. Thus this class
   * needs to be thread safe for multiple writers. A single stage could have have multiple executors
   * running concurrently.
   */
  public static class StreamingModeExecutionState
      extends DataflowOperationContext.DataflowExecutionState {

    // AtomicLong is used because this value is written in two places:
    // 1. The sampling thread calls takeSample to increment the time spent in this state
    // 2. The reporting thread calls extractUpdate which reads the current sum *AND* sets it to 0.
    private final AtomicLong totalMillisInState = new AtomicLong();

    // The worker that created this state.  Used to report lulls back to the worker.
    private final StreamingDataflowWorker worker;

    public StreamingModeExecutionState(
        NameContext nameContext,
        String stateName,
        MetricsContainer metricsContainer,
        ProfileScope profileScope,
        StreamingDataflowWorker worker) {
      // TODO: Take in the requesting step name and side input index for streaming.
      super(nameContext, stateName, null, null, metricsContainer, profileScope);
      this.worker = worker;
    }

    /**
     * Take sample is only called by the ExecutionStateSampler thread. It is the only place that
     * increments totalMillisInState, however the reporting thread periodically calls extractUpdate
     * which will read the sum and reset it to 0, so totalMillisInState does have multiple writers.
     */
    @Override
    public void takeSample(long millisSinceLastSample) {
      totalMillisInState.addAndGet(millisSinceLastSample);
    }

    /**
     * Extract updates in the form of a {@link CounterUpdate}.
     *
     * <p>Non-final updates are extracted periodically and report the physical value as a delta.
     * This requires setting the totalMillisInState back to 0.
     *
     * <p>Final updates should never be requested from a Streaming job since the work unit never
     * completes.
     */
    @Nullable
    @Override
    public CounterUpdate extractUpdate(boolean isFinalUpdate) {
      // Streaming reports deltas, so isFinalUpdate doesn't matter, and should never be true.
      long sum = totalMillisInState.getAndSet(0);
{'fraction_non_alphanumeric': 0.07789095266626722, 'fraction_numerical': 0.004793289394847214, 'mean_word_length': 4.302096569250318, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, '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': '6208937', 'n_tokens_mistral': 4707, 'n_tokens_neox': 4565, 'n_words': 1313}