Delete unused proxies after expiration time passes

README.md

@@ -118,12 +118,14 @@ system properties will be used to specify the upstream proxy.
 Command-line Arguments
 ----------------------
 
- - -port <port>
+ - -port \<port\>
   - Port on which the API listens. Default value is 8080.
  - -address <address>
   - Address to which the API is bound. Default value is 0.0.0.0.
- - -proxyPortRange <from>-<to>
-  - Range of ports reserved for proxies. Only applies if *port* parameter is not supplied in the POST request. Default values are <port>+1 to <port>+500+1.
+ - -proxyPortRange \<from\>-\<to\>
+  - Range of ports reserved for proxies. Only applies if *port* parameter is not supplied in the POST request. Default values are \<port\>+1 to \<port\>+500+1.
+ - -ttl \<seconds\>
+  - Proxy will be automatically deleted after a specified time period. Off by default.
 
 Embedded Mode
 -------------

pom.xml

@@ -39,6 +39,20 @@
         <developerConnection>scm:git:git@github.com:lightbody/browsermob-proxy.git</developerConnection>
         <url>git@github.com:lightbody/browsermob-proxy.git</url>
     </scm>
+
+    <repositories>
+        <repository>
+            <id>guiceyfruit.release</id>
+            <name>GuiceyFruit Release Repository</name>
+            <url>http://guiceyfruit.googlecode.com/svn/repo/releases/</url>
+            <snapshots>
+                <enabled>false</enabled>
+            </snapshots>
+            <releases>
+                <enabled>true</enabled>
+            </releases>
+        </repository>
+    </repositories>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
@@ -249,6 +263,12 @@
             <artifactId>guice-servlet</artifactId>
             <version>3.0</version>
         </dependency>
+
+        <dependency>
+            <groupId>org.guiceyfruit</groupId>
+            <artifactId>guiceyfruit-core</artifactId>
+            <version>2.0</version>                       
+        </dependency>
 
         <dependency>
             <groupId>net.jcip</groupId>

src/main/java/net/lightbody/bmp/proxy/Main.java

@@ -4,20 +4,6 @@
 import com.google.inject.Injector;
 import com.google.inject.servlet.GuiceServletContextListener;
 import com.google.sitebricks.SitebricksModule;
-
-import net.lightbody.bmp.exception.JettyException;
-import net.lightbody.bmp.proxy.bricks.ProxyResource;
-import net.lightbody.bmp.proxy.guice.ConfigModule;
-import net.lightbody.bmp.proxy.guice.JettyModule;
-import net.lightbody.bmp.proxy.util.StandardFormatter;
-
-import org.eclipse.jetty.server.Server;
-import org.eclipse.jetty.servlet.ServletContextHandler;
-import org.eclipse.jetty.util.log.Log;
-import org.slf4j.LoggerFactory;
-
-import javax.servlet.ServletContextEvent;
-
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -28,6 +14,17 @@
 import java.util.logging.Level;
 import java.util.logging.LogManager;
 import java.util.logging.Logger;
