RetryingExecutorService.java

Go to the documentation of this file.

/*
 * 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.apache.commons.logging.LogFactory;
import org.ros.exception.RosRuntimeException;

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;

public class RetryingExecutorService {

  private static final boolean DEBUG = false;
  private static final Log log = LogFactory.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) {
        if (DEBUG) {
          log.info("Retry requested.");
        }
        scheduledExecutorService.schedule(new Runnable() {
          @Override
          public void run() {
            submit(callable);
          }
        }, retryDelay, retryTimeUnit);
      } else {
        latches.get(callable).countDown();
      }
    }
  }

  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);
  }

  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();
      }
    }
  }

  public void setRetryDelay(long delay, TimeUnit unit) {
    retryDelay = delay;
    retryTimeUnit = unit;
  }

  public void shutdown(long timeout, TimeUnit unit) throws InterruptedException {
    running = false;
    for (CountDownLatch latch : latches.values()) {
      latch.await(timeout, unit);
    }
    retryLoop.cancel();
  }
}