From 4e1e38a0fd7aa996ec0373de5de61ca784ebc193 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Mon, 12 Oct 2020 19:36:32 +0100
Subject: [PATCH] fix init bug

---
 Cargo.lock   |  2 +-
 Cargo.toml   |  3 +++
 src/main.rs  | 14 +++++++++++---
 src/msg.rs   | 19 +++++++------------
 src/state.rs |  4 +---
 5 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1784f20..0cc9ee9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -616,7 +616,7 @@ dependencies = [
 
 [[package]]
 name = "markov"
-version = "0.7.1"
+version = "0.7.2"
 dependencies = [
  "async-compression",
  "bzip2-sys",
diff --git a/Cargo.toml b/Cargo.toml
index c9b5940..f748059 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -43,6 +43,9 @@ hog-buffer = []
 # Enable the /api/ route
 api = []
 
+# Do not wait 2 seconds before starting worker tasks after server
+instant-init = []
+
 [profile.release]
 opt-level = 3
 lto = "fat"
diff --git a/src/main.rs b/src/main.rs
index 73f6f67..caf058a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -281,8 +281,15 @@ async fn main() {
 	    tokio::join![
 		server,
 		async move {
-		    trace!("Init set");
-		    init.set().expect("Failed to initialise saver")
+		    cfg_if! {
+			if #[cfg(feature="instant-init")] {
+			    trace!("Setting init");
+			} else {
+			    trace!("Setting init in 2 seconds for good measure.");
+			    tokio::time::delay_for(tokio::time::Duration::from_secs(2)).await;
+			}
+		    }
+		    init.set().expect("Failed to initialise saver")	
 		},
 	    ];
 	}).await;
@@ -290,7 +297,8 @@ async fn main() {
 	// Cleanup
 	async move {
 	    trace!("Cleanup");
-
+	    debug!("Waiting on {} tasks now", tasks.len());
+	    
 	    join_all(tasks).await;
 	}
     }.await;
diff --git a/src/msg.rs b/src/msg.rs
index a5bfc08..70eced5 100644
--- a/src/msg.rs
+++ b/src/msg.rs
@@ -68,9 +68,10 @@ impl Initialiser
     pub async fn wait(&mut self) -> Result<(), InitWaitError>
     {
 	if !*self.rx.borrow() {
-	    self.rx.recv().await
-		.ok_or_else(|| InitWaitError)
-		.and_then(|x| if x {Ok(())} else {Err(InitWaitError)})
+	    while !self.rx.recv().await.ok_or_else(|| InitWaitError)? {
+		tokio::task::yield_now().await;
+	    }
+	    Ok(())
 	} else {
 	    Ok(())
 	}
@@ -95,14 +96,8 @@ impl Future for Initialiser
 {
     type Output = Result<(), InitWaitError>;
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-	
-	if !*self.rx.borrow() {
-	    let rx = self.rx.recv();
-	    tokio::pin!(rx);
-	    rx.poll(cx).map(|x| x.ok_or_else(|| InitWaitError)
-			    .and_then(|x| if x {Ok(())} else {Err(InitWaitError)}))
-	} else {
-	    Poll::Ready(Ok(()))
-	}
+	let uhh = self.wait();
+	tokio::pin!(uhh);
+	uhh.poll(cx)
     }
 }
diff --git a/src/state.rs b/src/state.rs
index 4888dd4..2125ef3 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -58,9 +58,7 @@ impl State
     pub async fn on_init_save(&mut self) -> Result<(), ShutdownError>
     {
 	tokio::select!{
-	    Ok(()) = self.save_begin.wait() => {
-		Ok(())
-	    }
+	    Ok(()) = self.save_begin.wait() => Ok(()),
 	    Some(true) = self.shutdown_recv.recv() => {
 		debug!("on_init_save(): shutdown received");
 		Err(ShutdownError)