cleanup

okhttp-clients/src/main/java/com/palantir/remoting3/okhttp/ConcurrencyLimitingInterceptor.java

@@ -19,6 +19,10 @@
 import com.google.common.collect.ImmutableSet;
 import com.netflix.concurrency.limits.Limiter;
 import java.io.IOException;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
 import okhttp3.Interceptor;
 import okhttp3.Response;
 import okhttp3.ResponseBody;
@@ -78,28 +82,47 @@ private static Response wrapResponse(Limiter.Listener listener, Response respons
         }
         ResponseBody currentBody = response.body();
         ResponseBody newResponseBody =
-                ResponseBody.create(currentBody.contentType(), currentBody.contentLength(),
-                        new ReleaseConcurrencyLimitBufferedSource(currentBody.source(), listener));
+                ResponseBody.create(
+                        currentBody.contentType(),
+                        currentBody.contentLength(),
+                        wrapSource(currentBody.source(), listener));
         return response.newBuilder()
                 .body(newResponseBody)
                 .build();
     }
 
-    private static final class ReleaseConcurrencyLimitBufferedSource extends ForwardingBufferedSource {
+    private static BufferedSource wrapSource(BufferedSource currentSource, Limiter.Listener listener) {
+        return (BufferedSource) Proxy.newProxyInstance(
+                BufferedSource.class.getClassLoader(),
+                new Class<?>[] { BufferedSource.class },
+                new ReleaseConcurrencyLimitProxy(currentSource, listener));
+    }
+
+    /**
+     * This proxy enables e.g. Okio to make additive additions to their API without breaking us.
+     */
+    private static final class ReleaseConcurrencyLimitProxy implements InvocationHandler {
         private final BufferedSource delegate;
         private final Limiter.Listener listener;
+        private boolean closed = false;
 
-        private ReleaseConcurrencyLimitBufferedSource(BufferedSource delegate, Limiter.Listener listener) {
-            super(delegate);
-            this.listener = listener;
+        private ReleaseConcurrencyLimitProxy(BufferedSource delegate, Limiter.Listener listener) {
             this.delegate = delegate;
+            this.listener = listener;
         }
 
         @Override
-        public void close() throws IOException {
-            listener.onSuccess();
-            delegate.close();
+        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+            if (method.getName().equals("close") && !closed) {
+                closed = true;
+                listener.onSuccess();
+            }
+
+            try {
+                return method.invoke(delegate, args);
+            } catch (InvocationTargetException e) {
+                throw e.getCause();
+            }
         }
     }
-
 }

okhttp-clients/src/main/java/com/palantir/remoting3/okhttp/OkHttpClients.java

@@ -190,6 +190,7 @@ private static RemotingOkHttpClient createInternal(
         client.sslSocketFactory(config.sslSocketFactory(), config.trustManager());
 
         // Intercept calls to augment request meta data
+        client.addInterceptor(new ConcurrencyLimitingInterceptor());
         client.addInterceptor(InstrumentedInterceptor.create(registry, hostMetrics, serviceClass));
         client.addInterceptor(OkhttpTraceInterceptor.INSTANCE);
         client.addInterceptor(UserAgentInterceptor.of(augmentUserAgent(userAgent, serviceClass)));