Merge pull request #165 from galgeek/amqpWaiter

AMQP URL Waiter

contrib/src/main/java/org/archive/crawler/event/AMQPUrlReceivedEvent.java

@@ -0,0 +1,47 @@
+/*
+ *  This file is part of the Heritrix web crawler (crawler.archive.org).
+ *
+ *  Licensed to the Internet Archive (IA) by one or more individual
+ *  contributors.
+ *
+ *  The IA licenses this file to You under the Apache License, Version 2.0
+ *  (the "License"); you may not use this file except in compliance with
+ *  the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.archive.crawler.event;
+
+import org.archive.crawler.frontier.AMQPUrlReceiver;
+import org.archive.modules.CrawlURI;
+import org.springframework.context.ApplicationEvent;
+
+import com.rabbitmq.client.AMQP;
+import com.rabbitmq.client.AMQP.BasicProperties;
+
+/**
+ * ApplicationEvent published when AMQPUrlReceiver receives a URL.
+ * Other modules can observe this event to learn when AMQPUrlReceiver receives a URL.
+ *
+ * @contributor galgeek
+ */
+public class AMQPUrlReceivedEvent extends ApplicationEvent {
+    private static final long serialVersionUID = 1L;
+
+    protected CrawlURI curi;
+    public CrawlURI getCuri() {
+        return curi;
+    }
+
+    public AMQPUrlReceivedEvent(AMQPUrlReceiver source, CrawlURI curi) {
+        super(source);
+        this.curi = curi;
+    }
+}

contrib/src/main/java/org/archive/crawler/frontier/AMQPUrlReceiver.java

@@ -32,6 +32,7 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import org.apache.commons.httpclient.URIException;
+import org.archive.crawler.event.AMQPUrlReceivedEvent;
 import org.archive.crawler.event.CrawlStateEvent;
 import org.archive.crawler.postprocessor.CandidatesProcessor;
 import org.archive.modules.CrawlURI;
@@ -45,7 +46,11 @@ import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
+import org.springframework.context.ApplicationEvent;
 import org.springframework.context.ApplicationListener;
+import org.springframework.beans.BeansException;
 import org.springframework.context.Lifecycle;
 
 import com.rabbitmq.client.AMQP.BasicProperties;
@@ -60,7 +65,8 @@ import com.rabbitmq.client.ShutdownSignalException;
 /**
  * @contributor nlevitt
  */
-public class AMQPUrlReceiver implements Lifecycle, ApplicationListener<CrawlStateEvent> {
+public class AMQPUrlReceiver
+	implements Lifecycle, ApplicationContextAware, ApplicationListener<CrawlStateEvent> {
 
     @SuppressWarnings("unused")
     private static final long serialVersionUID = 2L;
@@ -70,6 +76,11 @@ public class AMQPUrlReceiver implements Lifecycle, ApplicationListener<CrawlStat
 
     public static final String A_RECEIVED_FROM_AMQP = "receivedFromAMQP";
 
+    protected ApplicationContext appCtx;
+    public void setApplicationContext(ApplicationContext appCtx) throws BeansException {
+        this.appCtx = appCtx;
+    }
+
     protected CandidatesProcessor candidates;
     public CandidatesProcessor getCandidates() {
         return candidates;
@@ -331,6 +342,7 @@ public class AMQPUrlReceiver implements Lifecycle, ApplicationListener<CrawlStat
                     CrawlURI curi = makeCrawlUri(jo);
                     KeyedProperties.clearAllOverrideContexts();
                     candidates.runCandidateChain(curi, null);
+                    appCtx.publishEvent(new AMQPUrlReceivedEvent(AMQPUrlReceiver.this, curi));
                 } catch (URIException e) {
                     logger.log(Level.WARNING,
                             "problem creating CrawlURI from json received via AMQP "

contrib/src/main/java/org/archive/modules/AMQPUrlWaiter.java

@@ -0,0 +1,82 @@
+/*
+ *  This file is part of the Heritrix web crawler (crawler.archive.org).
+ *
+ *  Licensed to the Internet Archive (IA) by one or more individual
+ *  contributors.
+ *
+ *  The IA licenses this file to You under the Apache License, Version 2.0
+ *  (the "License"); you may not use this file except in compliance with
+ *  the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.archive.modules;
+
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import org.archive.crawler.event.AMQPUrlReceivedEvent;
+import org.archive.crawler.event.StatSnapshotEvent;
+import org.archive.crawler.framework.CrawlController;
+import org.archive.crawler.framework.CrawlStatus;
+import org.archive.crawler.framework.Frontier;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationListener;
+
+/**
+ * Bean to enforce a wait for Umbra's amqp queue
+ *
+ * @contributor galgeek
+ */
+public class AMQPUrlWaiter implements ApplicationListener<ApplicationEvent> {
+
+    public AMQPUrlWaiter() {}
+
+    static protected final Logger logger = Logger.getLogger(AMQPUrlWaiter.class.getName());
+
+    protected int urlsReceived = 0;
+
+    protected CrawlController controller;
+    public CrawlController getCrawlController() {
+        return this.controller;
+    }
+    @Autowired
+    public void setCrawlController(CrawlController controller) {
+        this.controller = controller;
+    }
+
+    protected Frontier frontier;
+    public Frontier getFrontier() {
+        return this.frontier;
+    }
+    /** Autowired frontier, needed to determine when a url is finished. */
+    @Autowired
+    public void setFrontier(Frontier frontier) {
+        this.frontier = frontier;
+    }
+
+    @Override
+    public void onApplicationEvent(ApplicationEvent event) {
+        if (event instanceof AMQPUrlReceivedEvent) {
+            urlsReceived += 1;
+        } else if (event instanceof StatSnapshotEvent) {
+            checkAMQPUrlWait();
+        }
+    }
+
+    protected void checkAMQPUrlWait() {
+        if (frontier.isEmpty() && urlsReceived > 0) {
+            logger.info("frontier is empty and we have received " + urlsReceived +
+                        " urls from AMQP, stopping crawl with status " + CrawlStatus.FINISHED);
+            controller.requestCrawlStop(CrawlStatus.FINISHED);
+        }
+    }
+}