- track uncommitted crawl ids in separate key

- cleanup uncommitted crawl id keys when crawl is canceled
- simplify final exit checks, for operations on final crawler exit

Author: ikreymer

src/crawler.ts

@@ -644,6 +644,7 @@ export class Crawler {
       if (!finished) {
         if (canceled) {
           status = "canceled";
+          await this.cleanupOnCancel();
         } else if (stopped) {
           status = "done";
           logger.info("Crawl gracefully stopped on request");
@@ -1612,7 +1613,17 @@ self.__bx_behaviors.selectMainBehavior();
 
   async checkCanceled() {
     if (this.crawlState && (await this.crawlState.isCrawlCanceled())) {
+      await this.cleanupOnCancel();
       await this.setStatusAndExit(ExitCodes.Success, "canceled");
+      return true;
+    }
+
+    return false;
+  }
+
+  async cleanupOnCancel() {
+    if (this.deduping) {
+      await this.crawlState.clearUncommitted();
     }
   }
 
@@ -1659,8 +1670,7 @@ self.__bx_behaviors.selectMainBehavior();
       return false;
     }
 
-    if (await this.crawlState.isCrawlCanceled()) {
-      await this.setStatusAndExit(ExitCodes.Success, "canceled");
+    if (await this.checkCanceled()) {
       return false;
     }
 
@@ -1728,6 +1738,9 @@ self.__bx_behaviors.selectMainBehavior();
         await this.crawlState.addSourceWACZForDedupe(filename);
       }
     }
+    if (this.deduping) {
+      await this.crawlState.addUncommited();
+    }
 
     if (POST_CRAWL_STATES.includes(initState)) {
       logger.info("crawl already finished, running post-crawl tasks", {
@@ -1905,8 +1918,7 @@ self.__bx_behaviors.selectMainBehavior();
     }
 
     const generateFiles =
-      !this.params.dryRun &&
-      (!this.interruptReason || this.finalExit || this.uploadAndDeleteLocal);
+      !this.params.dryRun && (this.finalExit || this.uploadAndDeleteLocal);
 
     if (
       (this.params.generateCDX || this.params.generateWACZ) &&
@@ -1949,13 +1961,18 @@ self.__bx_behaviors.selectMainBehavior();
       }
     }
 
-    if (this.finalExit && generateFiles && this.deduping) {
+    // from here, actions that should happen on final crawler exit (not temp interrupt)
+    if (!this.finalExit) {
+      return;
+    }
+
+    if (this.deduping) {
       // commit crawl data to main index
       logger.info("Committing dedupe index");
       await this.crawlState.commitDedupeDone();
     }
 
-    if (this.finalExit && generateFiles && this.params.saveProfile) {
+    if (this.params.saveProfile && generateFiles) {
       const resource = await this.browser.saveProfile(
         this.params.saveProfile,
         this.storage,
@@ -1966,7 +1983,7 @@ self.__bx_behaviors.selectMainBehavior();
       }
     }
 
-    if (this.params.waitOnDone && (!this.interruptReason || this.finalExit)) {
+    if (this.params.waitOnDone) {
       this.done = true;
       logger.info("All done, waiting for signal...");
       await this.crawlState.setStatus("done");

src/util/constants.ts

@@ -25,6 +25,7 @@ export const EXTRACT_TEXT_TYPES = ["to-pages", "to-warc", "final-to-warc"];
 export const DUPE_ALL_HASH_KEY = "alldupes";
 export const DUPE_ALL_CRAWLS = "allcrawls";
 export const DUPE_ALL_COUNTS = "allcounts";
+export const DUPE_UNCOMMITTED = "uncommittedcrawls";
 
 export enum BxFunctionBindings {
   BehaviorLogFunc = "__bx_log",

src/util/state.ts

@@ -10,6 +10,7 @@ import {
   DUPE_ALL_HASH_KEY,
   DUPE_ALL_CRAWLS,
   DUPE_ALL_COUNTS,
+  DUPE_UNCOMMITTED,
 } from "./constants.js";
 import { ScopedSeed } from "./seeds.js";
 import { Frame } from "puppeteer-core";
@@ -316,12 +317,38 @@ export class RedisDedupeIndex {
     }
 
     // add to crawls list
+    await this.dedupeRedis.srem(DUPE_UNCOMMITTED, crawlId);
     await this.dedupeRedis.sadd(DUPE_ALL_CRAWLS, crawlId);
 
     // add counts
     await this.addCrawlCounts(crawlId);
   }
 
+  // ADD UNCOMITTED CRAWL
+  async addUncommited() {
+    if (this.crawlId) {
+      await this.dedupeRedis.sadd(DUPE_UNCOMMITTED, this.crawlId);
+    }
+  }
+
+  // CLEAR UNCOMMITTED
+  async clearUncommitted() {
+    if (this.crawlId) {
+      await this.dedupeRedis.srem(DUPE_UNCOMMITTED, this.crawlId);
+      await this.deleteCrawlDedupeKeys(this.crawlId);
+    }
+  }
+
+  async clearAllUncommitted() {
+    while (true) {
+      const crawlId = await this.dedupeRedis.spop(DUPE_UNCOMMITTED);
+      if (!crawlId) {
+        break;
+      }
+      await this.deleteCrawlDedupeKeys(crawlId);
+    }
+  }
+
   // GET OR ADD INDIVIDUAL HASHES
 
   async getHashDupe(hash: string): Promise<DedupeEntry | null> {
@@ -675,15 +702,19 @@ export class RedisDedupeIndex {
           // ignore
         }
       }
-      await this.dedupeRedis.del(
-        `h:${crawlId}`,
-        `c:${crawlId}:wacz`,
-        `h:${crawlId}:counts`,
-      );
+      await this.deleteCrawlDedupeKeys(crawlId);
     }
 
     await this.dedupeRedis.del(TO_REMOVE_CRAWLS);
   }
+
+  private async deleteCrawlDedupeKeys(crawlId: string) {
+    await this.dedupeRedis.del(
+      `h:${crawlId}`,
+      `c:${crawlId}:wacz`,
+      `h:${crawlId}:counts`,
+    );
+  }
 }
 
 // ============================================================================