diff --git a/engine/src/main/java/org/archive/crawler/prefetch/QuotaEnforcer.java b/engine/src/main/java/org/archive/crawler/prefetch/QuotaEnforcer.java index 9fd434e8..91756e57 100644 --- a/engine/src/main/java/org/archive/crawler/prefetch/QuotaEnforcer.java +++ b/engine/src/main/java/org/archive/crawler/prefetch/QuotaEnforcer.java @@ -39,7 +39,6 @@ import org.springframework.beans.factory.annotation.Autowired; * blocks the current URI's processing with S_BLOCKED_BY_QUOTA. * * @author gojomo - * @version $Date$, $Revision$ */ public class QuotaEnforcer extends Processor { @SuppressWarnings("unused") @@ -57,21 +56,29 @@ public class QuotaEnforcer extends Processor { protected static final int SUCCESS_KB = 1; protected static final int RESPONSES = 2; protected static final int RESPONSE_KB = 3; + protected static final int NOVEL_KB = 4; + protected static final int NOVEL_URLS = 5; private static final String SERVER_MAX_FETCH_SUCCESSES = "serverMaxFetchSuccesses"; private static final String SERVER_MAX_SUCCESS_KB = "serverMaxSuccessKb"; private static final String SERVER_MAX_FETCH_RESPONSES = "serverMaxFetchResponses"; private static final String SERVER_MAX_ALL_KB = "serverMaxAllKb"; + private static final String SERVER_MAX_NOVEL_KB = "serverMaxNovelKb"; + private static final String SERVER_MAX_NOVEL_URLS = "serverMaxNovelUrls"; private static final String HOST_MAX_FETCH_SUCCESSES = "hostMaxFetchSuccesses"; private static final String HOST_MAX_SUCCESS_KB = "hostMaxSuccessKb"; private static final String HOST_MAX_FETCH_RESPONSES = "hostMaxFetchResponses"; private static final String HOST_MAX_ALL_KB = "hostMaxAllKb"; + private static final String HOST_MAX_NOVEL_KB = "hostMaxNovelKb"; + private static final String HOST_MAX_NOVEL_URLS = "hostMaxNovelUrls"; private static final String GROUP_MAX_FETCH_SUCCESSES = "groupMaxFetchSuccesses"; private static final String GROUP_MAX_SUCCESS_KB = "groupMaxSuccessKb"; private static final String GROUP_MAX_FETCH_RESPONSES = "groupMaxFetchResponses"; private static final String GROUP_MAX_ALL_KB = "groupMaxAllKb"; + private static final String GROUP_MAX_NOVEL_KB = "serverMaxNovelKb"; + private static final String GROUP_MAX_NOVEL_URLS = "serverMaxNovelUrls"; protected static final String[][] keys = new String[][] { { @@ -79,22 +86,27 @@ public class QuotaEnforcer extends Processor { SERVER_MAX_FETCH_SUCCESSES, SERVER_MAX_SUCCESS_KB, SERVER_MAX_FETCH_RESPONSES, - SERVER_MAX_ALL_KB + SERVER_MAX_ALL_KB, + SERVER_MAX_NOVEL_KB, + SERVER_MAX_NOVEL_URLS, }, { //"host" HOST_MAX_FETCH_SUCCESSES, HOST_MAX_SUCCESS_KB, HOST_MAX_FETCH_RESPONSES, - HOST_MAX_ALL_KB - , + HOST_MAX_ALL_KB, + HOST_MAX_NOVEL_KB, + HOST_MAX_NOVEL_URLS, }, { //"group" GROUP_MAX_FETCH_SUCCESSES, GROUP_MAX_SUCCESS_KB, GROUP_MAX_FETCH_RESPONSES, - GROUP_MAX_ALL_KB + GROUP_MAX_ALL_KB, + GROUP_MAX_NOVEL_KB, + GROUP_MAX_NOVEL_URLS, } }; @@ -158,6 +170,26 @@ public class QuotaEnforcer extends Processor { kp.put(SERVER_MAX_ALL_KB,max); } + { + setServerMaxNovelKb(-1L); // no limit + } + public long getServerMaxNovelKb() { + return (Long) kp.get(SERVER_MAX_NOVEL_KB); + } + public void setServerMaxNovelKb(long max) { + kp.put(SERVER_MAX_NOVEL_KB, max); + } + + { + setServerMaxNovelUrls(-1L); // no limit + } + public long getServerMaxNovelUrls() { + return (Long) kp.get(SERVER_MAX_NOVEL_URLS); + } + public void setServerMaxNovelUrls(long max) { + kp.put(SERVER_MAX_NOVEL_URLS, max); + } + /** * Maximum number of fetch successes (e.g. 200 responses) to collect from * one host. Default is -1, meaning no limit. @@ -200,6 +232,26 @@ public class QuotaEnforcer extends Processor { kp.put(HOST_MAX_FETCH_RESPONSES,max); } + { + setHostMaxNovelKb(-1L); // no limit + } + public long getHostMaxNovelKb() { + return (Long) kp.get(HOST_MAX_NOVEL_KB); + } + public void setHostMaxNovelKb(long max) { + kp.put(HOST_MAX_NOVEL_KB, max); + } + + { + setHostMaxNovelUrls(-1L); // no limit + } + public long getHostMaxNovelUrls() { + return (Long) kp.get(HOST_MAX_NOVEL_URLS); + } + public void setHostMaxNovelUrls(long max) { + kp.put(HOST_MAX_NOVEL_URLS, max); + } + /** * Maximum amount of response content (incl. error responses) in KB to * collect from one host. Default is -1, meaning no limit. @@ -269,6 +321,27 @@ public class QuotaEnforcer extends Processor { public void setGroupMaxAllKb(long max) { kp.put(GROUP_MAX_ALL_KB,max); } + + { + setGroupMaxNovelKb(-1L); // no limit + } + public long getGroupMaxNovelKb() { + return (Long) kp.get(GROUP_MAX_NOVEL_KB); + } + public void setGroupMaxNovelKb(long max) { + kp.put(GROUP_MAX_NOVEL_KB, max); + } + + { + setGroupMaxNovelUrls(-1L); // no limit + } + public long getGroupMaxNovelUrls() { + return (Long) kp.get(GROUP_MAX_NOVEL_URLS); + } + public void setGroupMaxNovelUrls(long max) { + kp.put(GROUP_MAX_NOVEL_URLS, max); + } + /** * Whether an over-quota situation should result in the containing queue @@ -357,8 +430,10 @@ public class QuotaEnforcer extends Processor { substats.getSuccessBytes()/1024, substats.getFetchResponses(), substats.getTotalBytes()/1024, + substats.getNovelBytes()/1024, + substats.getNovelUrls(), }; - for(int q=SUCCESSES; q<=RESPONSE_KB; q++) { + for(int q=SUCCESSES; q<=NOVEL_URLS; q++) { String key = keys[CAT][q]; if (applyQuota(curi, key, actuals[q])) { return true; diff --git a/engine/src/test/java/org/archive/crawler/prefetch/QuotaEnforcerTest.java b/engine/src/test/java/org/archive/crawler/prefetch/QuotaEnforcerTest.java index 0dcf72fd..9411ce7f 100644 --- a/engine/src/test/java/org/archive/crawler/prefetch/QuotaEnforcerTest.java +++ b/engine/src/test/java/org/archive/crawler/prefetch/QuotaEnforcerTest.java @@ -19,16 +19,402 @@ package org.archive.crawler.prefetch; +import java.io.File; +import java.io.IOException; +import java.io.PrintWriter; +import java.util.HashMap; +import java.util.Map; + +import javax.management.openmbean.CompositeData; + +import org.apache.commons.httpclient.URIException; import org.archive.crawler.framework.CrawlerProcessorTestBase; +import org.archive.crawler.framework.Frontier; +import org.archive.crawler.framework.Frontier.FrontierGroup; +import org.archive.crawler.frontier.FrontierJournal; +import org.archive.modules.CoreAttributeConstants; +import org.archive.modules.CrawlURI; +import org.archive.modules.ProcessResult; +import org.archive.modules.deciderules.DecideRule; +import org.archive.modules.fetcher.DefaultServerCache; +import org.archive.modules.fetcher.FetchStats; +import org.archive.modules.fetcher.FetchStats.Stage; +import org.archive.modules.fetcher.FetchStatusCodes; +import org.archive.modules.net.CrawlHost; +import org.archive.modules.net.CrawlServer; +import org.archive.net.UURIFactory; +import org.archive.util.ObjectIdentityCache; +import org.archive.util.ObjectIdentityMemCache; + +import junit.framework.Assert; /** * Unit test for {@link QuotaEnforcer}. * - * @author pjack + * @contributor pjack + * @contributor nlevitt */ public class QuotaEnforcerTest extends CrawlerProcessorTestBase { - - // TODO TESTME! + static class MockFetchStats extends FetchStats { + private static final long serialVersionUID = 1l; + public void setNovelBytes(long n) { + novelBytes = n; + } + public void setNovelUrls(long n) { + novelUrls = n; + } + } + + static class MockServerCache extends DefaultServerCache { + private static final long serialVersionUID = 1l; + public void setHostFor(String host, CrawlHost crawlHost) { + ((ObjectIdentityMemCache) hosts).getMap().put(host, crawlHost); + } + public void setServerFor(String h, CrawlServer crawlServer) { + ((ObjectIdentityMemCache) servers).getMap().put(h, crawlServer); + } + } + + interface CanSetSubstats { + public void setSubstats(FetchStats stats); + } + + static class MockCrawlHost extends CrawlHost implements CanSetSubstats { + private static final long serialVersionUID = 1l; + public MockCrawlHost(String hostname) { + super(hostname); + } + @Override + public void setSubstats(FetchStats stats) { + this.substats = stats; + } + } + + static class MockCrawlServer extends CrawlServer implements CanSetSubstats { + private static final long serialVersionUID = 1l; + public MockCrawlServer(String h) { + super(h); + } + @Override + public void setSubstats(FetchStats stats) { + this.substats = stats; + } + } + + static class MockFrontierGroup implements FrontierGroup, CanSetSubstats { + private static final long serialVersionUID = 1l; + + protected FetchStats substats = new FetchStats(); + + // the only method that's used in this mock class + @Override + public FetchStats getSubstats() { + return substats; + } + + @Override + public void setSubstats(FetchStats stats) { + this.substats = stats; + } + + @Override + public void tally(CrawlURI curi, Stage stage) { + } + @Override + public void setIdentityCache(ObjectIdentityCache cache) { + } + @Override + public String getKey() { + return null; + } + @Override + public void makeDirty() { + } + } + + static class MockFrontier implements Frontier { + protected Map hostGroups = new HashMap(); + + // the only method that's used in this mock class + @Override + public FrontierGroup getGroup(CrawlURI curi) { + String host; + try { + host = curi.getUURI().getHost(); + MockFrontierGroup group = hostGroups.get(host); + if (group == null) { + group = new MockFrontierGroup(); + hostGroups.put(host, group); + } + return group; + } catch (URIException e) { + Assert.fail(); + return null; + } + } + + @Override + public void start() { + } + @Override + public void stop() { + } + @Override + public boolean isRunning() { + return false; + } + @Override + public void reportTo(PrintWriter writer) throws IOException { + } + @Override + public void shortReportLineTo(PrintWriter pw) throws IOException { + } + @Override + public Map shortReportMap() { + return null; + } + @Override + public String shortReportLegend() { + return null; + } + @Override + public CrawlURI next() throws InterruptedException { + return null; + } + @Override + public boolean isEmpty() { + return false; + } + @Override + public void schedule(CrawlURI caURI) { + } + @Override + public void finished(CrawlURI cURI) { + } + @Override + public long discoveredUriCount() { + return 0; + } + @Override + public long queuedUriCount() { + return 0; + } + @Override + public long futureUriCount() { + return 0; + } + @Override + public long deepestUri() { + return 0; + } + @Override + public long averageDepth() { + return 0; + } + @Override + public float congestionRatio() { + return 0; + } + @Override + public long finishedUriCount() { + return 0; + } + @Override + public long succeededFetchCount() { + return 0; + } + @Override + public long failedFetchCount() { + return 0; + } + @Override + public long disregardedUriCount() { + return 0; + } + @Override + public void importURIs(String params) throws IOException { + } + @Override + public long importRecoverFormat(File source, boolean applyScope, boolean includeOnly, boolean forceFetch, + String acceptTags) throws IOException { + return 0; + } + @Override + public CompositeData getURIsList(String marker, int numberOfMatches, String regex, boolean verbose) { + return null; + } + @Override + public long deleteURIs(String queueRegex, String match) { + return 0; + } + @Override + public void deleted(CrawlURI curi) { + } + @Override + public void considerIncluded(CrawlURI curi) { + } + @Override + public void pause() { + } + @Override + public void unpause() { + } + @Override + public void terminate() { + } + @Override + public FrontierJournal getFrontierJournal() { + return null; + } + @Override + public String getClassKey(CrawlURI cauri) { + return null; + } + @Override + public DecideRule getScope() { + return null; + } + @Override + public void run() { + } + @Override + public void requestState(State target) { + } + @Override + public void beginDisposition(CrawlURI curi) { + } + @Override + public void endDisposition() { + } + } + + // separate methods to make it easier to know what failed + public void testHostNovelKbForceRetire() throws URIException, InterruptedException { + testNovel("kb", "host", true); + } + public void testServerNovelKbForceRetire() throws URIException, InterruptedException { + testNovel("kb", "server", true); + } + public void testGroupNovelKbForceRetire() throws URIException, InterruptedException { + testNovel("kb", "group", true); + } + public void testHostNovelKbNoForceRetire() throws URIException, InterruptedException { + testNovel("kb", "host", false); + } + public void testServerNovelKbNoForceRetire() throws URIException, InterruptedException { + testNovel("kb", "server", false); + } + public void testGroupNovelKbNoForceRetire() throws URIException, InterruptedException { + testNovel("kb", "group", false); + } + + public void testHostNovelUrlsForceRetire() throws URIException, InterruptedException { + testNovel("urls", "host", true); + } + public void testServerNovelUrlsForceRetire() throws URIException, InterruptedException { + testNovel("urls", "server", true); + } + public void testGroupNovelUrlsForceRetire() throws URIException, InterruptedException { + testNovel("urls", "group", true); + } + public void testHostNovelUrlsNoForceRetire() throws URIException, InterruptedException { + testNovel("urls", "host", false); + } + public void testServerNovelUrlsNoForceRetire() throws URIException, InterruptedException { + testNovel("urls", "server", false); + } + public void testGroupNovelUrlsNoForceRetire() throws URIException, InterruptedException { + testNovel("urls", "group", false); + } + + protected void testNovel(String urlsOrKb, String hostServerOrGroup, boolean forceRetire) throws URIException, InterruptedException { + QuotaEnforcer qe = new QuotaEnforcer(); + MockServerCache serverCache = new MockServerCache(); + qe.setServerCache(serverCache); + MockFrontier frontier = new MockFrontier(); + qe.setFrontier(frontier); + qe.setForceRetire(forceRetire); + + // sanity check + assertTrue("urls".equals(urlsOrKb) || "kb".equals(urlsOrKb)); + assertTrue("host".equals(hostServerOrGroup) || "server".equals(hostServerOrGroup) || "group".equals(hostServerOrGroup)); + + if ("host".equals(hostServerOrGroup)) { + if ("urls".equals(urlsOrKb)) { + qe.setHostMaxNovelUrls(1); + assertEquals(1, qe.getHostMaxNovelUrls()); + } else if ("kb".equals(urlsOrKb)) { + qe.setHostMaxNovelKb(100); + assertEquals(100, qe.getHostMaxNovelKb()); + } + } else if ("server".equals(hostServerOrGroup)) { + if ("urls".equals(urlsOrKb)) { + qe.setServerMaxNovelUrls(1); + assertEquals(1, qe.getServerMaxNovelUrls()); + } else { + qe.setServerMaxNovelKb(100); + assertEquals(100, qe.getServerMaxNovelKb()); + } + } else if ("group".equals(hostServerOrGroup)) { + if ("urls".equals(urlsOrKb)) { + qe.setGroupMaxNovelUrls(1); + assertEquals(1, qe.getGroupMaxNovelUrls()); + } else { + qe.setGroupMaxNovelKb(100); + assertEquals(100, qe.getGroupMaxNovelKb()); + } + } + + // nothing accumulated yet + CrawlURI curi = new CrawlURI(UURIFactory.getInstance("http://example.com/1")); + ProcessResult res = qe.process(curi); + assertEquals(ProcessResult.PROCEED, res); + if (forceRetire) { + assertNull(curi.getData().get(CoreAttributeConstants.A_FORCE_RETIRE)); + } else { + assertEquals(FetchStatusCodes.S_UNATTEMPTED, curi.getFetchStatus()); + } + + // we do all this to set only the stats value we're testing, to avoid + // quotas checking the wrong thing but tests passing anyway + CanSetSubstats thing; + if ("host".equals(hostServerOrGroup)) { + thing = new MockCrawlHost("example.com"); + serverCache.setHostFor("example.com", (CrawlHost) thing); + } else if ("server".equals(hostServerOrGroup)) { + thing = new MockCrawlServer("example.com"); + serverCache.setServerFor("example.com", (CrawlServer) thing); + } else { + thing = (MockFrontierGroup) frontier.getGroup(curi); + } + MockFetchStats stats = new MockFetchStats(); + if ("urls".equals(urlsOrKb)) { + stats.setNovelUrls(1); + } else { + stats.setNovelBytes(200000); + } + thing.setSubstats(stats); + + // another uri from same host should hit quota + curi = new CrawlURI(UURIFactory.getInstance("http://example.com/2")); + res = qe.process(curi); + assertEquals(ProcessResult.FINISH, res); + if (forceRetire) { + assertTrue((Boolean) curi.getData().get(CoreAttributeConstants.A_FORCE_RETIRE)); + } else { + assertEquals(FetchStatusCodes.S_BLOCKED_BY_QUOTA, curi.getFetchStatus()); + } + + // some other host has not hit quota + curi = new CrawlURI(UURIFactory.getInstance("http://example.org/")); + res = qe.process(curi); + assertEquals(ProcessResult.PROCEED, res); + if (forceRetire) { + assertNull(curi.getData().get(CoreAttributeConstants.A_FORCE_RETIRE)); + } else { + assertEquals(FetchStatusCodes.S_UNATTEMPTED, curi.getFetchStatus()); + } + } }