package com.hpcloud.util.concurrent;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

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

import com.codahale.metrics.Gauge;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.RatioGauge;
import com.codahale.metrics.Timer;

/**
 * ThreadPoolExecutor instrumented with extensive metrics.
 * 
 * @author Jonathan Halterman
 */
public class InstrumentedThreadPoolExecutor extends ThreadPoolExecutor {
  private final Logger log = LoggerFactory.getLogger(getClass());
  private final String name;
  private final Meter requestRate;
  private final Meter rejectedRate;
  private final Timer executionTimer;
  private final ThreadLocal<Long> startTime = new ThreadLocal<Long>();
  private final Set<ExecutionListener> listeners = new HashSet<ExecutionListener>();

  InstrumentedThreadPoolExecutor(MetricRegistry metricRegistry, String name, int corePoolSize,
      int maximumPoolSize, long keepAliveTime, TimeUnit unit,
      final BlockingQueue<Runnable> workQueue, ThreadFactory factory) {
    super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, factory);
    this.name = name;
    requestRate = metricRegistry.meter(MetricRegistry.name(getClass(), "request"));
    rejectedRate = metricRegistry.meter(MetricRegistry.name(getClass(), "rejected"));
    executionTimer = metricRegistry.timer(MetricRegistry.name(getClass(), "execution"));
    metricRegistry.register(MetricRegistry.name(getClass(), "queue.size"), new Gauge<Integer>() {
      @Override
      public Integer getValue() {
        return getQueue().size();
      }
    });
    metricRegistry.register(MetricRegistry.name(getClass(), "threads.count"), new Gauge<Integer>() {
      @Override
      public Integer getValue() {
        return getPoolSize();
      }
    });
    metricRegistry.register(MetricRegistry.name(getClass(), "threads.active"),
        new Gauge<Integer>() {
          @Override
          public Integer getValue() {
            return getActiveCount();
          }
        });
    metricRegistry.register(MetricRegistry.name(getClass(), "threads.idle"), new Gauge<Integer>() {
      @Override
      public Integer getValue() {
        return getPoolSize() - getActiveCount();
      }
    });
    metricRegistry.register(MetricRegistry.name(getClass(), "threads.percent-active"),
        new RatioGauge() {
          @Override
          protected Ratio getRatio() {
            return Ratio.of(getPoolSize(), getActiveCount());
          }
        });

    setRejectedExecutionHandler(new RejectedExecutionHandler() {
      @Override
      public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
        rejectedRate.mark();
        if (!workQueue.offer(r))
          log.warn("Thread pool {} rejected work.", InstrumentedThreadPoolExecutor.this.name);
        throw new RejectedExecutionException();
      }
    });
  }

  public void addListener(ExecutionListener listener) {
    listeners.add(listener);
  }

  @Override
  public void execute(Runnable r) {
    requestRate.mark();
    super.execute(r);
  }

  public void removeListener(ExecutionListener listener) {
    listeners.remove(listener);
  }

  @Override
  protected void afterExecute(Runnable r, Throwable t) {
    for (ExecutionListener listener : listeners)
      listener.afterExecute(r, t);
    long duration = System.nanoTime() - startTime.get();
    executionTimer.update(duration, TimeUnit.NANOSECONDS);
  }

  @Override
  protected void beforeExecute(Thread t, Runnable r) {
    for (ExecutionListener listener : listeners)
      listener.beforeExecute(t, r);
    startTime.set(System.nanoTime());
  }
}

@@ -1545,24 +1545,30 @@
 (getClass(),
+ name,
  %22request%22))
@@ -1640,16 +1640,22 @@
 Class(),
+ name,
  %22reject
@@ -1734,16 +1734,22 @@
 Class(),
+ name,
  %22execut
@@ -1815,16 +1815,22 @@
 Class(),
+ name,
  %22queue.
@@ -1832,24 +1832,32 @@
 ueue.size%22),
+%0A       
  new Gauge%3CI
@@ -1866,32 +1866,36 @@
 eger%3E() %7B%0A      
+ 
+   
 @Override%0A      
@@ -1880,32 +1880,36 @@
       @Override%0A
+    
       public Int
@@ -1926,32 +1926,36 @@
 lue() %7B%0A        
+  
+  
 return getQueue(
@@ -1962,38 +1962,46 @@
 ).size();%0A      
-%7D%0A
+    %7D%0A  
     
+  
 %7D);%0A    metricRe
@@ -2039,32 +2039,38 @@
 name(getClass(),
+ name,
  %22threads.count%22
@@ -2071,16 +2071,24 @@
 count%22),
+%0A       
  new Gau
@@ -2101,32 +2101,36 @@
 eger%3E() %7B%0A      
+    
 @Override%0A      
@@ -2121,32 +2121,36 @@
 @Override%0A      
+   
+ 
 public Integer g
@@ -2161,32 +2161,36 @@
 lue() %7B%0A        
+    
 return getPoolSi
@@ -2193,34 +2193,42 @@
 olSize();%0A      
+ 
-%7D%0A
+   %7D%0A    
     %7D);%0A    metr
@@ -2270,32 +2270,38 @@
 name(getClass(),
+ name,
  %22threads.active
@@ -2505,32 +2505,38 @@
 name(getClass(),
+ name,
  %22threads.idle%22)
@@ -2536,16 +2536,24 @@
 .idle%22),
+%0A       
  new Gau
@@ -2566,32 +2566,36 @@
 eger%3E() %7B%0A      
+  
+  
 @Override%0A      
@@ -2580,32 +2580,36 @@
       @Override%0A
+    
       public Int
@@ -2626,32 +2626,36 @@
 lue() %7B%0A        
+  
+  
 return getPoolSi
@@ -2677,34 +2677,42 @@
 eCount();%0A      
-%7D%0A
+    %7D%0A    
     %7D);%0A    metr
@@ -2762,16 +2762,22 @@
 Class(),
+ name,
  %22thread

Add thread pool name to instrumented thread pool metric names.