com.facebook.concurrency.UnstoppableExecutorServiceCore Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2012 Facebook, 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 com.facebook.concurrency;

import com.facebook.collections.ListMapper;
import com.facebook.collectionsbase.Mapper;
import org.joda.time.DateTimeUtils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * core class that Unstoppable[Scheduled]ExecutorService delegates termination
 * methods to in order to guard shutdown
 */
class UnstoppableExecutorServiceCore {
  private final AtomicInteger remaining = new AtomicInteger(0);
  private volatile boolean isShutdown = false;

  public List registerRunnableList(List taskList) {
    if (isShutdown) {
      throw new RejectedExecutionException("executor shutdown already");
    }

    List result = new ArrayList();

    for (Runnable task : taskList) {
      result.add(new TrackedRunnableImpl(task));
    }

    return result;
  }

  public  List> registerCallableList(
    Collection> taskList
  ) {
    if (isShutdown()) {
      throw new RejectedExecutionException("executor shutdown already");
    }

    List> result = new ArrayList>();

    for (Callable task : taskList) {
      result.add(new TrackedCallableImpl(task));
    }

    return result;
  }

  public TrackedRunnable registerTask(Runnable task) {
    if (isShutdown()) {
      throw new RejectedExecutionException("executor shutdown already");
    }

    return new TrackedRunnableImpl(task);
  }

  public  TrackedCallable registerTask(final Callable task) {
    if (isShutdown()) {
      throw new RejectedExecutionException("executor shutdown already");
    }

    return new TrackedCallableImpl(task);
  }

  private void decrementRemaining() {
    if (remaining.decrementAndGet() == 0) {
      synchronized (remaining) {
        remaining.notifyAll();
      }
    }
  }

  public synchronized void shutdown() {
    if (isShutdown) {
      return;
    }

    isShutdown = true;
  }

  public List shutdownNow() {
    // for now, shutdownNow() is equivalent to shutdown()
    shutdown();

    // TODO: we can track started tasks and actually interrupt them

    return Collections.emptyList();
  }

  public boolean isShutdown() {
    return isShutdown;
  }

  public boolean isTerminated() {
    assert remaining.get() >= 0;

    return remaining.get() == 0;
  }

  public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
    if (!isShutdown) {
      return false;
    }

    long start = DateTimeUtils.currentTimeMillis();

    synchronized (remaining) {
      while (remaining.get() > 0) {
        // timed wait due to likely lost notifications, so relatively short also
        remaining.wait(50);

        long elapsedMillis = DateTimeUtils.currentTimeMillis() - start;

        if (elapsedMillis > unit.toMillis(timeout)) {
          return false;
        }
      }
    }

    return true;
  }

  public  List> trackFutureList(
    List> futureList, List completableList
  ) {
    return ListMapper.map(futureList, new FutureMapper(completableList));
  }

  public  Future trackFuture(Future future, Completable task) {
    return new TrackedFuture(future, task);
  }

  public  ScheduledFuture trackScheduledFuture(
    ScheduledFuture future, Completable task
  ) {
    return new TrackedScheduledFuture(future, task);
  }

  private class TrackedRunnableImpl implements TrackedRunnable {
    private final Runnable delegate;
    private final AtomicBoolean hasCompleted = new AtomicBoolean(false);

    private TrackedRunnableImpl(Runnable delegate) {
      this.delegate = delegate;
      remaining.incrementAndGet();
    }

    @Override
    public void run() {
      try {
        delegate.run();
      } finally {
        complete();
      }
    }

    public void complete() {
      if (hasCompleted.compareAndSet(false, true)) {
        decrementRemaining();
      }
    }
  }

  private class TrackedCallableImpl implements TrackedCallable {

    private final Callable delegate;
    private final AtomicBoolean hasCompleted = new AtomicBoolean(false);

    private TrackedCallableImpl(Callable delegate) {
      this.delegate = delegate;
      remaining.incrementAndGet();
    }

    @Override
    public V call() throws Exception {
      try {
        return delegate.call();
      } finally {
        complete();
      }
    }

    public void complete() {
      if (hasCompleted.compareAndSet(false, true)) {
        decrementRemaining();
      }
    }
  }

  private class TrackedFuture extends WrappedFuture {
    private final Completable task;

    private TrackedFuture(Future delegate, Completable task) {
      super(delegate);
      this.task = task;
    }

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
      task.complete();

      return super.cancel(mayInterruptIfRunning);
    }
  }

  private class TrackedScheduledFuture extends WrappedScheduledFuture {
    private final Completable task;

    private TrackedScheduledFuture(
      ScheduledFuture delegate, Completable task
    ) {
      super(delegate);
      this.task = task;
    }

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
      task.complete();

      return super.cancel(mayInterruptIfRunning);
    }
  }

  private class FutureMapper implements Mapper, Future> {
    private final List completableList;
    private int index = 0;

    private FutureMapper(List completableList) {
      this.completableList = completableList;
    }

    @Override
    public Future map(Future input) {
      TrackedFuture trackedFuture = new TrackedFuture(input, completableList.get(index));

      index++;

      return trackedFuture;
    }
  }
}