Make ThrottleRequestFilter preserve interfaces of AsyncHandlers, close AsyncHttpClient#1314

Author: Markus Heiden

client/src/main/java/org/asynchttpclient/filter/AsyncHandlerWrapper.java

@@ -0,0 +1,50 @@
+package org.asynchttpclient.filter;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.Semaphore;
+
+import org.asynchttpclient.AsyncHandler;
+
+/**
+ * Wrapper for {@link AsyncHandler}s to release a permit on {@link AsyncHandler#onCompleted()}.
+ * This is done via a dynamic proxy to preserve all interfaces of the wrapped handler.
+ */
+public class ReleasePermitOnComplete {
+   /**
+    * Wrap handler to release the permit of the semaphore on {@link AsyncHandler#onCompleted()}.
+    */
+   public static <T> AsyncHandler<T> wrap(final AsyncHandler<T> handler, final Semaphore available) {
+      Class<?> handlerClass = handler.getClass();
+      ClassLoader classLoader = handlerClass.getClassLoader();
+      Class<?>[] interfaces = allInterfaces(handlerClass);
+
+      return (AsyncHandler<T>) Proxy.newProxyInstance(classLoader, interfaces, new InvocationHandler() {
+         @Override
+         public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+            try {
+               return method.invoke(handler, args);
+            } finally {
+               if ("onCompleted".equals(method.getName())) {
+                  available.release();
+               }
+            }
+         }
+      });
+   }
+
+   /**
+    * Extract all interfaces of a class.
+    */
+   static Class<?>[] allInterfaces(Class<?> handlerClass) {
+      Set<Class<?>> allInterfaces = new HashSet<>();
+      for (Class<?> clazz = handlerClass; clazz != null; clazz = clazz.getSuperclass()) {
+         Collections.addAll(allInterfaces, clazz.getInterfaces());
+      }
+      return allInterfaces.toArray(new Class[allInterfaces.size()]);
+   }
+}

client/src/main/java/org/asynchttpclient/filter/ReleasePermitOnComplete.java

@@ -19,11 +19,11 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * A {@link org.asynchttpclient.filter.RequestFilter} throttles requests and block when the number of permits is reached, waiting for
- * the response to arrives before executing the next request.
+ * A {@link org.asynchttpclient.filter.RequestFilter} throttles requests and block when the number of permits is reached,
+ * waiting for the response to arrives before executing the next request.
  */
 public class ThrottleRequestFilter implements RequestFilter {
-    private final static Logger logger = LoggerFactory.getLogger(ThrottleRequestFilter.class);
+    private static final Logger logger = LoggerFactory.getLogger(ThrottleRequestFilter.class);
     private final Semaphore available;
     private final int maxWait;
 
@@ -32,33 +32,34 @@ public ThrottleRequestFilter(int maxConnections) {
     }
 
     public ThrottleRequestFilter(int maxConnections, int maxWait) {
-      this(maxConnections, maxWait, false);
+        this(maxConnections, maxWait, false);
     }
 
     public ThrottleRequestFilter(int maxConnections, int maxWait, boolean fair) {
-      this.maxWait = maxWait;
-      available = new Semaphore(maxConnections, fair);
+        this.maxWait = maxWait;
+        available = new Semaphore(maxConnections, fair);
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public <T> FilterContext<T> filter(FilterContext<T> ctx) throws FilterException {
-
         try {
             if (logger.isDebugEnabled()) {
                 logger.debug("Current Throttling Status {}", available.availablePermits());
             }
             if (!available.tryAcquire(maxWait, TimeUnit.MILLISECONDS)) {
                 throw new FilterException(String.format("No slot available for processing Request %s with AsyncHandler %s",
-                        ctx.getRequest(), ctx.getAsyncHandler()));
+                      ctx.getRequest(), ctx.getAsyncHandler()));
             }
         } catch (InterruptedException e) {
-            throw new FilterException(String.format("Interrupted Request %s with AsyncHandler %s", ctx.getRequest(), ctx.getAsyncHandler()));
+            throw new FilterException(String.format("Interrupted Request %s with AsyncHandler %s",
+                  ctx.getRequest(), ctx.getAsyncHandler()));
         }
 
-        return new FilterContext.FilterContextBuilder<>(ctx).asyncHandler(new AsyncHandlerWrapper<>(ctx.getAsyncHandler(), available))
-                .build();
+        return new FilterContext.FilterContextBuilder<>(ctx)
+              .asyncHandler(ReleasePermitOnComplete.wrap(ctx.getAsyncHandler(), available))
+              .build();
     }
-}
+}

client/src/main/java/org/asynchttpclient/filter/ThrottleRequestFilter.java

@@ -1,8 +1,8 @@
 package org.asynchttpclient.extras.guava;
 
-import org.asynchttpclient.filter.AsyncHandlerWrapper;
 import org.asynchttpclient.filter.FilterContext;
 import org.asynchttpclient.filter.FilterException;
+import org.asynchttpclient.filter.ReleasePermitOnComplete;
 import org.asynchttpclient.filter.RequestFilter;
 import org.asynchttpclient.filter.ThrottleRequestFilter;
 import org.slf4j.Logger;
@@ -56,8 +56,9 @@ public <T> FilterContext<T> filter(FilterContext<T> ctx) throws FilterException
             throw new FilterException(String.format("Interrupted Request %s with AsyncHandler %s", ctx.getRequest(), ctx.getAsyncHandler()));
         }
 
-        return new FilterContext.FilterContextBuilder<>(ctx).asyncHandler(new AsyncHandlerWrapper<>(ctx.getAsyncHandler(), available))
-                .build();
+        return new FilterContext.FilterContextBuilder<>(ctx)
+              .asyncHandler(ReleasePermitOnComplete.wrap(ctx.getAsyncHandler(), available))
+              .build();
     }
 
     private <T> void attemptRateLimitedPermitAcquistion(FilterContext<T> ctx, long startOfWait) throws FilterException {