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.windmill.work.processing.failures;
  
  import java.util.concurrent.TimeUnit;
  import java.util.function.Consumer;
  import java.util.function.Supplier;
  import javax.annotation.Nullable;
  import javax.annotation.concurrent.ThreadSafe;
  import org.apache.beam.runners.dataflow.worker.KeyTokenInvalidException;
  import org.apache.beam.runners.dataflow.worker.WorkItemCancelledException;
  import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
  import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
  import org.apache.beam.runners.dataflow.worker.streaming.Work;
  import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
  import org.apache.beam.sdk.annotations.Internal;
  import org.apache.beam.sdk.util.UserCodeException;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /** Processes a failure that occurs during user processing of {@link Work}. */
  @ThreadSafe
  @Internal
  public final class WorkFailureProcessor {
    private static final Logger LOG = LoggerFactory.getLogger(WorkFailureProcessor.class);
    private static final Duration MAX_LOCAL_PROCESSING_RETRY_DURATION = Duration.standardMinutes(5);
    private static final int DEFAULT_RETRY_LOCALLY_MS = 10000;
  
    private final BoundedQueueExecutor workUnitExecutor;
    private final FailureTracker failureTracker;
    private final HeapDumper heapDumper;
    private final Supplier<Instant> clock;
    private final int retryLocallyDelayMs;
  
    private WorkFailureProcessor(
        BoundedQueueExecutor workUnitExecutor,
        FailureTracker failureTracker,
        HeapDumper heapDumper,
        Supplier<Instant> clock,
        int retryLocallyDelayMs) {
      this.workUnitExecutor = workUnitExecutor;
      this.failureTracker = failureTracker;
      this.heapDumper = heapDumper;
      this.clock = clock;
      this.retryLocallyDelayMs = retryLocallyDelayMs;
    }
  
    public static WorkFailureProcessor create(
        BoundedQueueExecutor workUnitExecutor,
        FailureTracker failureTracker,
        HeapDumper heapDumper,
        Supplier<Instant> clock) {
      return new WorkFailureProcessor(
          workUnitExecutor, failureTracker, heapDumper, clock, DEFAULT_RETRY_LOCALLY_MS);
    }
  
    @VisibleForTesting
    public static WorkFailureProcessor forTesting(
        BoundedQueueExecutor workUnitExecutor,
        FailureTracker failureTracker,
        HeapDumper heapDumper,
        Supplier<Instant> clock,
        int retryLocallyDelayMs) {
      return new WorkFailureProcessor(
          workUnitExecutor,
          failureTracker,
          heapDumper,
          clock,
          retryLocallyDelayMs >= 0 ? retryLocallyDelayMs : DEFAULT_RETRY_LOCALLY_MS);
    }
  
    /** Returns whether an exception was caused by a {@link OutOfMemoryError}. */
    private static boolean isOutOfMemoryError(@Nullable Throwable t) {
      while (t != null) {
        if (t instanceof OutOfMemoryError) {
          return true;
        }
        t = t.getCause();
      }
      return false;
    }
  
    /**
     * Processes failures caused by thrown exceptions that occur during execution of {@link Work}. May
     * attempt to retry execution of the {@link Work} or drop it if it is invalid.
     */
    public void logAndProcessFailure(
        String computationId,
        ExecutableWork executableWork,
        Throwable t,
        Consumer<Work> onInvalidWork) {
      if (shouldRetryLocally(computationId, executableWork.work(), t)) {
        // Try again after some delay and at the end of the queue to avoid a tight loop.
        executeWithDelay(retryLocallyDelayMs, executableWork);
      } else {
        // Consider the item invalid. It will eventually be retried by Windmill if it still needs to
        // be processed.
        onInvalidWork.accept(executableWork.work());
      }
    }
  
    private String tryToDumpHeap() {
      return heapDumper
          .dumpAndGetHeap()
          .map(heapDump -> "written to '" + heapDump + "'")
          .orElseGet(() -> "not written");
    }
  
    private void executeWithDelay(long delayMs, ExecutableWork executableWork) {
      Uninterruptibles.sleepUninterruptibly(delayMs, TimeUnit.MILLISECONDS);
      workUnitExecutor.forceExecute(
          executableWork, executableWork.work().getSerializedWorkItemSize());
    }
  
    private boolean shouldRetryLocally(String computationId, Work work, Throwable t) {
      @Nullable final Throwable cause = t.getCause();
      Throwable parsedException = (t instanceof UserCodeException && cause != null) ? cause : t;
      if (KeyTokenInvalidException.isKeyTokenInvalidException(parsedException)) {
        LOG.debug(
            "Execution of work for computation '{}' on key '{}' failed due to token expiration. "
                + "Work will not be retried locally.",
            computationId,
            work.getWorkItem().getKey().toStringUtf8());
      } else if (WorkItemCancelledException.isWorkItemCancelledException(parsedException)) {
        LOG.debug(
            "Execution of work for computation '{}' on key '{}' failed. "
                + "Work will not be retried locally.",
            computationId,
            work.getWorkItem().getShardingKey());
      } else {
        LastExceptionDataProvider.reportException(parsedException);
        LOG.debug("Failed work: {}", work);
        Duration elapsedTimeSinceStart = new Duration(work.getStartTime(), clock.get());
        if (!failureTracker.trackFailure(computationId, work.getWorkItem(), parsedException)) {
          LOG.error(
              "Execution of work for computation '{}' on key '{}' failed with uncaught exception, "
                  + "and Windmill indicated not to retry locally.",
              computationId,
              work.getWorkItem().getKey().toStringUtf8(),
              parsedException);
        } else if (isOutOfMemoryError(parsedException)) {
          String heapDump = tryToDumpHeap();
          LOG.error(
              "Execution of work for computation '{}' for key '{}' failed with out-of-memory. "
                  + "Work will not be retried locally. Heap dump {}.",
              computationId,
              work.getWorkItem().getKey().toStringUtf8(),
              heapDump,
              parsedException);
        } else if (elapsedTimeSinceStart.isLongerThan(MAX_LOCAL_PROCESSING_RETRY_DURATION)) {
          LOG.error(
              "Execution of work for computation '{}' for key '{}' failed with uncaught exception, "
                  + "and it will not be retried locally because the elapsed time since start {} "
                  + "exceeds {}.",
              computationId,
              work.getWorkItem().getKey().toStringUtf8(),
              elapsedTimeSinceStart,
              MAX_LOCAL_PROCESSING_RETRY_DURATION,
              parsedException);
        } else {
          LOG.error(
              "Execution of work for computation '{}' on key '{}' failed with uncaught exception. "
                  + "Work will be retried locally.",
              computationId,
              work.getWorkItem().getKey().toStringUtf8(),
              parsedException);
          return true;
        }
      }
  
      return false;
    }
  }
  