/*
 * 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.List;
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.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;
  }

  public void logAndProcessFailureBatch(
      String computationId,
      List<ExecutableWork> executableWorks,
      Throwable t,
      Consumer<Work> onInvalidWork)
      throws Throwable {
    List<ExecutableWork> worksToRetryLocally = new java.util.ArrayList<>();

    for (ExecutableWork executableWork : executableWorks) {
      switch (evaluateRetry(computationId, executableWork.work(), t)) {
        case DO_NOT_RETRY:
          // Consider the item invalid. It will eventually be retried by Windmill if it still needs
          // to be processed.
          onInvalidWork.accept(executableWork.work());
          break;
        case RETRY_LOCALLY:
          // Try again after some delay and at the end of the queue to avoid a tight loop.
          worksToRetryLocally.add(executableWork);
          break;
        case RETHROW_THROWABLE:
          throw t;
      }
    }

    executeWithDelay(worksToRetryLocally);
  }

  private void executeWithDelay(List<ExecutableWork> worksToRetryLocally) {
    if (!worksToRetryLocally.isEmpty()) {
      // Sleep ONCE for the entire batch delay to avoid sequential thread blocks
      Uninterruptibles.sleepUninterruptibly(retryLocallyDelayMs, TimeUnit.MILLISECONDS);
      for (ExecutableWork ew : worksToRetryLocally) {
        workUnitExecutor.forceExecute(ew, ew.work().getSerializedWorkItemSize());
      }
    }
  }

  private String tryToDumpHeap() {
    return heapDumper
        .dumpAndGetHeap()
        .map(heapDump -> "written to '" + heapDump + "'")
        .orElseGet(() -> "not written");
  }

  private enum RetryEvaluation {
    DO_NOT_RETRY,
    RETRY_LOCALLY,
    RETHROW_THROWABLE,
  }

  private RetryEvaluation evaluateRetry(String computationId, Work work, Throwable t) {
    if (work.isFailed()) {
      LOG.debug(
          "Execution of work for computation '{}' on sharding key '{}' failed. "
              + "Work is already marked as failed, not retrying locally.",
          computationId,
          work.getWorkItem().getShardingKey());
      return RetryEvaluation.DO_NOT_RETRY;
    }
    @Nullable final Throwable cause = t.getCause();
    Throwable parsedException = (t instanceof UserCodeException && cause != null) ? cause : t;

    LastExceptionDataProvider.reportException(parsedException);
    LOG.debug("Failed work: {}", work);
    Duration elapsedTimeSinceStart = new Duration(work.getStartTime(), clock.get());
    if (isOutOfMemoryError(parsedException)) {
      String heapDump = tryToDumpHeap();
      LOG.error(
          "Execution of work for computation '{}' for sharding key '{}' failed with out-of-memory. "
              + "Work will not be retried locally. Heap dump {}.",
          computationId,
          work.getWorkItem().getShardingKey(),
          heapDump,
          parsedException);
      return RetryEvaluation.RETHROW_THROWABLE;
    }

    if (!failureTracker.trackFailure(computationId, work.getWorkItem(), parsedException)) {
      LOG.error(
          "Execution of work for computation '{}' on sharding key '{}' failed with uncaught exception, "
              + "and Windmill indicated not to retry locally.",
          computationId,
          work.getWorkItem().getShardingKey(),
          parsedException);
      return RetryEvaluation.DO_NOT_RETRY;
    }
    if (elapsedTimeSinceStart.isLongerThan(MAX_LOCAL_PROCESSING_RETRY_DURATION)) {
      LOG.error(
          "Execution of work for computation '{}' for sharding key '{}' failed with uncaught exception, "
              + "and it will not be retried locally because the elapsed time since start {} "
              + "exceeds {}.",
          computationId,
          work.getWorkItem().getShardingKey(),
          elapsedTimeSinceStart,
          MAX_LOCAL_PROCESSING_RETRY_DURATION,
          parsedException);
      return RetryEvaluation.DO_NOT_RETRY;
    }
    LOG.error(
        "Execution of work for computation '{}' on sharding key '{}' failed with uncaught exception. "
            + "Work will be retried locally.",
        computationId,
        work.getWorkItem().getShardingKey(),
        parsedException);
    return RetryEvaluation.RETRY_LOCALLY;
  }
}
