TaskRetryUtil.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2020-2021 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * 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 org.sleuthkit.autopsy.threadutils;

import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Level;
import java.util.logging.Logger;

public class TaskRetryUtil {

    private static final AtomicLong totalTasks = new AtomicLong();
    private static final AtomicLong totalTaskRetries = new AtomicLong();
    private static final AtomicLong totalTaskAttemptTimeOuts = new AtomicLong();
    private static final AtomicLong totalFailedTasks = new AtomicLong();

    public static class TaskAttempt {

        private final Long delay;
        private final Long timeOut;
        private final TimeUnit timeUnit;

        public TaskAttempt() {
            this.delay = 0L;
            this.timeOut = 0L;
            this.timeUnit = TimeUnit.SECONDS;
        }

        public TaskAttempt(Long delay, TimeUnit timeUnit) {
            if (delay == null || delay < 0) {
                throw new IllegalArgumentException(String.format("Argument for delay parameter = %d, must be zero or any positive integer", delay));
            }
            if (timeUnit == null) {
                throw new IllegalArgumentException("Argument for timeUnit parameter is null");
            }
            this.delay = delay;
            this.timeOut = 0L;
            this.timeUnit = timeUnit;
        }

        public TaskAttempt(Long delay, Long timeOut, TimeUnit timeUnit) {
            if (delay == null || delay < 0) {
                throw new IllegalArgumentException(String.format("Argument for delay parameter = %d, must be zero or any positive integer", delay));
            }
            if (timeOut == null || timeOut < 0) {
                throw new IllegalArgumentException(String.format("Argument for timeOut parameter = %d, must be zero or any positive integer", delay));
            }
            if (timeUnit == null) {
                throw new IllegalArgumentException("Argument for timeUnit parameter is null");
            }
            this.delay = delay;
            this.timeOut = timeOut;
            this.timeUnit = timeUnit;
        }

        public Long getDelay() {
            return delay;
        }

        public Long getTimeout() {
            return timeOut;
        }

        public TimeUnit getTimeUnit() {
            return timeUnit;
        }

    }

    public interface Terminator {

        boolean stopTaskAttempts();

    }

    public static <T> T attemptTask(Callable<T> task, List<TaskAttempt> attempts, ScheduledThreadPoolExecutor executor, Terminator terminator, Logger logger, String taskDesc) throws InterruptedException {
        /*
         * Attempt the task.
         */
        T result = null;
        String taskDescForLog = taskDesc != null ? taskDesc : "Task";
        int attemptCounter = 0;
        if (attempts.size() > 0) {
            totalTasks.incrementAndGet();
        }
        while (result == null && attemptCounter < attempts.size()) {
            if (terminator != null && terminator.stopTaskAttempts()) {
                if (logger != null) {
                    logger.log(Level.WARNING, String.format("Attempts to execute '%s' terminated ", taskDescForLog));
                }
                break;
            }
            TaskAttempt attempt = attempts.get(attemptCounter);
            if (attemptCounter > 0) {
                totalTaskRetries.incrementAndGet();
            }
            ScheduledFuture<T> future = executor.schedule(task, attempt.getDelay(), attempt.getTimeUnit());
            try {
                result = future.get(attempt.getDelay() + attempt.getTimeout(), attempt.getTimeUnit());
            } catch (InterruptedException ex) {
                if (logger != null) {
                    logger.log(Level.SEVERE, String.format("Interrupted executing '%s'", taskDescForLog), ex);
                }
                throw ex;
            } catch (ExecutionException ex) {
                if (logger != null) {
                    logger.log(Level.SEVERE, String.format("Error executing '%s'", taskDescForLog), ex);
                }
            } catch (TimeoutException ex) {
                if (logger != null) {
                    logger.log(Level.SEVERE, String.format("Time out executing '%s'", taskDescForLog), ex);
                }
                totalTaskAttemptTimeOuts.incrementAndGet();
                future.cancel(true);
            }
            ++attemptCounter;
        }

        /*
         * If the task required more than one attempt, log it.
         */
        if (logger != null && attemptCounter > 1) {
            if (result != null) {
                logger.log(Level.WARNING, String.format("'%s' succeeded after %d attempts", taskDescForLog, attemptCounter));
            } else {
                logger.log(Level.SEVERE, String.format("'%s' failed after %d attempts", taskDescForLog, attemptCounter));
            }
        }

        /*
         * If the task failed, count it as a failed task.
         */
        if (result == null) {
            if (terminator == null || !terminator.stopTaskAttempts()) {
                totalFailedTasks.incrementAndGet();
            }
        }

        return result;
    }

    public static long getTotalTasksCount() {
        return totalTasks.get();
    }

    public static long getTotalTaskRetriesCount() {
        return totalTaskRetries.get();
    }

    public static long getTotalTaskAttemptTimeOutsCount() {
        return totalTaskAttemptTimeOuts.get();
    }

    public static long getTotalFailedTasksCount() {
        return totalFailedTasks.get();
    }

    private TaskRetryUtil() {
    }

}