com.palantir.common.concurrent.MultiplexingCompletionService Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of atlasdb-commons Show documentation
Palantir open source project
There is a newer version: 0.1152.0
/*
 * (c) Copyright 2018 Palantir Technologies Inc. All rights reserved.
 *
 * 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.palantir.common.concurrent;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.palantir.common.streams.KeyedStream;
import com.palantir.logsafe.exceptions.SafeIllegalStateException;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RunnableFuture;
import java.util.concurrent.TimeUnit;

/**
 * A MultiplexingCompletionService is much like a {@link java.util.concurrent.ExecutorCompletionService}, but
 * supports multiple delegate {@link ExecutorService}s feeding in to a single shared {@link BlockingQueue}.
 * {@link #poll()} operations will retrieve the results of computations that are done first (regardless of which
 * actual underlying executor service they may have been scheduled on).
 *
 * Maintaining separate executors may see application in improving monitoring and bounding of thread pools that
 * have several distinct use cases.
 *
 * @param  key type
 * @param  return type of tasks that are to be submitted
 */
public final class MultiplexingCompletionService {
    private final ImmutableMap executors;
    private final BlockingQueue>> taskQueue;

    private MultiplexingCompletionService(
            ImmutableMap executors, BlockingQueue>> taskQueue) {
        this.executors = executors;
        this.taskQueue = taskQueue;
    }

    public static  MultiplexingCompletionService create(Map executors) {
        return new MultiplexingCompletionService<>(ImmutableMap.copyOf(executors), new LinkedBlockingQueue<>());
    }

    public static  MultiplexingCompletionService createFromCheckedExecutors(
            Map executors) {
        return create(KeyedStream.stream(executors)
                .map(CheckedRejectionExecutorService::getUnderlyingExecutor)
                .collectToMap());
    }

    /**
     * Submits a task to be run on a specific executor.
     *
     * @param key to identify which executor the task should be run on
     * @param task to be run on the relevant executor
     * @return future associated with submitting the task to the correct executor
     *
     * @throws IllegalStateException if the key provided is not associated with any executor
     * @throws CheckedRejectedExecutionException if the task could not be submitted to the relevant executor
     */
    public Future> submit(K key, Callable task) throws CheckedRejectedExecutionException {
        ExecutorService targetExecutor = executors.get(key);
        if (targetExecutor == null) {
            throw new SafeIllegalStateException(
                    "The key provided to the multiplexing completion service doesn't exist!");
        }
        return submitAndPrepareForQueueing(targetExecutor, key, task);
    }

    public Future> poll() {
        return taskQueue.poll();
    }

    public Future> poll(long timeout, TimeUnit unit) throws InterruptedException {
        return taskQueue.poll(timeout, unit);
    }

    private Future> submitAndPrepareForQueueing(ExecutorService delegate, K key, Callable callable)
            throws CheckedRejectedExecutionException {
        FutureTask> futureTask = new FutureTask<>(() -> Maps.immutableEntry(key, callable.call()));
        try {
            delegate.submit(new QueueTask(futureTask), null);
        } catch (RejectedExecutionException e) {
            throw new CheckedRejectedExecutionException(e);
        }
        return futureTask;
    }

    private final class QueueTask extends FutureTask> {
        private final RunnableFuture> runnable;

        private QueueTask(RunnableFuture> runnable) {
            super(runnable, null);
            this.runnable = runnable;
        }

        @Override
        protected void done() {
            taskQueue.add(runnable);
        }
    }
}