From aefe6254c821c7ebc5d4d24cb024337261acfb56 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Wed, 17 Apr 2024 07:31:55 +0000
Subject: [PATCH] Do not flush database after each insert

---
 src/notifier.rs |  2 --
 src/schedule.rs | 19 ++++++++++++-------
 src/server.rs   |  5 ++++-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/notifier.rs b/src/notifier.rs
index 00b0591..bc6c841 100644
--- a/src/notifier.rs
+++ b/src/notifier.rs
@@ -114,7 +114,6 @@ async fn wakeup(
                 info!("delivered notification for {}", device_token);
                 schedule
                     .insert_token_now(&key_device_token)
-                    .await
                     .context("Failed to update latest notification timestamp")?;
                 metrics.heartbeat_notifications_total.inc();
             }
@@ -136,7 +135,6 @@ async fn wakeup(
             // to avoid busy looping.
             schedule
                 .insert_token_now(&key_device_token)
-                .await
                 .with_context(|| format!("Failed to update token timestamp: {err:?}"))?;
         }
     }
diff --git a/src/schedule.rs b/src/schedule.rs
index bac9fb3..5fb1208 100644
--- a/src/schedule.rs
+++ b/src/schedule.rs
@@ -40,20 +40,24 @@ impl Schedule {
     ///
     /// This should also be called after successful notification
     /// to update latest notification time.
-    pub async fn insert_token(&self, token: &str, now: u64) -> Result<()> {
+    pub fn insert_token(&self, token: &str, now: u64) -> Result<()> {
         self.db.insert(token.as_bytes(), &u64::to_be_bytes(now))?;
-        self.db.flush_async().await?;
         let mut heap = self.heap.lock().unwrap();
         heap.push((Reverse(now), token.to_owned()));
         Ok(())
     }
 
-    pub async fn insert_token_now(&self, token: &str) -> Result<()> {
+    pub fn insert_token_now(&self, token: &str) -> Result<()> {
         let now = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
             .unwrap_or_default()
             .as_secs();
-        self.insert_token(token, now).await
+        self.insert_token(token, now)
+    }
+
+    pub async fn flush(&self) -> Result<()> {
+        self.db.flush_async().await?;
+        Ok(())
     }
 
     /// Removes token from the schedule.
@@ -87,13 +91,14 @@ mod tests {
         let db_path = dir.path().join("db.sled");
         let schedule = Schedule::new(&db_path)?;
 
-        schedule.insert_token("foo", 10).await?;
-        schedule.insert_token("bar", 20).await?;
+        schedule.insert_token("foo", 10)?;
+        schedule.insert_token("bar", 20)?;
 
         let (first_timestamp, first_token) = schedule.pop().unwrap();
         assert_eq!(first_timestamp, 10);
         assert_eq!(first_token, "foo");
-        schedule.insert_token("foo", 30).await?;
+        schedule.insert_token("foo", 30)?;
+        schedule.flush().await?;
 
         // Reopen to test persistence.
         drop(schedule);
diff --git a/src/server.rs b/src/server.rs
index da4a17e..5f2a9e6 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -30,7 +30,10 @@ async fn register_device(mut req: tide::Request<State>) -> tide::Result<tide::Re
     info!("register_device {}", query.token);
 
     let schedule = req.state().schedule();
-    schedule.insert_token_now(&query.token).await?;
+    schedule.insert_token_now(&query.token)?;
+
+    // Flush database to ensure we don't lose this token in case of restart.
+    schedule.flush().await?;
 
     req.state().metrics().heartbeat_registrations_total.inc();