package com.atlassian.bitbucket.mesh.concurrent;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.util.concurrent.*;

import static java.util.Objects.requireNonNull;

/**
 * A {@link RejectedExecutionHandler} that blocks the caller that is submitting a task until the task can be queued
 * or until a timeout occurs. If a timeout occurs, a {@link RejectedExecutionException} is thrown.
 */
public class CallerBlocksPolicy implements RejectedExecutionHandler {

    private static final Logger log = LoggerFactory.getLogger(CallerBlocksPolicy.class);

    private final long timeoutMs;

    public CallerBlocksPolicy(@Nonnull Duration timeout) {
        timeoutMs = requireNonNull(timeout, "timeout").toMillis();
    }

    @Override
    public void rejectedExecution(@Nonnull Runnable runnable, @Nonnull ThreadPoolExecutor executor) {
        if (executor.isShutdown()) {
            throw new RejectedExecutionException("Executor has been shut down");
        }
        try {
            log.debug("All pool threads are busy; queuing task with {} ms timeout", timeoutMs);

            long start = System.nanoTime();
            if (executor.getQueue().offer(runnable, timeoutMs, TimeUnit.MILLISECONDS)) {
                log.debug("Task was queued after {} ms", TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start));
            } else {
                throw new RejectedExecutionException("Timed out queuing task; all threads are busy");
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RejectedExecutionException("Interrupted while queuing task", e);
        }
    }
}
