package com.atlassian.stash.internal.event;

import com.atlassian.event.config.EventThreadPoolConfiguration;
import com.atlassian.stash.internal.util.StackException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Date;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * A customized {@code ThreadPoolExecutor} which can be configured by {@code EventThreadPoolConfiguration} and
 * which logs its active threads when an event is {@link #rejectedExecution(Runnable, ThreadPoolExecutor) rejected}.
 * <p>
 * This class also tracks a number of statistics relevant for monitoring of event processing, in particular how
 * many events have been rejected and when an event has been rejected most recently.
 * <p>
 * Note: Because {@code ThreadPoolExecutor} doesn't expose its active threads, this class maintains a map between
 * in-flight tasks and the threads running them.
 *
 * @since 3.2.6, 3.3.2, 3.4
 */
public class EventThreadPoolExecutor extends ThreadPoolExecutor implements RejectedExecutionHandler {

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

    private final Map<Runnable, Thread> activeThreads;
    private final RejectedExecutionHandler delegate;
    private final long rejectionCooldown;
    private final AtomicLong rejectionCount;

    private volatile long lastRejectionLogged;
    private volatile long lastRejectionTimestamp;

    public EventThreadPoolExecutor(EventThreadPoolConfiguration configuration, ThreadFactory threadFactory,
                                   int maxQueueSize, long rejectionCooldown) {
        super(configuration.getCorePoolSize(), configuration.getMaximumPoolSize(),
                configuration.getKeepAliveTime(), configuration.getTimeUnit(),
                new LinkedBlockingQueue<>(maxQueueSize), threadFactory);

        this.rejectionCooldown = TimeUnit.MINUTES.toNanos(rejectionCooldown);

        activeThreads = new ConcurrentHashMap<>(configuration.getMaximumPoolSize(), 1.0F);
        delegate = getRejectedExecutionHandler(); //Get the default rejection handler before we replace it
        rejectionCount = new AtomicLong();

        setRejectedExecutionHandler(this);
    }

    public Date getLastRejectionDate() {
        return lastRejectionTimestamp > 0 ? new Date(lastRejectionTimestamp) : null;
    }

    public long getRejectionCount() {
        return rejectionCount.get();
    }

    /**
     * {@link #maybeLogThreads() Logs} stack traces for all {@link #activeThreads active} event processor threads,
     * on a configurable cooldown, and then delegates to the default rejection handler for further handling.
     *
     * @param r        the task that could not be added to the queue
     * @param executor the executor the task was submitted to
     */
    @Override
    public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
        lastRejectionTimestamp = System.currentTimeMillis();
        rejectionCount.incrementAndGet();
        maybeLogThreads();

        delegate.rejectedExecution(r, executor);
    }

    /**
     * Removes the {@code Runnable} (and, thereby, the thread that was running it) from the {@link #activeThreads
     * active threads}.
     *
     * @param task   the completed task
     * @param thrown the exception thrown by the task's processing, which may be {@code null} and is ignored regardless
     */
    @Override
    protected void afterExecute(Runnable task, Throwable thrown) {
        activeThreads.remove(task);
    }

    /**
     * Adds the {@code Runnable} and its thread to the {@link #activeThreads active thread map}.
     *
     * @param thread the thread the task will run on
     * @param task   the task about to be started
     */
    @Override
    protected void beforeExecute(Thread thread, Runnable task) {
        activeThreads.put(task, thread);
    }

    private void maybeLogThreads() {
        long check = System.nanoTime();

        if (check - lastRejectionLogged > rejectionCooldown) {
            synchronized (activeThreads) {
                if (check - lastRejectionLogged > rejectionCooldown) {
                    lastRejectionLogged = System.nanoTime();

                    log.warn("The event queue is full. Stacks for the processing threads follow:");
                    for (Thread thread : activeThreads.values()) {
                        //Leverage the logger's ability to handle stack traces when they're on exceptions
                        log.warn("Stack trace for {}", thread.getName(), new StackException(thread));
                    }
                }
            }
        }
    }
}
