From c7395c33ee0a69035c9ca3b2e0129f96fabf4643 Mon Sep 17 00:00:00 2001
From: mreedell <michael_reedell@comcast.com>
Date: Mon, 25 Aug 2014 09:36:12 -0400
Subject: [PATCH] Changing throttle configuration to accept read/write limits
 instead of GlobalTrafficShapingHandler.

- Proxy server accepts read/write limits in bytes/second.
- Added getters for read/write limits.
---
 .../littleshoot/proxy/HttpProxyServer.java    |  9 ++++
 .../proxy/HttpProxyServerBootstrap.java       |  9 ++--
 .../proxy/impl/DefaultHttpProxyServer.java    | 46 ++++++++++++++-----
 3 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/main/java/org/littleshoot/proxy/HttpProxyServer.java b/src/main/java/org/littleshoot/proxy/HttpProxyServer.java
index 28976cc8d..4c00532bc 100644
--- a/src/main/java/org/littleshoot/proxy/HttpProxyServer.java
+++ b/src/main/java/org/littleshoot/proxy/HttpProxyServer.java
@@ -39,4 +39,13 @@ public interface HttpProxyServer {
      * @return
      */
     InetSocketAddress getListenAddress();
+
+    /**
+     * <p>
+     * Set the read/write throttle bandwidths (in bytes/second) for this proxy.
+     * </p>
+     * @param readThrottleBytesPerSecond
+     * @param writeThrottleBytesPerSecond
+     */
+    void setThrottle(long readThrottleBytesPerSecond, long writeThrottleBytesPerSecond);
 }
diff --git a/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java b/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java
index 77d0321e6..7e512fe37 100644
--- a/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java
+++ b/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java
@@ -1,7 +1,5 @@
 package org.littleshoot.proxy;
 
-import io.netty.handler.traffic.GlobalTrafficShapingHandler;
-
 import java.net.InetSocketAddress;
 
 /**
@@ -287,12 +285,13 @@ HttpProxyServerBootstrap withConnectTimeout(
 
     /**
      * <p>
-     * Specify a global traffic shaping handler for this proxy server.
+     * Specify the read and/or write bandwidth throttles for this proxy server. 0 indicates not throttling.
      * </p>
-     * @param globalTrafficShapingHandler
+     * @param readThrottleBytesPerSecond
+     * @param writeThrottleBytesPerSecond
      * @return
      */
-    HttpProxyServerBootstrap withGlobalTrafficShapingHandler(GlobalTrafficShapingHandler globalTrafficShapingHandler);
+    HttpProxyServerBootstrap withThrottling(long readThrottleBytesPerSecond, long writeThrottleBytesPerSecond);
 
     /**
      * <p>
diff --git a/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java b/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java
index ad79c6463..d9b219efc 100644
--- a/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java
+++ b/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java
@@ -35,6 +35,8 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 
@@ -183,6 +185,10 @@ public static HttpProxyServerBootstrap bootstrapFromFile(String path) {
      *            server
      * @param serverResolver
      *            the {@link HostResolver} to use for resolving server addresses
+     * @param readThrottleBytesPerSecond
+     *            read throttle bandwidth
+     * @param writeThrottleBytesPerSecond
+     *            write throttle bandwidth
      */
     private DefaultHttpProxyServer(ServerGroup serverGroup,
             TransportProtocol transportProtocol,
@@ -198,7 +204,8 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
             Collection<ActivityTracker> activityTrackers,
             int connectTimeout,
             HostResolver serverResolver,
-            GlobalTrafficShapingHandler globalTrafficShapingHandler) {
+            long readThrottleBytesPerSecond,
+            long writeThrottleBytesPerSecond) {
         this.serverGroup = serverGroup;
         this.transportProtocol = transportProtocol;
         this.address = address;
@@ -215,7 +222,9 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
         }
         this.connectTimeout = connectTimeout;
         this.serverResolver = serverResolver;
