ThrottleAspect.groovy

ThrottleAspect.groovy

package throttle

import java.util.concurrent.Executors
import java.util.concurrent.ScheduledExecutorService
import java.util.concurrent.ScheduledFuture
import java.util.concurrent.Semaphore
import java.util.concurrent.TimeUnit

import javax.annotation.PreDestroy

import org.aspectj.lang.ProceedingJoinPoint
import org.aspectj.lang.annotation.Around
import org.aspectj.lang.annotation.Aspect


@Aspect
class ThrottleAspect {
	
	private static class ThrottleDecorator {
		int count = 0
		Semaphore semaphore
		ScheduledFuture timer
		
		public ThrottleDecorator(int count, long periodMillis, ScheduledExecutorService scheduler) {
			this.count = count
			this.semaphore = new Semaphore(count)
			this.timer = scheduler.scheduleAtFixedRate({
				//println "Tic (${periodMillis}) - ${semaphore.availablePermits()} permits available before release"
				synchronized(semaphore) {
					semaphore.release(count - semaphore.availablePermits())
				}
			}, 0, periodMillis, TimeUnit.MILLISECONDS)
		}
		
		public Object proceed(ProceedingJoinPoint pjp) {
			if(timer.isCancelled()) {
				throw new RuntimeException("Invalid timer!")
			}
			semaphore.acquire()
			return pjp.proceed()
		}
	}

	Map throttles = [:]
	ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1)
	
	@Around("@annotation(throttleConfig)")
	public Object doThrottleAction(ProceedingJoinPoint pjp, Throttle throttleConfig) throws Throwable {
		// throttle per instance, per method
		String throttleKey = "${pjp.target.hashCode()}:${pjp.signature.toLongString()}"		
		
		synchronized(throttles) {
			if(!throttles.containsKey(throttleKey)) {
				throttles[throttleKey] = new ThrottleDecorator(throttleConfig.count(), throttleConfig.periodMillis(), scheduler)
			}
		}
		ThrottleDecorator t = throttles[throttleKey]
		return t.proceed(pjp)
	}
	
	@PreDestroy
	public void shutdown() {
		scheduler.shutdownNow()
	}
}