diff --git a/browsermob-core/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java b/browsermob-core/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java deleted file mode 100644 index 64de4c572..000000000 --- a/browsermob-core/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java +++ /dev/null @@ -1,98 +0,0 @@ -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.ThreadFactory; -import java.util.concurrent.TimeUnit; - -public class ExpirableMap extends ConcurrentHashMap{ - public final static int DEFAULT_CHECK_INTERVAL = 10*60; - public final static int DEFAULT_TTL = 30*60; - private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor( - new ThreadFactory() { - @Override - public Thread newThread(Runnable r) { - Thread t = new Thread(r); - t.setDaemon(true); - return t; - } - } - ); - private final long ttl; - private final Map expires; - private final OnExpire onExpire; - - public ExpirableMap(int ttl, int checkInterval, OnExpire onExpire) { - this.ttl = ttl*1000; - this.onExpire = onExpire; - expires = new HashMap<>(); - scheduler.scheduleWithFixedDelay(new Worker(), checkInterval, checkInterval, TimeUnit.SECONDS); - } - - public ExpirableMap(int ttl, OnExpire onExpire) { - this(ttl, DEFAULT_CHECK_INTERVAL, onExpire); - } - - public ExpirableMap(OnExpire 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); - } - } - - private class Worker implements Runnable{ - - @Override - public void run() { - Map m; - synchronized(ExpirableMap.this){ - m = new HashMap<>(expires); - } - Long now = new Date().getTime(); - for(Entry 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{ - public abstract void run(V value); - } -} diff --git a/browsermob-core/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java b/browsermob-core/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java deleted file mode 100644 index 90cf8a24c..000000000 --- a/browsermob-core/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java +++ /dev/null @@ -1,51 +0,0 @@ -package net.lightbody.bmp.proxy; - -import java.util.HashSet; -import java.util.Set; -import net.lightbody.bmp.proxy.util.ExpirableMap; -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 strings = new HashSet<>(); - private ExpirableMap m; - - @Before - public void setUp() throws Exception { - m = new ExpirableMap<>(1, 1, new ExpirableMap.OnExpire(){ - @Override - public void run(String s) { - ExpirableMapTest.this.strings.add(s); - } - }); - } - - @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")); - - } -} diff --git a/browsermob-rest/pom.xml b/browsermob-rest/pom.xml index b22c3b71d..d85b5f9a7 100644 --- a/browsermob-rest/pom.xml +++ b/browsermob-rest/pom.xml @@ -68,6 +68,11 @@ 3.2 + + com.google.guava + guava + + org.apache.logging.log4j log4j-api diff --git a/browsermob-rest/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java b/browsermob-rest/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java index 83cdb25b5..30e70bd60 100644 --- a/browsermob-rest/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java +++ b/browsermob-rest/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java @@ -1,22 +1,29 @@ package net.lightbody.bmp.proxy; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.RemovalListener; +import com.google.common.cache.RemovalNotification; import com.google.inject.Inject; import com.google.inject.Provider; import com.google.inject.Singleton; import com.google.inject.name.Named; +import net.lightbody.bmp.exception.ProxyExistsException; +import net.lightbody.bmp.exception.ProxyPortsExhaustedException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.lang.ref.WeakReference; import java.net.InetAddress; import java.net.UnknownHostException; import java.util.Collection; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; - -import net.lightbody.bmp.exception.ProxyExistsException; -import net.lightbody.bmp.exception.ProxyPortsExhaustedException; -import net.lightbody.bmp.proxy.util.ExpirableMap; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; @Singleton public class ProxyManager { @@ -26,27 +33,99 @@ public class ProxyManager { private final int minPort; private final int maxPort; private final Provider proxyServerProvider; + // retain a reference to the Cache to allow the ProxyCleanupTask to .cleanUp(), since asMap() is just a view into the cache. + // it would seem to make sense to pass the newly-built Cache directly to the ProxyCleanupTask and have it retain a WeakReference to it, and + // only maintain a reference to the .asMap() result in this class. puzzlingly, however, the Cache can actually get garbage collected + // before the .asMap() view of it does. + private final Cache proxyCache; private final ConcurrentMap proxies; + /** + * Interval at which expired proxy checks will actively clean up expired proxies. Proxies may still be cleaned up when accessing the + * proxies map. + */ + private static final int EXPIRED_PROXY_CLEANUP_INTERVAL_SECONDS = 60; + + // Initialize-on-demand a single thread executor that will create a daemon thread to clean up expired proxies. Since the resulting executor + // is a singleton, there will at most one thread to service all ProxyManager instances. + private static class ScheduledExecutorHolder { + private static final ScheduledExecutorService expiredProxyCleanupExecutor = Executors.newSingleThreadScheduledExecutor(new ThreadFactory() { + @Override + public Thread newThread(Runnable r) { + Thread thread = Executors.defaultThreadFactory().newThread(r); + thread.setName("expired-proxy-cleanup-thread"); + thread.setDaemon(true); + return thread; + } + }); + } + + // static inner class to prevent leaking ProxyManager instances to the cleanup task + private static class ProxyCleanupTask implements Runnable { + // using a WeakReference that will indicate to us when the Cache (and thus its ProxyManager) has been garbage + // collected, allowing this cleanup task to kill itself + private final WeakReference> proxyCache; + + public ProxyCleanupTask(Cache cache) { + this.proxyCache = new WeakReference>(cache); + } + + @Override + public void run() { + Cache cache = proxyCache.get(); + if (cache != null) { + try { + cache.cleanUp(); + } catch (RuntimeException e) { + LOG.warn("Error occurred while attempting to clean up expired proxies", e); + } + } else { + // the cache instance was garbage collected, so it no longer needs to be cleaned up. throw an exception + // to prevent the scheduled executor from re-scheduling this cleanup + LOG.info("Proxy Cache was garbage collected. No longer cleaning up expired proxies for unused ProxyManager."); + + throw new RuntimeException("Exiting ProxyCleanupTask"); + } + } + } + @Inject - public ProxyManager(Provider proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort, @Named("ttl") Integer ttl) { + public ProxyManager(Provider proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort, final @Named("ttl") Integer ttl) { this.proxyServerProvider = proxyServerProvider; this.minPort = minPort; this.maxPort = maxPort; - this.lastPort = maxPort; - this.proxies = ttl > 0 ? - new ExpirableMap(ttl, new ExpirableMap.OnExpire(){ - @Override - public void run(ProxyServer proxy) { - try { - LOG.debug("Expiring ProxyServer `{}`...", proxy.getPort()); + this.lastPort = maxPort; + if (ttl > 0) { + // proxies should be evicted after the specified ttl, so set up an evicting cache and a listener to stop the proxies when they're evicted + RemovalListener removalListener = new RemovalListener () { + public void onRemoval(RemovalNotification removal) { + try { + ProxyServer proxy = removal.getValue(); + if (proxy != null) { + LOG.info("Expiring ProxyServer on port {} after {} seconds without activity", proxy.getPort(), ttl); proxy.stop(); - } catch (Exception ex) { - LOG.warn("Error while stopping an expired proxy", ex); } + } catch (Exception ex) { + LOG.warn("Error while stopping an expired proxy on port " + removal.getKey(), ex); } - }) : - new ConcurrentHashMap(); + } + }; + + this.proxyCache = CacheBuilder.newBuilder() + .expireAfterAccess(ttl, TimeUnit.SECONDS) + .removalListener(removalListener) + .build(); + + this.proxies = proxyCache.asMap(); + + // schedule the asynchronous proxy cleanup task + ScheduledExecutorHolder.expiredProxyCleanupExecutor.scheduleWithFixedDelay(new ProxyCleanupTask(proxyCache), + EXPIRED_PROXY_CLEANUP_INTERVAL_SECONDS, EXPIRED_PROXY_CLEANUP_INTERVAL_SECONDS, TimeUnit.SECONDS); + } else { + this.proxies = new ConcurrentHashMap(); + // nothing to timeout, so no Cache + this.proxyCache = null; + } } public ProxyServer create(Map options, Integer port, String bindAddr) { diff --git a/browsermob-rest/src/test/java/net/lightbody/bmp/proxy/ExpiringProxyTest.java b/browsermob-rest/src/test/java/net/lightbody/bmp/proxy/ExpiringProxyTest.java new file mode 100644 index 000000000..9b59273c3 --- /dev/null +++ b/browsermob-rest/src/test/java/net/lightbody/bmp/proxy/ExpiringProxyTest.java @@ -0,0 +1,78 @@ +package net.lightbody.bmp.proxy; + +import com.google.inject.Provider; +import org.junit.Before; +import org.junit.Test; + +import java.util.Collections; +import java.util.Random; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +public class ExpiringProxyTest { + @Test + public void testExpiredProxyStops() throws InterruptedException { + int minPort = new Random().nextInt(50000) + 10000; + + ProxyManager proxyManager = new ProxyManager(new Provider() { + @Override + public ProxyServer get() { + return new ProxyServer(); + } + }, + minPort, + minPort + 100, + 2); + + ProxyServer proxy = proxyManager.create(Collections.emptyMap()); + int port = proxy.getPort(); + + ProxyServer retrievedProxy = proxyManager.get(port); + + assertEquals("ProxyManager did not return the expected proxy instance", proxy, retrievedProxy); + + Thread.sleep(2500); + + // explicitly create a new proxy to cause a write to the cache. cleanups happen on "every" write and "occasional" reads, so force a cleanup by writing. + int newPort = proxyManager.create(Collections.emptyMap()).getPort(); + proxyManager.delete(newPort); + + ProxyServer expiredProxy = proxyManager.get(port); + + assertNull("ProxyManager did not expire proxy as expected", expiredProxy); + } + + @Test + public void testZeroTtlProxyDoesNotExpire() throws InterruptedException { + int minPort = new Random().nextInt(50000) + 10000; + + ProxyManager proxyManager = new ProxyManager(new Provider() { + @Override + public ProxyServer get() { + return new ProxyServer(); + } + }, + minPort, + minPort + 100, + 0); + + ProxyServer proxy = proxyManager.create(Collections.emptyMap()); + int port = proxy.getPort(); + + ProxyServer retrievedProxy = proxyManager.get(port); + + assertEquals("ProxyManager did not return the expected proxy instance", proxy, retrievedProxy); + + Thread.sleep(2500); + + // explicitly create a new proxy to cause a write to the cache. cleanups happen on "every" write and "occasional" reads, so force a cleanup by writing. + int newPort = proxyManager.create(Collections.emptyMap()).getPort(); + proxyManager.delete(newPort); + + ProxyServer nonExpiredProxy = proxyManager.get(port); + + assertEquals("ProxyManager did not return the expected proxy instance", proxy, nonExpiredProxy); + } + +} diff --git a/pom.xml b/pom.xml index bd04eabd3..18dac187f 100644 --- a/pom.xml +++ b/pom.xml @@ -167,6 +167,12 @@ ${slf4j.version} + + com.google.guava + guava + 18.0 + + org.apache.logging.log4j log4j-api