-        this.globalTrafficShapingHandler = globalTrafficShapingHandler;
+
+        ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(50);
+        this.globalTrafficShapingHandler = new GlobalTrafficShapingHandler(executor, writeThrottleBytesPerSecond, readThrottleBytesPerSecond);
     }
 
     boolean isTransparent() {
@@ -242,9 +251,18 @@ public HostResolver getServerResolver() {
     public InetSocketAddress getListenAddress() {
         return address;
     }
-    
-    public GlobalTrafficShapingHandler getGlobalTrafficShapingHandler() {
-        return globalTrafficShapingHandler;
+
+    @Override
+    public void setThrottle(long readThrottleBytesPerSecond, long writeThrottleBytesPerSecond) {
+        globalTrafficShapingHandler.configure(writeThrottleBytesPerSecond, readThrottleBytesPerSecond);
+    }
+
+    public long getReadThrottle() {
+        return globalTrafficShapingHandler.getReadLimit();
+    }
+
+    public long getWriteThrottle() {
+        return globalTrafficShapingHandler.getWriteLimit();
     }
 
     @Override
@@ -256,7 +274,8 @@ public HttpProxyServerBootstrap clone() {
                 chainProxyManager,
                 mitmManager, filtersSource, transparent,
                 idleConnectionTimeout, activityTrackers, connectTimeout,
-                serverResolver, globalTrafficShapingHandler);
+                serverResolver, globalTrafficShapingHandler.getReadLimit(),
+                globalTrafficShapingHandler.getWriteLimit());
     }
 
     @Override
@@ -560,7 +579,8 @@ private static class DefaultHttpProxyServerBootstrap implements
         private Collection<ActivityTracker> activityTrackers = new ConcurrentLinkedQueue<ActivityTracker>();
         private int connectTimeout = 40000;
         private HostResolver serverResolver = new DefaultHostResolver();
-        private GlobalTrafficShapingHandler globalTrafficShapingHandler;
+        private long readThrottleBytesPerSecond;
+        private long writeThrottleBytesPerSecond;
 
         private DefaultHttpProxyServerBootstrap() {
         }
@@ -578,7 +598,7 @@ private DefaultHttpProxyServerBootstrap(
                 boolean transparent, int idleConnectionTimeout,
                 Collection<ActivityTracker> activityTrackers,
                 int connectTimeout, HostResolver serverResolver,
-                GlobalTrafficShapingHandler globalTrafficShapingHandler) {
+                long readThrottleBytesPerSecond, long  writeThrottleBytesPerSecond) {
             this.original = original;
             this.transportProtocol = transportProtocol;
             this.address = address;
@@ -595,7 +615,8 @@ private DefaultHttpProxyServerBootstrap(
             }
             this.connectTimeout = connectTimeout;
             this.serverResolver = serverResolver;
-            this.globalTrafficShapingHandler = globalTrafficShapingHandler;
+            this.readThrottleBytesPerSecond = readThrottleBytesPerSecond;
+            this.writeThrottleBytesPerSecond = writeThrottleBytesPerSecond;
         }
 
         private DefaultHttpProxyServerBootstrap(Properties props) {
@@ -745,8 +766,9 @@ public HttpProxyServerBootstrap plusActivityTracker(
         }
 
         @Override
-        public HttpProxyServerBootstrap withGlobalTrafficShapingHandler(GlobalTrafficShapingHandler globalTrafficShapingHandler) {
-            this.globalTrafficShapingHandler = globalTrafficShapingHandler;
+        public HttpProxyServerBootstrap withThrottling(long readThrottleBytesPerSecond, long writeThrottleBytesPerSecond) {
+            this.readThrottleBytesPerSecond = readThrottleBytesPerSecond;
+            this.writeThrottleBytesPerSecond = writeThrottleBytesPerSecond;
             return this;
         }
 
@@ -771,7 +793,7 @@ transportProtocol, determineListenAddress(),
                     proxyAuthenticator, chainProxyManager, mitmManager,
                     filtersSource, transparent,
                     idleConnectionTimeout, activityTrackers, connectTimeout,
-                    serverResolver, globalTrafficShapingHandler);
+                    serverResolver, readThrottleBytesPerSecond, writeThrottleBytesPerSecond);
         }
 
         private InetSocketAddress determineListenAddress() {