Separate global state handler from request tracing

pom.xml

@@ -3,7 +3,7 @@
     <groupId>org.littleshoot</groupId>
     <artifactId>littleproxy</artifactId>
     <packaging>jar</packaging>
-    <version>1.1.4.0-VGS-SNAPSHOT</version>
+    <version>1.1.5.0-VGS-SNAPSHOT</version>
     <name>LittleProxy</name>
     <description>
         LittleProxy is a high performance HTTP proxy written in Java and using the Netty networking framework.

src/main/java/org/littleshoot/proxy/GlobalStateHandler.java

@@ -13,13 +13,6 @@
  */
 public interface GlobalStateHandler {
 
-  /**
-   * Serializes global state to channel.
-   *
-   * @param channel client connection channel
-   */
-  void persistToChannel(Channel channel);
-
   /**
    * Deserializes global state from channel.
    *

src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java

@@ -211,6 +211,21 @@ HttpProxyServerBootstrap withClientToProxyExHandler(
     HttpProxyServerBootstrap withProxyToServerExHandler(
         ExceptionHandler proxyToServerExHandler);
 
+
+    /**
+     * <p>
+     * Specify a {@link RequestTracer} to trace proxy requests
+     * </p>
+     *
+     * <p>
+     * Default = null
+     * </p>
+     *
+     * @param requestTracer
+     * @return proxy server bootstrap
+     */
+    HttpProxyServerBootstrap withRequestTracer(RequestTracer requestTracer);
+
     /**
      * <p>
      * Specify an {@link GlobalStateHandler} to customize a global state based on channel attributes

src/main/java/org/littleshoot/proxy/RequestTracer.java

@@ -0,0 +1,18 @@
+package org.littleshoot.proxy;
+
+import io.netty.channel.Channel;
+
+public interface RequestTracer {
+
+  /**
+   * Start tracing proxy request
+   * @param channel
+   */
+  void start(Channel channel);
+
+  /**
+   * Request is served. Finish tracing.
+   * @param channel
+   */
+  void finish(Channel channel);
+}

src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java

@@ -4,7 +4,6 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
-import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.http.DefaultHttpRequest;
 import io.netty.handler.codec.http.FullHttpRequest;
@@ -771,17 +770,6 @@ protected void exceptionCaught(Throwable cause) {
         }
     }
 
-    @Override
-    public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
-        try {
-            if (this.proxyServer.getGlobalStateHandler() != null) {
-                this.proxyServer.getGlobalStateHandler().persistToChannel(ctx.channel());
-            }
-        } finally {
-            super.channelRegistered(ctx);
-        }
-    }
-
     /***************************************************************************
      * Connection Management
      **************************************************************************/
