Merge pull request adamfisk#170 from jekh/request-sent-callback-after…

…-content-sent Call proxyToServerRequestSent() after LastHttpContent is sent
vyarosh · Mar 10, 2015 · 2bf48bc · 2bf48bc
2 parents bf3fd76 + c6d0885
commit 2bf48bc
Show file tree

Hide file tree

Showing 4 changed files with 238 additions and 112 deletions.
diff --git a/src/main/java/org/littleshoot/proxy/HttpFilters.java b/src/main/java/org/littleshoot/proxy/HttpFilters.java
@@ -97,7 +97,7 @@ public interface HttpFilters {
     void proxyToServerRequestSending();
 
     /**
-     * Informs filter that proxy to server request has been sent.
+     * Informs filter that the HTTP request, including any content, has been sent.
      */
     void proxyToServerRequestSent();
 

diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java
@@ -760,34 +760,40 @@ protected abstract class RequestWrittenMonitor extends
         public void write(ChannelHandlerContext ctx,
                 Object msg, ChannelPromise promise)
                 throws Exception {
-
             HttpRequest originalRequest = null;
             if (msg instanceof HttpRequest) {
                 originalRequest = (HttpRequest) msg;
             }
 
-            try {
-                if (null != originalRequest) {
-                    requestWritten(originalRequest);
-                }
-            } catch (Throwable t) {
-                LOG.warn("Unable to record bytesRead", t);
-            } finally {
-                if (null != originalRequest) {
-                    getHttpFiltersFromProxyServer(originalRequest)
-                            .proxyToServerRequestSending();
-                }
+            if (null != originalRequest) {
+                requestWriting(originalRequest);
+            }
 
-                super.write(ctx, msg, promise);
+            super.write(ctx, msg, promise);
 
-                if (null != originalRequest) {
-                    getHttpFiltersFromProxyServer(originalRequest)
-                            .proxyToServerRequestSent();
-                }
+            if (null != originalRequest) {
+                requestWritten(originalRequest);
+            }
+
+            if (msg instanceof HttpContent) {
+                contentWritten((HttpContent) msg);
             }
         }
 
+        /**
+         * Invoked immediately before an HttpRequest is written.
+         */
+        protected abstract void requestWriting(HttpRequest httpRequest);
+
+        /**
+         * Invoked immediately after an HttpRequest has been sent.
+         */
         protected abstract void requestWritten(HttpRequest httpRequest);
+
+        /**
+         * Invoked immediately after an HttpContent has been sent.
+         */
+        protected abstract void contentWritten(HttpContent httpContent);
     }
 
     /**
@@ -805,7 +811,7 @@ public void write(ChannelHandlerContext ctx,
                     responseWritten(((HttpResponse) msg));
                 }
             } catch (Throwable t) {
-                LOG.warn("Unable to record bytesRead", t);
+                LOG.warn("Error while invoking responseWritten callback", t);
             } finally {
                 super.write(ctx, msg, promise);
             }

diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java
@@ -18,6 +18,7 @@
 import io.netty.handler.codec.http.HttpRequestEncoder;
 import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.codec.http.HttpResponseDecoder;
+import io.netty.handler.codec.http.LastHttpContent;
 import io.netty.handler.timeout.IdleStateHandler;
 import io.netty.handler.traffic.GlobalTrafficShapingHandler;
 import io.netty.util.ReferenceCounted;
@@ -902,12 +903,29 @@ protected void bytesWritten(int numberOfBytes) {
 
     private RequestWrittenMonitor requestWrittenMonitor = new RequestWrittenMonitor() {
         @Override
-        protected void requestWritten(HttpRequest httpRequest) {
+        protected void requestWriting(HttpRequest httpRequest) {
             FullFlowContext flowContext = new FullFlowContext(clientConnection,
                     ProxyToServerConnection.this);
-            for (ActivityTracker tracker : proxyServer
-                    .getActivityTrackers()) {
-                tracker.requestSentToServer(flowContext, httpRequest);
+            try {
+                for (ActivityTracker tracker : proxyServer
+                        .getActivityTrackers()) {
+                    tracker.requestSentToServer(flowContext, httpRequest);
+                }
+            } catch (Throwable t) {
+                LOG.warn("Error while invoking ActivityTracker on request", t);
+            }
+
+            currentFilters.proxyToServerRequestSending();
+        }
+
+        @Override
+        protected void requestWritten(HttpRequest httpRequest) {
+        }
+
+        @Override
+        protected void contentWritten(HttpContent httpContent) {
+            if (httpContent instanceof LastHttpContent) {
+                currentFilters.proxyToServerRequestSent();
             }
         }
     };