/*
 * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 */

package io.airbyte.workers.temporal;

import io.temporal.activity.ActivityExecutionContext;
import io.temporal.client.ActivityCompletionException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public interface CancellationHandler {

  void checkAndHandleCancellation(Runnable onCancellationCallback);

  class TemporalCancellationHandler implements CancellationHandler {

    private static final Logger LOGGER = LoggerFactory.getLogger(TemporalCancellationHandler.class);

    private final ActivityExecutionContext activityContext;

    public TemporalCancellationHandler(ActivityExecutionContext activityContext) {
      this.activityContext = activityContext;
    }

    /**
     * Check for a cancellation/timeout status and run any callbacks necessary to shut down underlying
     * processes. This method should generally be run frequently within an activity so a change in
     * cancellation status is respected. This will only be effective if the cancellation type for the
     * workflow is set to
     * {@link io.temporal.activity.ActivityCancellationType#WAIT_CANCELLATION_COMPLETED}; otherwise, the
     * activity will be killed automatically as part of cleanup without removing underlying processes.
     *
     * @param onCancellationCallback a runnable that will only run when Temporal indicates the activity
     *        should be killed (cancellation or timeout).
     */
    @Override
    public void checkAndHandleCancellation(final Runnable onCancellationCallback) {
      try {
        /**
         * Heartbeat is somewhat misleading here. What it does is check the current Temporal activity's
         * context and throw an exception if the sync has been cancelled or timed out. The input to this
         * heartbeat function is available as a field in thrown ActivityCompletionExceptions, which we
         * aren't using for now.
         *
         * We should use this only as a check for the ActivityCompletionException. See
         * {@link TemporalUtils#withBackgroundHeartbeat} for where we actually send heartbeats to ensure
         * that we don't time out the activity.
         */
        activityContext.heartbeat(null);
      } catch (final ActivityCompletionException e) {
        onCancellationCallback.run();
        LOGGER.warn("Job either timed out or was cancelled.");
      }
    }

  }

}
