From b5b83b955ae18c5ab3b19fc13e1e300fa526e121 Mon Sep 17 00:00:00 2001
From: Aravinth Manivannan <realaravinth@batsense.net>
Date: Thu, 19 Oct 2023 10:10:33 +0530
Subject: [PATCH] fix: make archive shutdown responsive

---
 src/api/v1/mcaptcha/hooks.rs |  2 ++
 src/archive.rs               | 31 +++++++++++++++----------------
 src/mcaptcha.rs              | 14 +++-----------
 3 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/src/api/v1/mcaptcha/hooks.rs b/src/api/v1/mcaptcha/hooks.rs
index c8985eb..5474bce 100644
--- a/src/api/v1/mcaptcha/hooks.rs
+++ b/src/api/v1/mcaptcha/hooks.rs
@@ -147,6 +147,7 @@ async fn download(
 mod tests {
     use crate::api::v1::bench::Submission;
     use crate::api::v1::bench::SubmissionType;
+    use crate::api::v1::get_random;
     use crate::errors::*;
     use crate::mcaptcha::PerformanceAnalytics;
     use crate::mcaptcha::Secret;
@@ -181,6 +182,7 @@ mod tests {
 
         let payload = super::MCaptchaInstance {
             url: mcaptcha_instance.clone(),
+            auth_token: get_random(23),
         };
 
         let resp = test::call_service(
diff --git a/src/archive.rs b/src/archive.rs
index ad4b190..33c6103 100644
--- a/src/archive.rs
+++ b/src/archive.rs
@@ -17,7 +17,6 @@
 use std::future::Future;
 use std::path::{Path, PathBuf};
 
-
 use serde::{Deserialize, Serialize};
 use sqlx::types::time::OffsetDateTime;
 use sqlx::types::Uuid;
@@ -218,27 +217,27 @@ impl Archiver {
     ) -> ServiceResult<(Sender<bool>, JoinHandle<()>)> {
         let (tx, mut rx) = oneshot::channel();
 
+        fn can_run(rx: &mut oneshot::Receiver<bool>) -> bool {
+            match rx.try_recv() {
+                Err(TryRecvError::Empty) => true,
+                _ => false,
+            }
+        }
+
         let job = async move {
             loop {
-                //            let rx = self.rx.as_mut().unwrap();
-                match rx.try_recv() {
-                    // The channel is currently empty
-                    Ok(_) => {
+                if !can_run(&mut rx) {
+                    log::info!("Killing archive loop: received signal");
+                    break;
+                }
+
+                for _ in 0..data.settings.publish.duration {
+                    if !can_run(&mut rx) {
                         log::info!("Killing archive loop: received signal");
                         break;
                     }
-                    Err(TryRecvError::Empty) => {
-                        let _ = self.archive(&data).await;
-
-                        tokio::time::sleep(std::time::Duration::new(
-                            data.settings.publish.duration,
-                            0,
-                        ))
-                        .await;
-                    }
-                    Err(TryRecvError::Closed) => break,
+                    tokio::time::sleep(std::time::Duration::new(1, 0)).await;
                 }
-
                 let _ = self.archive(&data).await;
             }
         };
diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs
index 0135bb5..c57e695 100644
--- a/src/mcaptcha.rs
+++ b/src/mcaptcha.rs
@@ -111,17 +111,9 @@ impl MCaptchaClient for MCaptchaClientReqwest {
             auth_token: String,
         }
 
-        let msg = S {
-            secret,
-            auth_token,
-        };
+        let msg = S { secret, auth_token };
         mcaptcha.set_path("/api/v1/survey/secret");
-        self.client
-            .post(mcaptcha)
-            .json(&msg)
-            .send()
-            .await
-            .unwrap();
+        self.client.post(mcaptcha).json(&msg).send().await.unwrap();
         Ok(())
     }
     async fn download_benchmarks(
@@ -192,7 +184,7 @@ pub mod tests {
         ) -> ServiceResult<()> {
             mcaptcha.set_path("/api/v1/survey/secret");
             let mut x = self.client.write().unwrap();
-            x.insert(mcaptcha.to_string(), secret.secret.to_owned());
+            x.insert(mcaptcha.to_string(), secret);
             drop(x);
             Ok(())
         }