package ru.yandex.qe.jetty.metrics;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.atomic.AtomicLong;

import org.eclipse.jetty.util.annotation.Name;
import org.eclipse.jetty.util.thread.QueuedThreadPool;

import ru.yandex.monlib.metrics.labels.Labels;
import ru.yandex.monlib.metrics.primitives.Rate;
import ru.yandex.monlib.metrics.registry.MetricRegistry;
import ru.yandex.qe.dispenser.solomon.SolomonHolder;

public class InstrumentedQueuedThreadPool extends QueuedThreadPool {

    private static final AtomicLong POOL_INDEX = new AtomicLong(0);

    private final MetricRegistry metricRegistry;
    private final Rate rejectRate;
    private final String poolIndex;

    public InstrumentedQueuedThreadPool(@Name("maxThreads") int maxThreads, @Name("minThreads") int minThreads, @Name("idleTimeout") int idleTimeout,
                                        @Name("solomonHolder") SolomonHolder solomonHolder, @Name("queueSize") int queueSize) {
        super(maxThreads, minThreads, idleTimeout, new ArrayBlockingQueue<Runnable>(queueSize));
        this.poolIndex = String.valueOf(POOL_INDEX.incrementAndGet());
        this.metricRegistry = solomonHolder.getRootRegistry();
        this.rejectRate = this.metricRegistry.rate("jetty.pool.reject_rate", Labels.of("pool_index", this.poolIndex));
    }

    @Override
    protected void doStart() throws Exception {
        super.doStart();
        metricRegistry.lazyGaugeDouble("jetty.pool.utilization", Labels.of("pool_index", this.poolIndex), () -> {
            int used = getThreads() - getIdleThreads();
            int total = getThreads();
            if (total == 0) {
                return 1.0d;
            }
            return (double) used / (double) total;
        });
        metricRegistry.lazyGaugeDouble("jetty.pool.utilization_max", Labels.of("pool_index", this.poolIndex), () -> {
            int used = getThreads() - getIdleThreads();
            int total = getMaxThreads();
            if (total == 0) {
                return 1.0d;
            }
            return (double) used / (double) total;
        });
        metricRegistry.lazyGaugeInt64("jetty.pool.size", Labels.of("pool_index", this.poolIndex),
                () -> (long) getThreads());
    }

    @Override
    public void execute(Runnable job) {
        try {
            super.execute(job);
        } catch (RejectedExecutionException e) {
            rejectRate.inc();
            throw e;
        }
    }
}
