Add a failure and complete hook for better diagnostics

src/main/java/nl/aerius/print/DriverHook.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright the State of the Netherlands
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see http://www.gnu.org/licenses/.
+ */
+package nl.aerius.print;
+
+import com.intuit.karate.driver.DevToolsDriver;
+
+@FunctionalInterface
+public interface DriverHook {
+  void accept(DevToolsDriver driver, String url, String phase, Throwable cause);
+}

src/main/java/nl/aerius/print/ExportJob.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -54,6 +55,9 @@ public class ExportJob {
 
   private boolean saved;
 
+  private DriverHook completeHook;
+  private DriverHook failureHook;
+
   public static ExportJob create() {
     return new ExportJob();
   }
@@ -121,24 +125,8 @@ public SnapshotJob snapshot(final Map<String, Object> dimensions) {
 
     LOG.info("Exporting graphic from: {}", url);
 
-    final DevToolsDriver chrome = fetchChrome();
-    try {
-      chrome.setUrl(url);
-
-      // Set a default wait
-      if (waitForComplete == null) {
-        completeViaSleep();
-      }
-      waitForComplete(chrome);
-
-      name = handle + ".png";
-      exportResult = chrome.screenshot();
-    } catch (final RuntimeException e) {
-      LOG.error("Unrecoverable failure while sending snapshot job to chrome instance.", e);
-      throw new RuntimeException("Could not finish web document snapshot.", e);
-    } finally {
-      chrome.quit();
-    }
+    name = handle + ".png";
+    exportResult = runExport(true, d -> d.screenshot(), "snapshot");
 
     return new SnapshotJob(this);
   }
@@ -154,23 +142,8 @@ public PrintJob print(final Map<String, Object> printParams) {
     ensureHandle();
     exported = true;
 
-    final DevToolsDriver chrome = fetchChrome();
-    try {
-      chrome.setUrl(url);
-
-      if (waitForComplete == null) {
-        completeViaIndicator();
-      }
-      waitForComplete(chrome);
-
-      name = handle + ".pdf";
-      exportResult = chrome.pdf(printParams);
-    } catch (final RuntimeException e) {
-      LOG.error("Unrecoverable failure while sending print job to chrome instance.", e);
-      throw new RuntimeException("Could not finish PDF export.", e);
-    } finally {
-      chrome.quit();
-    }
+    name = handle + ".pdf";
+    exportResult = runExport(false, d -> d.pdf(printParams), "print");
 
     return new PrintJob(this);
   }
@@ -222,6 +195,18 @@ public ExportJob driverOptions(final Map<String, Object> driverOptions) {
     return this;
   }
 
+  public ExportJob completeHook(final DriverHook hook) {
+    checkExported();
+    this.completeHook = hook;
+    return this;
+  }
+
+  public ExportJob failureHook(final DriverHook hook) {
+    checkExported();
+    this.failureHook = hook;
+    return this;
+  }
+
   /**
    * TODO Provide graceful failure functionality
    */
@@ -260,6 +245,45 @@ private DevToolsDriver fetchChrome() {
     return chrome;
   }
 
+  private byte[] runExport(final boolean useSleepWait, final Function<DevToolsDriver, byte[]> exporter,
+      final String failurePhase) {
+    final DevToolsDriver chrome = fetchChrome();
+    try {
+      chrome.setUrl(url);
+
+      if (waitForComplete == null) {
+        if (useSleepWait) {
+          completeViaSleep();
+        } else {
+          completeViaIndicator();
+        }
+      }
+      waitForComplete(chrome);
+
+      if (completeHook != null) {
+        try {
+          completeHook.accept(chrome, url, "complete", null);
+        } catch (final RuntimeException ignored) {
+          LOG.warn("Failure during completeHook execution, ignoring.", ignored);
+        }
+      }
+
+      return exporter.apply(chrome);
+    } catch (final RuntimeException e) {
+      if (failureHook != null) {
+        try {
+          failureHook.accept(chrome, url, failurePhase, e);
+        } catch (final RuntimeException ignored) {
+          LOG.warn("Failure during failureHook execution, ignoring.", ignored);
+        }
+      }
+      LOG.error("Unrecoverable failure while executing export.", e);
+      throw e;
+    } finally {
+      chrome.quit();
+    }
+  }
+
   public String outputDocument() {
     return outputDocument;
   }