@@ -803,6 +791,10 @@ public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
     private void initChannelPipeline(ChannelPipeline pipeline) {
         LOG.debug("Configuring ChannelPipeline");
 
+        if (proxyServer.getRequestTracer() != null) {
+            pipeline.addLast("requestTracerHandler", new RequestTracerHandler(this));
+        }
+
         if (proxyServer.getGlobalStateHandler() != null) {
             pipeline.addLast("inboundGlobalStateHandler", new InboundGlobalStateHandler(this));
         }

src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java

@@ -33,6 +33,7 @@
 import org.littleshoot.proxy.MitmManagerFactory;
 import org.littleshoot.proxy.ProxyAuthenticator;
 import org.littleshoot.proxy.ExceptionHandler;
+import org.littleshoot.proxy.RequestTracer;
 import org.littleshoot.proxy.SslEngineSource;
 import org.littleshoot.proxy.TransportProtocol;
 import org.littleshoot.proxy.UnknownTransportProtocolException;
@@ -114,6 +115,7 @@ public class DefaultHttpProxyServer implements HttpProxyServer {
     private final MitmManagerFactory mitmManagerFactory;
     private final ExceptionHandler clientToProxyExHandler;
     private final ExceptionHandler proxyToServerExHandler;
+    private final RequestTracer requestTracer;
     private final GlobalStateHandler globalStateHandler;
     private final HttpFiltersSource filtersSource;
     private final FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer;
@@ -250,6 +252,7 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
             MitmManagerFactory mitmManagerFactory,
             ExceptionHandler clientToProxyExHandler,
             ExceptionHandler proxyToServerExHandler,
+            RequestTracer requestTracer,
             GlobalStateHandler globalStateHandler,
             HttpFiltersSource filtersSource,
             FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer,
@@ -276,6 +279,7 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
         this.mitmManagerFactory = mitmManagerFactory;
         this.clientToProxyExHandler = clientToProxyExHandler;
         this.proxyToServerExHandler = proxyToServerExHandler;
+        this.requestTracer = requestTracer;
         this.globalStateHandler = globalStateHandler;
         this.filtersSource = filtersSource;
         this.unrecoverableFailureHttpResponseComposer = unrecoverableFailureHttpResponseComposer;
@@ -414,6 +418,7 @@ public HttpProxyServerBootstrap clone() {
                     mitmManagerFactory,
                     clientToProxyExHandler,
                     proxyToServerExHandler,
+                    requestTracer,
                     globalStateHandler,
                     filtersSource,
                     unrecoverableFailureHttpResponseComposer,
@@ -605,6 +610,10 @@ protected GlobalStateHandler getGlobalStateHandler() {
         return globalStateHandler;
     }
 
+    protected RequestTracer getRequestTracer() {
+        return requestTracer;
+    }
+
     protected SslEngineSource getSslEngineSource() {
         return sslEngineSource;
     }
@@ -649,6 +658,7 @@ private static class DefaultHttpProxyServerBootstrap implements HttpProxyServerB
         private MitmManagerFactory mitmManagerFactory = null;
         private ExceptionHandler clientToProxyExHandler = null;
         private ExceptionHandler proxyToServerExHandler = null;
+        private RequestTracer requestTracer = null;
         private GlobalStateHandler globalStateHandler = null;
         private HttpFiltersSource filtersSource = new HttpFiltersSourceAdapter();
         private FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer = new DefaultFailureHttpResponseComposer();
@@ -683,6 +693,7 @@ private DefaultHttpProxyServerBootstrap(
                 MitmManagerFactory mitmManagerFactory,
                 ExceptionHandler clientToProxyExHandler,
                 ExceptionHandler proxyToServerExHandler,
+                RequestTracer requestTracer,
                 GlobalStateHandler globalStateHandler,
                 HttpFiltersSource filtersSource,
                 FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer,
@@ -708,6 +719,7 @@ private DefaultHttpProxyServerBootstrap(
             this.mitmManagerFactory = mitmManagerFactory;
             this.clientToProxyExHandler = clientToProxyExHandler;
             this.proxyToServerExHandler = proxyToServerExHandler;
+            this.requestTracer = requestTracer;
             this.globalStateHandler = globalStateHandler;
             this.filtersSource = filtersSource;
             this.unrecoverableFailureHttpResponseComposer = unrecoverableFailureHttpResponseComposer;
@@ -856,14 +868,20 @@ public HttpProxyServerBootstrap withClientToProxyExHandler(
             return this;
         }
 
+        @Override
+        public HttpProxyServerBootstrap withRequestTracer(
+            RequestTracer requestTracer) {
+            this.requestTracer = requestTracer;
+            return this;
+        }
+
         @Override
         public HttpProxyServerBootstrap withCustomGlobalState(
             GlobalStateHandler globalStateHandler) {
             this.globalStateHandler = globalStateHandler;
             return this;
         }
 
-
         @Override
         public HttpProxyServerBootstrap withFiltersSource(
                 HttpFiltersSource filtersSource) {
@@ -980,7 +998,7 @@ private DefaultHttpProxyServer build() {
                     transportProtocol, determineListenAddress(),
                     sslEngineSource, authenticateSslClients,
                     proxyAuthenticator, chainProxyManager, mitmManagerFactory,
-                    clientToProxyExHandler, proxyToServerExHandler, globalStateHandler,
+                    clientToProxyExHandler, proxyToServerExHandler, requestTracer, globalStateHandler,
                     filtersSource, unrecoverableFailureHttpResponseComposer, transparent,
                     idleConnectionTimeout, activityTrackers, connectTimeout,
                     serverResolver, readThrottleBytesPerSecond, writeThrottleBytesPerSecond,

src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java

@@ -743,6 +743,41 @@ public void channelRead(ChannelHandlerContext ctx, Object msg)
         protected abstract void responseRead(HttpResponse httpResponse);
     }
 
+    @Sharable
+    protected class RequestTracerHandler extends ChannelDuplexHandler {
+
+        private final ProxyConnection<HttpRequest> clientToProxyConnection;
+
+        RequestTracerHandler(ProxyConnection<HttpRequest> clientToProxyConnection) {
+            this.clientToProxyConnection = clientToProxyConnection;
+        }
+
+        @Override
+        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+            try {
+                proxyServer.getRequestTracer().start(clientToProxyConnection.channel);
+            } catch (Throwable t) {
+                LOG.warn("Unable to start tracing request", t);
+            } finally {
+                super.channelRead(ctx, msg);
+            }
+        }
+
+        @Override
+        public void write(ChannelHandlerContext ctx,
+                          Object msg, ChannelPromise promise) throws Exception {
+            try {
+                super.write(ctx, msg, promise);
+            } finally {
+                try {
+                    proxyServer.getRequestTracer().finish(clientToProxyConnection.channel);
+                } catch (Throwable t) {
+                    LOG.warn("Unable to finish request tracing", t);
+                }
+            }
+        }
+    }
+
     @Sharable
     protected class InboundGlobalStateHandler extends
         ChannelInboundHandlerAdapter {