package com.bytedance.hadoop.hdfs.server.quota; import com.bytedance.hadoop.hdfs.server.NNProxy; import com.google.common.base.Function; import com.google.common.base.Preconditions; import org.apache.hadoop.classification.InterfaceAudience; import org.apache.hadoop.classification.InterfaceStability; import org.apache.hadoop.ipc.StandbyException; import java.lang.reflect.InvocationHandler; import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.util.concurrent.atomic.AtomicLong; @InterfaceAudience.Private @InterfaceStability.Stable public class ThrottleInvocationHandler implements InvocationHandler { final Object underlying; final Function<Method, AtomicLong> opCounter; final long threshold; public ThrottleInvocationHandler(Object underlying, Function<Method, AtomicLong> opCounter, long threshold) { this.underlying = underlying; this.opCounter = opCounter; this.threshold = threshold; } @Override public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { AtomicLong counter = opCounter.apply(method); Preconditions.checkState(counter != null); long current = counter.getAndIncrement(); try { if (current > threshold) { NNProxy.proxyMetrics.throttledOps.incr(); throw new StandbyException("Too many requests (" + current + "/" + threshold + "), try later"); } Object ret = method.invoke(underlying, args); NNProxy.proxyMetrics.successOps.incr(); return ret; } catch (InvocationTargetException e) { NNProxy.proxyMetrics.failedOps.incr(); throw e.getCause(); } finally { counter.decrementAndGet(); } } }