From 39ab1e7683742022444e2926cd3df47327c6eb50 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Fri, 30 Apr 2021 19:27:21 +0530
Subject: [PATCH] Log error that causes `flush_db` to fail

---
 server/src/coredb/mod.rs | 18 ++++++++++++++++--
 server/src/dbnet/mod.rs  | 24 ++++++++++++------------
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/server/src/coredb/mod.rs b/server/src/coredb/mod.rs
index 2b0c153a..41149915 100644
--- a/server/src/coredb/mod.rs
+++ b/server/src/coredb/mod.rs
@@ -46,6 +46,16 @@ use std::sync::Arc;
 use tokio;
 use tokio::sync::Notify;
 
+#[macro_export]
+macro_rules! flush_db {
+    ($db:expr) => {
+        crate::coredb::CoreDB::flush_db(&$db, None)
+    };
+    ($db:expr, $file:expr) => {
+        crate::coredb::CoreDB::flush_db(&$db, Some($file))
+    };
+}
+
 /// This is a thread-safe database handle, which on cloning simply
 /// gives another atomic reference to the `shared` which is a `Shared` object
 #[derive(Debug, Clone)]
@@ -356,12 +366,16 @@ impl CoreDB {
         self.shared.table.read()
     }
     /// Flush the contents of the in-memory table onto disk
-    pub fn flush_db(&self) -> TResult<()> {
+    pub fn flush_db(&self, file: Option<flock::FileLock>) -> TResult<()> {
         let data = match self.acquire_write() {
             Some(wlock) => wlock,
             None => return Err("Can no longer flush data; coretable is poisoned".into()),
         };
-        diskstore::write_to_disk(&PERSIST_FILE, &data.coremap)?;
+        if let Some(mut file) = file {
+            diskstore::flush_data(&mut file, &data.coremap)?;
+        } else {
+            diskstore::write_to_disk(&PERSIST_FILE, &data.coremap)?;
+        }
         Ok(())
     }
 
diff --git a/server/src/dbnet/mod.rs b/server/src/dbnet/mod.rs
index 8d354328..745b3d64 100644
--- a/server/src/dbnet/mod.rs
+++ b/server/src/dbnet/mod.rs
@@ -62,6 +62,7 @@ use tokio::sync::Semaphore;
 use tokio::sync::{broadcast, mpsc};
 pub mod connection;
 mod tls;
+use crate::flush_db;
 
 /// Responsible for gracefully shutting down the server instead of dying randomly
 // Sounds very sci-fi ;)
@@ -390,29 +391,28 @@ pub async fn run(
         }
     }
     server.finish_with_termsig().await;
-    if let Ok(_) = db.flush_db() {
-        log::info!("Successfully saved data to disk");
-        ()
-    } else {
-        log::error!("Failed to flush data to disk");
+    if let Some(mut lock) = lock {
+        if let Err(e) = lock.unlock() {
+            log::error!("Failed to release lock on data file with '{}'", e);
+            process::exit(0x100);
+        }
+    }
+    if let Err(e) = flush_db!(db) {
+        log::error!("Failed to flush data to disk with '{}'", e);
         loop {
             // Keep looping until we successfully write the in-memory table to disk
             log::warn!("Press enter to try again...");
             io::stdout().flush().unwrap();
             io::stdin().read(&mut [0]).unwrap();
-            if let Ok(_) = db.flush_db() {
+            if let Ok(_) = flush_db!(db) {
                 log::info!("Successfully saved data to disk");
                 break;
             } else {
                 continue;
             }
         }
-    }
-    if let Some(mut lock) = lock {
-        if let Err(e) = lock.unlock() {
-            log::error!("Failed to release lock on data file with '{}'", e);
-            process::exit(0x100);
-        }
+    } else {
+        log::info!("Successfully saved data to disk");
     }
     terminal::write_info("Goodbye :)\n").unwrap();
 }