org.ros.concurrent.RetryingExecutorService.java Source code

Introduction

Here is the source code for org.ros.concurrent.RetryingExecutorService.java

Source

/*
 * Copyright (C) 2011 Google Inc.
 *
 * 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.ros.concurrent;

import com.google.common.collect.Maps;

import org.apache.commons.logging.Log;
import org.ros.exception.RosRuntimeException;
import org.ros.log.RosLogFactory;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Wraps an {@link ScheduledExecutorService} to execute {@link Callable}s with
 * retries.
 *
 * @author damonkohler@google.com (Damon Kohler)
 */
public class RetryingExecutorService {

    private static final Log log = RosLogFactory.getLog(RetryingExecutorService.class);

    private static final long DEFAULT_RETRY_DELAY = 5;
    private static final TimeUnit DEFAULT_RETRY_TIME_UNIT = TimeUnit.SECONDS;

    private final ScheduledExecutorService scheduledExecutorService;
    private final RetryLoop retryLoop;
    private final Map<Callable<Boolean>, CountDownLatch> latches;
    private final Map<Future<Boolean>, Callable<Boolean>> callables;
    private final CompletionService<Boolean> completionService;
    private final Object mutex;

    private long retryDelay;
    private TimeUnit retryTimeUnit;
    private boolean running;

    private class RetryLoop extends CancellableLoop {
        @Override
        public void loop() throws InterruptedException {
            Future<Boolean> future = completionService.take();
            final Callable<Boolean> callable = callables.remove(future);
            boolean retry;
            try {
                retry = future.get();
            } catch (ExecutionException e) {
                throw new RosRuntimeException(e.getCause());
            }
            if (retry) {
                log.debug("Retry requested in RetryingExecutorService.");
                scheduledExecutorService.schedule(new Runnable() {
                    @Override
                    public void run() {
                        submit(callable);
                    }
                }, retryDelay, retryTimeUnit);
            } else {
                latches.get(callable).countDown();
            }
        }
    }

    /**
     * @param scheduledExecutorService
     *          the {@link ExecutorService} to wrap
     */
    public RetryingExecutorService(ScheduledExecutorService scheduledExecutorService) {
        this.scheduledExecutorService = scheduledExecutorService;
        retryLoop = new RetryLoop();
        latches = Maps.newConcurrentMap();
        callables = Maps.newConcurrentMap();
        completionService = new ExecutorCompletionService<Boolean>(scheduledExecutorService);
        mutex = new Object();
        retryDelay = DEFAULT_RETRY_DELAY;
        retryTimeUnit = DEFAULT_RETRY_TIME_UNIT;
        running = true;
        // TODO(damonkohler): Unify this with the passed in ExecutorService.
        scheduledExecutorService.execute(retryLoop);
    }

    /**
     * Submit a new {@link Callable} to be executed. The submitted
     * {@link Callable} should return {@code true} to be retried, {@code false}
     * otherwise.
     *
     * @param callable
     *          the {@link Callable} to execute
     * @throws RejectedExecutionException
     *           if the {@link RetryingExecutorService} is shutting down
     */
    public void submit(Callable<Boolean> callable) {
        synchronized (mutex) {
            if (running) {
                Future<Boolean> future = completionService.submit(callable);
                latches.put(callable, new CountDownLatch(1));
                callables.put(future, callable);
            } else {
                throw new RejectedExecutionException();
            }
        }
    }

    /**
     * @param delay
     *          the delay in units of {@code unit}
     * @param unit
     *          the {@link TimeUnit} of the delay
     */
    public void setRetryDelay(long delay, TimeUnit unit) {
        retryDelay = delay;
        retryTimeUnit = unit;
    }

    /**
     * Stops accepting new {@link Callable}s and waits for all submitted
     * {@link Callable}s to finish within the specified timeout.
     *
     * @param timeout
     *          the timeout in units of {@code unit}
     * @param unit
     *          the {@link TimeUnit} of {@code timeout}
     * @throws InterruptedException
     */
    public void shutdown(long timeout, TimeUnit unit) throws InterruptedException {
        running = false;
        for (CountDownLatch latch : latches.values()) {
            latch.await(timeout, unit);
        }
        retryLoop.cancel();
    }
}