From b4ce27a1b973e74596f796c44c2302eb2aa028c8 Mon Sep 17 00:00:00 2001
From: mtkennerly <mtkennerly@gmail.com>
Date: Sat, 20 Apr 2024 20:50:10 -0400
Subject: [PATCH] Save progress when interrupted

---
 Cargo.lock   | 11 +++++++++++
 Cargo.toml   |  1 +
 src/main.rs  | 14 ++++++++++++++
 src/steam.rs |  5 +++++
 src/wiki.rs  | 10 +++++++++-
 5 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 512dd7f1..d8e62a67 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -928,6 +928,7 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_yaml",
+ "signal-hook",
  "thiserror",
  "tokio",
  "wikitext-parser",
@@ -1600,6 +1601,16 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "signal-hook"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801"
+dependencies = [
+ "libc",
+ "signal-hook-registry",
+]
+
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.1"
diff --git a/Cargo.toml b/Cargo.toml
index 1a0726e2..43e1cb7d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ regex = "1.10.2"
 serde = { version = "1.0.139", features = ["derive"] }
 serde_json = "1.0.108"
 serde_yaml = "0.8.25"
+signal-hook = "0.3.17"
 thiserror = "1.0.50"
 tokio = { version = "1.34.0", features = ["full"] }
 wikitext-parser = "0.3.3"
diff --git a/src/main.rs b/src/main.rs
index d546e4c4..6c45d419 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,13 @@ mod schema;
 mod steam;
 mod wiki;
 
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
+
+use once_cell::sync::Lazy;
+
 use crate::{
     manifest::{Manifest, ManifestOverride},
     resource::ResourceFile,
@@ -14,6 +21,11 @@ use crate::{
 };
 
 pub const REPO: &str = env!("CARGO_MANIFEST_DIR");
+static CANCEL: Lazy<Arc<AtomicBool>> = Lazy::new(|| Arc::new(AtomicBool::new(false)));
+
+pub fn should_cancel() -> bool {
+    CANCEL.load(Ordering::Relaxed)
+}
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -90,6 +102,8 @@ impl Error {
 async fn main() {
     let cli = cli::parse();
 
+    signal_hook::flag::register(signal_hook::consts::SIGINT, (*CANCEL).clone()).unwrap();
+
     let mut wiki_cache = WikiCache::load().unwrap();
     let mut wiki_meta_cache = WikiMetaCache::load().unwrap();
     let mut steam_cache = SteamCache::load().unwrap();
diff --git a/src/steam.rs b/src/steam.rs
index ccf4f70e..122bdff2 100644
--- a/src/steam.rs
+++ b/src/steam.rs
@@ -3,6 +3,7 @@ use std::{collections::BTreeMap, process::Command};
 use crate::{
     manifest::{placeholder, Os},
     resource::ResourceFile,
+    should_cancel,
     wiki::WikiCache,
     Error, State, REPO,
 };
@@ -36,6 +37,10 @@ impl SteamCache {
         });
 
         for app_id in app_ids {
+            if should_cancel() {
+                break;
+            }
+
             let latest = SteamCacheEntry::fetch_from_id(app_id)?;
             self.0.insert(
                 app_id,
diff --git a/src/wiki.rs b/src/wiki.rs
index eb68498a..0e9baf8d 100644
--- a/src/wiki.rs
+++ b/src/wiki.rs
@@ -7,7 +7,7 @@ use wikitext_parser::{Attribute, TextPiece};
 use crate::{
     manifest::{placeholder, Os, Store, Tag},
     resource::ResourceFile,
-    Error, Regularity, State,
+    should_cancel, Error, Regularity, State,
 };
 
 const SAVE_INTERVAL: u32 = 100;
@@ -212,6 +212,10 @@ impl WikiCache {
             .as_array()
             .ok_or(Error::WikiData("query.categorymembers"))?
         {
+            if should_cancel() {
+                break;
+            }
+
             let title = page["title"]
                 .as_str()
                 .ok_or(Error::WikiData("query.categorymembers[].title"))?;
@@ -277,6 +281,10 @@ impl WikiCache {
         });
 
         for title in &titles {
+            if should_cancel() {
+                break;
+            }
+
             let cached = self.0.get(title).cloned().unwrap_or_default();
 
             println!("Wiki: {}", title);