+import javax.servlet.ServletContextEvent;
+import net.lightbody.bmp.exception.JettyException;
+import net.lightbody.bmp.proxy.bricks.ProxyResource;
+import net.lightbody.bmp.proxy.guice.ConfigModule;
+import net.lightbody.bmp.proxy.guice.JettyModule;
+import net.lightbody.bmp.proxy.util.StandardFormatter;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.util.log.Log;
+import org.guiceyfruit.jsr250.Jsr250Module;
+import org.slf4j.LoggerFactory;
 
 public class Main {
     private static final String LOGGING_PROPERTIES_FILENAME = "conf/bmp-logging.properties";
@@ -38,7 +35,7 @@ public class Main {
     public static void main(String[] args) {
         configureJdkLogging();
 
-        final Injector injector = Guice.createInjector(new ConfigModule(args), new JettyModule(), new SitebricksModule() {
+        final Injector injector = Guice.createInjector(new ConfigModule(args), new Jsr250Module(), new JettyModule(), new SitebricksModule() {
             @Override
             protected void configureSitebricks() {
                 scan(ProxyResource.class.getPackage());
@@ -148,4 +145,4 @@ private static void configureDefaultLogger() {
         logger.addHandler(handler);
 	}
 }
-
+

src/main/java/net/lightbody/bmp/proxy/ProxyManager.java

@@ -3,15 +3,14 @@
 import com.google.inject.Inject;
 import com.google.inject.Provider;
 import com.google.inject.Singleton;
-
 import com.google.inject.name.Named;
-
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-
+import javax.annotation.PreDestroy;
+import net.lightbody.bmp.proxy.util.ExpirableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -22,15 +21,28 @@ public class ProxyManager {
     private int lastPort;
     private final int minPort; 
     private final int maxPort;
-    private Provider<ProxyServer> proxyServerProvider;
-    private ConcurrentHashMap<Integer, ProxyServer> proxies = new ConcurrentHashMap<Integer, ProxyServer>();
+    private final Provider<ProxyServer> proxyServerProvider;
+    private final ConcurrentHashMap<Integer, ProxyServer> proxies;
 
     @Inject
-    public ProxyManager(Provider<ProxyServer> proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort) {
+    public ProxyManager(Provider<ProxyServer> proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort, @Named("ttl") Integer ttl) {
         this.proxyServerProvider = proxyServerProvider;
         this.minPort = minPort;
         this.maxPort = maxPort;
         this.lastPort = maxPort; 
+        this.proxies = ttl > 0 ?                
+                new ExpirableMap<Integer, ProxyServer>(ttl, new ExpirableMap.OnExpire<ProxyServer>(){
+                    @Override
+                    public void run(ProxyServer proxy) {
+                        try {
+                            LOG.debug("Expiring ProxyServer `{}`...", proxy.getPort());
+                            proxy.stop();
+                        } catch (Exception ex) {
+                            LOG.warn("Error while stopping an expired proxy", ex);
+                        }
+                    }
+                }) : 
+                new ConcurrentHashMap<Integer, ProxyServer>();    
     }
 
     public ProxyServer create(Map<String, String> options, Integer port, String bindAddr) {
@@ -115,4 +127,11 @@ public void delete(int port) {
         ProxyServer proxy = proxies.remove(port);
         proxy.stop();
     }
+
+    @PreDestroy
+    public void stop(){
+        if(proxies instanceof ExpirableMap){                   
+            ((ExpirableMap)proxies).stop();
+        }
+    }
 }

src/main/java/net/lightbody/bmp/proxy/guice/ConfigModule.java

@@ -39,6 +39,12 @@ public void configure(Binder binder) {
                       .defaultsTo(8081, 8581)
                       .withValuesSeparatedBy('-');
 
+        ArgumentAcceptingOptionSpec<Integer> ttlSpec =
+                parser.accepts("ttl", "Time in seconds until an unused proxy is deleted")
+                      .withOptionalArg()
+                      .ofType(Integer.class)
+                      .defaultsTo(0);
+
         parser.acceptsAll(asList("help", "?"), "This help text");
 
         OptionSet options = parser.parse(args);
@@ -77,8 +83,9 @@ public void configure(Binder binder) {
         binder.bind(Key.get(Integer.class, new NamedImpl("port"))).toInstance(port);
         binder.bind(Key.get(String.class, new NamedImpl("address"))).toInstance(addressSpec.value(options));
         binder.bind(Key.get(Integer.class, new NamedImpl("minPort"))).toInstance(minPort);
-        binder.bind(Key.get(Integer.class, new NamedImpl("maxPort"))).toInstance(maxPort);   
-
+        binder.bind(Key.get(Integer.class, new NamedImpl("maxPort"))).toInstance(maxPort);                 
+        binder.bind(Key.get(Integer.class, new NamedImpl("ttl"))).toInstance(ttlSpec.value(options));
+
         /*
          * Init User Agent String Parser, update of the UAS datastore will run in background.
          */

src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java

@@ -0,0 +1,97 @@
+package net.lightbody.bmp.proxy.util;
+
+import java.util.Date;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+public class ExpirableMap<K,V> extends ConcurrentHashMap<K,V>{
+    public final static int DEFAULT_CHECK_INTERVAL = 10*60;
+    public final static int DEFAULT_TTL = 30*60;
+    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
+    private final long ttl;    
+    private final Map<K, Long> expires;
+    private final OnExpire<V> onExpire;
+
+    public ExpirableMap(int ttl, int checkInterval, OnExpire<V> onExpire) {
+        this.ttl = ttl*1000;
+        this.onExpire = onExpire;
+        expires = new HashMap<>();        
+        scheduler.scheduleWithFixedDelay(new Worker(), checkInterval, checkInterval, TimeUnit.SECONDS);
+    }
+
+    public ExpirableMap(int ttl, OnExpire<V> onExpire) {
+        this(ttl, DEFAULT_CHECK_INTERVAL, onExpire); 
+    }
+
+    public ExpirableMap(OnExpire<V> onExpire) {
+        this(DEFAULT_TTL, DEFAULT_CHECK_INTERVAL, onExpire);        
+    }   
+
+    public ExpirableMap() {
+        this(DEFAULT_TTL, DEFAULT_CHECK_INTERVAL, null);        
+    }  
+
+    @Override
+    public V putIfAbsent(K key, V value) {
+        synchronized(this){            
+            expires.put(key, new Date().getTime()+ttl);            
+            return super.putIfAbsent(key, value);            
+        }   
+    }        
+
+    @Override
+    public V put(K key, V value) {                        
+        synchronized(this){
+            expires.put(key, new Date().getTime()+ttl);
+            return super.put(key, value);
+        }        
+    }
+
+    public void stop() {
+        scheduler.shutdown();       
+        try {
+            scheduler.awaitTermination(10, TimeUnit.SECONDS);
+        } catch (InterruptedException ex) {
+            scheduler.shutdownNow();
+        }        
+    }
+
+    private class Worker implements Runnable{
+
+        @Override
+        public void run() {
+            Map<K, Long> m;
+            synchronized(ExpirableMap.this){
+                m = new HashMap<>(expires);
+            }
+            Long now = new Date().getTime();
+            for(Entry<K, Long> e : m.entrySet()){
+                if(e.getValue() > now){
+                    continue;
+                }
+                synchronized(ExpirableMap.this){
+                    Long expire = expires.get(e.getKey());
+                    if(expire == null){                            
+                        continue;
+                    }
+                    if(expire <= new Date().getTime()){
+                        expires.remove(e.getKey());                        
+                        V v = ExpirableMap.this.remove(e.getKey());                                            
+                        if(v != null && onExpire != null){
+                            onExpire.run(v);
+                        }
+                    }
+                }                
+            }
+        }
+
+    }
+
+    public interface OnExpire<V>{
+        public abstract void run(V value);
+    }
+}

src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java

@@ -0,0 +1,57 @@
+package net.lightbody.bmp.proxy;
+
+import java.util.HashSet;
+import java.util.Set;
+import net.lightbody.bmp.proxy.util.ExpirableMap;
+import org.junit.After;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ExpirableMapTest {    
+    private Set<String> strings = new HashSet<>();
+    private ExpirableMap<Integer, String> m;
+
+    @Before
+    public void setUp() throws Exception {
+        m = new ExpirableMap<>(1, 1, new ExpirableMap.OnExpire<String>(){
+            @Override
+            public void run(String s) {
+                ExpirableMapTest.this.strings.add(s);
+            }
+        });
+    }
+
+    @After
+    public void tearDown() throws Exception {  
+        m.stop();
+    }
+
+    @Test
+    public void testKeyExpiration() throws Exception {
+
+        m.put(1, "a");
+        m.put(1, "b");
+        String s = m.putIfAbsent(2, "c");
+
+        assertNull(s);
+
+        s = m.putIfAbsent(2, "d");
+
+        assertEquals("c", s);
+
+        Thread.sleep(2000);
+
+        assertEquals(0, m.size());
+
+        assertFalse(strings.contains("a"));
+
+        assertTrue(strings.contains("b"));
+
+        assertTrue(strings.contains("c"));
+
+    }
+}