Merge pull request #626 from internetarchive/checkpoint-on-shutdown

Add checkpoint on shutdown and --checkpoint CLI option

engine/src/main/java/org/archive/crawler/Heritrix.java

@@ -128,6 +128,10 @@ public class Heritrix {
                 "\"password\" (which leaves username as the default 'admin'), " +
                 "\"username:password\", or \"@filename\" for a file that " +
                 "includes the single line \"username:password\". ");
+        options.addOption("c", "checkpoint", true,
+                "Recovers from the given checkpoint. May only be used with the " +
+                "--run-job option. The special value 'latest' will recover the " +
+                "last checkpoint or if none exist will launch a new crawl.");
         options.addOption("j", "jobs-dir", true, "The jobs directory.  " +
                         "Defaults to ./jobs");
         options.addOption("l", "logging-properties", true, 
@@ -265,6 +269,11 @@ public class Heritrix {
             System.exit(1);
             authPassword = ""; // suppresses uninitialized warning
         }
+
+        if (cl.hasOption('c') && !cl.hasOption('r')) {
+            System.err.println("Cannot use --checkpoint without --run-job.");
+            System.exit(1);
+        }
         
         File jobsDir = null; 
         if (cl.hasOption('j')) {
@@ -374,12 +383,21 @@ public class Heritrix {
             }
             if (cl.hasOption('r')) {
                 String jobName = cl.getOptionValue('r');
-                engine.requestLaunch(jobName);
                 CrawlJob job = engine.getJob(jobName);
-                if (job == null || job.getCrawlController() == null) {
+                if (job == null) {
+                    System.err.println("Job not found: " + jobName);
+                    System.exit(1);
+                }
+                job.validateConfiguration();
+                if (cl.hasOption('c')) {
+                    job.getCheckpointService().setRecoveryCheckpointByName(cl.getOptionValue('c'));
+                }
+                job.launch();
+                if (job.getCrawlController() == null) {
                     System.err.println("Failed to launch job: " + jobName);
                     System.exit(1);
                 }
+
                 job.getCrawlController().requestCrawlResume();
                 engine.waitForNoRunningJobs(0);
                 engine.shutdown();

engine/src/main/java/org/archive/crawler/framework/CheckpointService.java

@@ -56,7 +56,7 @@ import org.springframework.validation.Validator;
  * build and before launch (setRecoveryCheckpointByName).
  * 
  * Offers optional automatic checkpointing at a configurable interval 
- * in minutes. 
+ * in minutes.
  * 
  * @author stack
  * @author gojomo
@@ -80,6 +80,8 @@ public class CheckpointService implements Lifecycle, ApplicationContextAware, Ha
     protected TimerTask checkpointTask = null; 
     protected ConfigPath checkpointsDir = 
         new ConfigPath("checkpoints subdirectory","checkpoints");
+    protected Thread shutdownHook;
+
     public ConfigPath getCheckpointsDir() {
         return checkpointsDir;
     }
@@ -106,6 +108,21 @@ public class CheckpointService implements Lifecycle, ApplicationContextAware, Ha
             setupCheckpointTask();
         }
     }
+
+    protected boolean checkpointOnShutdown = false;
+
+    public boolean getCheckpointOnShutdown() {
+        return checkpointOnShutdown;
+    }
+
+    /**
+     * Whether a checkpoint should be made when the JVM is shutdown.
+     * Default is false.
+     */
+    public void setCheckpointOnShutdown(boolean checkpointOnShutdown) {
+        this.checkpointOnShutdown = checkpointOnShutdown;
+        setupShutdownHook();
+    }
     
     protected boolean forgetAllButLatest = false;
     public boolean getForgetAllButLatest() {
@@ -179,6 +196,7 @@ public class CheckpointService implements Lifecycle, ApplicationContextAware, Ha
         }   
         this.isRunning = true; 
         setupCheckpointTask();
+        setupShutdownHook();
     }
     
     /**
@@ -221,12 +239,31 @@ public class CheckpointService implements Lifecycle, ApplicationContextAware, Ha
     public synchronized boolean isRunning() {
         return isRunning; 
     }
-    
+
+    protected synchronized void setupShutdownHook() {
+        if (!checkpointOnShutdown || shutdownHook != null || !isRunning) return;
+        shutdownHook = new Thread(() -> {
+            if (!checkpointOnShutdown) return;
+            LOGGER.info("Checkpointing on shutdown");
+            try {
+                // pause first to ensure no crawling occurs between the checkpoint and process termination
+                controller.requestCrawlPause();
+                requestCrawlCheckpoint();
+            } catch (Exception e) {
+                LOGGER.severe("Failed to checkpoint on shutdown: " + e.getMessage());
+            }
+        });
+        Runtime.getRuntime().addShutdownHook(shutdownHook);
+    }
     
     public synchronized void stop() {
         LOGGER.info("Cleaned up Checkpoint TimerThread.");
         this.timer.cancel();
-        this.isRunning = false; 
+        this.isRunning = false;
+        if (shutdownHook != null) {
+            Runtime.getRuntime().removeShutdownHook(shutdownHook);
+            shutdownHook = null;
+        }
     }
     
     /**

engine/src/main/resources/org/archive/crawler/restlet/profile-crawler-beans.cxml

@@ -564,6 +564,7 @@ http://example.example/example
  <bean id="checkpointService" 
    class="org.archive.crawler.framework.CheckpointService">
   <!-- <property name="checkpointIntervalMinutes" value="-1"/> -->
+  <!-- <property name="checkpointOnShutdown" value="true"/> -->
   <!-- <property name="checkpointsDir" value="checkpoints"/> -->
   <!-- <property name="forgetAllButLatest" value="true"/> -->
  </bean>