org.infinispan.executors.SemaphoreCompletionService Maven / Gradle / Ivy

Go to download
package org.infinispan.executors;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.TimeUnit;

import org.infinispan.IllegalLifecycleStateException;
import org.infinispan.util.concurrent.WithinThreadExecutor;
import org.infinispan.util.logging.Log;
import org.infinispan.util.logging.LogFactory;

/**
 * Executes tasks in the given executor, but never has more than {@code maxConcurrentTasks} tasks running at the same time.
 *
 * @author Dan Berindei
 * @since 7.2
 */
public class SemaphoreCompletionService implements CompletionService {
   private static final Log log = LogFactory.getLog(SemaphoreCompletionService.class);
   private static final boolean trace = log.isTraceEnabled();

   private final Executor executor;
   private final CustomSemaphore semaphore;
   private final BlockingQueue queue;
   private final BlockingQueue completionQueue = new LinkedBlockingQueue<>();
   private final boolean blocking;

   public SemaphoreCompletionService(Executor executor, int maxConcurrentTasks) {
      this.executor = executor;
      this.semaphore = new CustomSemaphore(maxConcurrentTasks);
      // Users of WithinThreadExecutor expect the tasks to execute in the thread that submitted them
      // But with a LinkedBlockingQueue, they could execute on any thread calling backgroundTaskFinished.
      this.blocking = executor instanceof WithinThreadExecutor;
      this.queue = blocking ? new SynchronousQueue<>() : new LinkedBlockingQueue<>();
   }

   public List> drainCompletionQueue() {
      List list = new ArrayList();
      completionQueue.drainTo(list);
      return list;
   }

   /**
    * When stopping, cancel any queued tasks.
    */
   public void cancelQueuedTasks() {
      ArrayList queuedTasks = new ArrayList();
      queue.drainTo(queuedTasks);
      for (QueueingTask task : queuedTasks) {
         task.cancel(false);
      }
   }

   /**
    * Called from a task to remove the permit that would otherwise be freed when the task finishes running
    *
    * When the asynchronous part of the task finishes, it must call {@link #backgroundTaskFinished(Callable)}
    * to make the permit available again.
    */
   public void continueTaskInBackground() {
      if (trace) log.tracef("Moving task to background, available permits %d", semaphore.availablePermits());
      // Prevent other tasks from running with this task's permit
      semaphore.removePermit();
   }

   /**
    * Signal that a task that called {@link #continueTaskInBackground()} has finished and
    * optionally execute another task on the just-freed thread.
    */
   public Future backgroundTaskFinished(final Callable cleanupTask) {
      QueueingTask futureTask = null;
      if (cleanupTask != null) {
         if (trace) log.tracef("Background task finished, executing cleanup task");
         futureTask = new QueueingTask(cleanupTask);
         executor.execute(futureTask);
      } else {
         semaphore.release();
         if (trace) log.tracef("Background task finished, available permits %d", semaphore.availablePermits());
         executeFront();
      }
      return futureTask;
   }

   @Override
   public Future submit(final Callable task) {
      QueueingTask futureTask = new QueueingTask(task);
      return doSubmit(futureTask);
   }

   @Override
   public Future submit(final Runnable task, T result) {
      QueueingTask futureTask = new QueueingTask(task, result);
      return doSubmit(futureTask);
   }

   private Future doSubmit(QueueingTask futureTask) {
      if (blocking) {
         try {
            semaphore.acquire();
         } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalLifecycleStateException();
         }
         try {
            futureTask.run();
         } finally {
            semaphore.release();
         }
      } else {
         queue.add(futureTask);
         if (trace) log.tracef("New task submitted, tasks in queue %d, available permits %d", queue.size(),
               semaphore.availablePermits());
         executeFront();
      }
      return futureTask;
   }

   private void executeFront() {
      while (!queue.isEmpty() && semaphore.tryAcquire()) {
         QueueingTask next = queue.poll();
         if (next != null) {
            // Execute the task, and it will release the permit when it finishes
            executor.execute(next);
            // Only execute one task, if there are other tasks and permits available, they will be scheduled
            // to be executed either by the threads that released the permits or by the threads that added
            // the tasks.
            return;
         } else {
            // Perform another iteration, in case someone adds a task and skips executing it just before
            // we release the permit
            semaphore.release();
         }
      }
   }

   @Override
   public Future take() throws InterruptedException {
      return completionQueue.take();
   }

   @Override
   public Future poll() {
      return completionQueue.poll();
   }

   @Override
   public Future poll(long timeout, TimeUnit unit) throws InterruptedException {
      return completionQueue.poll(timeout, unit);
   }

   private class QueueingTask extends FutureTask {

      public QueueingTask(Callable task) {
         super(task);
      }

      public QueueingTask(Runnable runnable, Object result) {
         super(runnable, (T) result);
      }

      @Override
      public void run() {
         try {
            QueueingTask next = this;
            do {
               next.runInternal();

               // Don't run another task if the current task called startBackgroundTask()
               // and there are no more permits available
               if (semaphore.availablePermits() < 0)
                  break;
               next = queue.poll();
            } while (next != null);
         } finally {
            semaphore.release();

            // In case we just got a new task between queue.poll() and semaphore.release()
            if (!queue.isEmpty()) {
               executeFront();
            }
         }
      }

      private void runInternal() {
         try {
            if (trace) log.tracef("Task started, tasks in queue %d, available permits %d", queue.size(), semaphore.availablePermits());
            super.run();
         } finally {
            completionQueue.offer(this);
            if (trace) log.tracef("Task finished, tasks in queue %d, available permits %d", queue.size(), semaphore.availablePermits());
         }
      }
   }

   /**
    * Extend {@code Semaphore} to expose the {@code reducePermits(int)} method.
    */
   private static class CustomSemaphore extends Semaphore {
      CustomSemaphore(int permits) {
         super(permits);
      }

      void removePermit() {
         super.reducePermits(1);
      }
   }
}