Merge branch 'master' into rotation

Author: AliceLanniste

Cargo.lock

@@ -9,7 +9,7 @@ crossbeam-channel = "0.4.0"
 crossbeam-utils = "0.7.0"
 log = "0.4.6"
 rand = "0.7.2"
-snap = "0.2.5"
+snap = "1.0.0"
 crc32c = "0.4.0"
 fs2 = "0.4.3"
 fxhash = "0.2.1"

Cargo.toml

@@ -91,6 +91,8 @@ pub trait DB {
 #[derive(Clone)]
 pub struct WickDB<S: Storage + Clone + 'static> {
     inner: Arc<DBImpl<S>>,
+    shutdown_batch_processing_thread: (Sender<()>, Receiver<()>),
+    shutdown_compaction_thread: (Sender<()>, Receiver<()>),
 }
 
 pub type WickDBIterator<S> = DBIterator<
@@ -152,15 +154,19 @@ impl<S: Storage + Clone> DB for WickDB<S> {
 
     fn close(&mut self) -> Result<()> {
         self.inner.is_shutting_down.store(true, Ordering::Release);
-        match &self.inner.db_lock {
-            Some(lock) => lock.unlock(),
-            None => Ok(()),
-        }
+        self.inner.schedule_close_batch();
+        let _ = self.shutdown_batch_processing_thread.1.recv();
+        // Send a signal to avoid blocking forever
+        let _ = self.inner.do_compaction.0.send(());
+        let _ = self.shutdown_compaction_thread.1.recv();
+        self.inner.close()?;
+        debug!("DB {} closed", &self.inner.db_name);
+        Ok(())
     }
 
     fn destroy(&mut self) -> Result<()> {
         let db = self.inner.clone();
-        db.is_shutting_down.store(true, Ordering::Release);
+        self.close()?;
         db.env.remove_dir(&db.db_name, true)
     }
 
@@ -173,6 +179,7 @@ impl<S: Storage + Clone> WickDB<S> {
     /// Create a new WickDB
     pub fn open_db(mut options: Options, db_name: &'static str, storage: S) -> Result<Self> {
         options.initialize(db_name.to_owned(), &storage);
+        debug!("Open db: '{}'", db_name);
         let mut db = DBImpl::new(options, db_name, storage);
         let (mut edit, should_save_manifest) = db.recover()?;
         let mut versions = db.versions.lock().unwrap();
@@ -188,16 +195,20 @@ impl<S: Storage + Clone> WickDB<S> {
         if should_save_manifest {
             edit.set_prev_log_number(0);
             edit.set_log_number(versions.log_number());
+            dbg!("log_and_apply in open_db");
             versions.log_and_apply(&mut edit)?;
         }
 
         let current = versions.current();
         db.delete_obsolete_files(versions);
         let wick_db = WickDB {
             inner: Arc::new(db),
+            shutdown_batch_processing_thread: crossbeam_channel::bounded(1),
+            shutdown_compaction_thread: crossbeam_channel::bounded(1),
         };
         wick_db.process_compaction();
         wick_db.process_batch();
+        // Schedule a compaction to current version for potential unfinished work
         wick_db.inner.maybe_schedule_compaction(current);
         Ok(wick_db)
     }
@@ -223,9 +234,17 @@ impl<S: Storage + Clone> WickDB<S> {
     // 5. Update sequence of version set
     fn process_batch(&self) {
         let db = self.inner.clone();
+        let shutdown = self.shutdown_batch_processing_thread.0.clone();
         thread::spawn(move || {
             loop {
                 if db.is_shutting_down.load(Ordering::Acquire) {
+                    // Cleanup all the batch queue
+                    let mut queue = db.batch_queue.lock().unwrap();
+                    while let Some(batch) = queue.pop_front() {
+                        let _ = batch.signal.send(Err(Error::DBClosed(
+                            "DB is closing. Clean up all the batch in queue".to_owned(),
+                        )));
+                    }
                     break;
                 }
                 let first = {
@@ -236,6 +255,9 @@ impl<S: Storage + Clone> WickDB<S> {
                     }
                     queue.pop_front().unwrap()
                 };
+                if first.stop_process {
+                    break;
+                }
                 let force = first.force_mem_compaction;
                 // TODO: The VersionSet is locked when processing `make_room_for_write`
                 match db.make_room_for_write(force) {
@@ -313,13 +335,16 @@ impl<S: Storage + Clone> WickDB<S> {
                     }
                 }
             }
+            shutdown.send(()).unwrap();
+            debug!("batch processing thread shut down");
         });
     }
 
     // Process a compaction work when receiving the signal.
     // The compaction might run recursively since we produce new table files.
     fn process_compaction(&self) {
         let db = self.inner.clone();
+        let shutdown = self.shutdown_compaction_thread.0.clone();
         thread::spawn(move || {
             while let Ok(()) = db.do_compaction.1.recv() {
                 if db.is_shutting_down.load(Ordering::Acquire) {
@@ -338,6 +363,8 @@ impl<S: Storage + Clone> WickDB<S> {
                 let current = db.versions.lock().unwrap().current();
                 db.maybe_schedule_compaction(current);
             }
+            shutdown.send(()).unwrap();
+            debug!("compaction thread shut down");
         });
     }
 }
@@ -389,9 +416,18 @@ unsafe impl<S: Storage + Clone> Send for DBImpl<S> {}
 impl<S: Storage + Clone> Drop for DBImpl<S> {
     #[allow(unused_must_use)]
     fn drop(&mut self) {
+        if !self.is_shutting_down.load(Ordering::Acquire) {
+            let _ = self.close();
+        }
+    }
+}
+
+impl<S: Storage + Clone> DBImpl<S> {
+    fn close(&self) -> Result<()> {
         self.is_shutting_down.store(true, Ordering::Release);
-        if let Some(lock) = self.db_lock.as_ref() {
-            lock.unlock();
+        match &self.db_lock {
+            Some(lock) => lock.unlock(),
+            None => Ok(()),
         }
     }
 }
@@ -492,13 +528,16 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
                 new_db.set_log_number(0);
                 new_db.set_next_file(2);
                 new_db.set_last_sequence(0);
+                // Create manifest
                 let manifest_filenum = 1;
                 let manifest_filename =
                     generate_filename(self.db_name, FileType::Manifest, manifest_filenum);
+                debug!("Create manifest file: {}", &manifest_filename);
                 let manifest = self.env.create(manifest_filename.as_str())?;
                 let mut manifest_writer = Writer::new(manifest);
                 let mut record = vec![];
                 new_db.encode_to(&mut record);
+                debug!("Append manifest record: {:?}", &new_db);
                 match manifest_writer.add_record(&record) {
                     Ok(()) => update_current(&self.env, self.db_name, manifest_filenum)?,
                     Err(e) => {
@@ -530,7 +569,7 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
         let prev_log = versions.prev_log_number();
         let all_files = self.env.list(self.db_name)?;
         let mut logs_to_recover = vec![];
-        for filename in all_files.iter() {
+        for filename in all_files {
             if let Some((file_type, file_number)) = parse_filename(filename) {
                 if file_type == FileType::Log && (file_number >= min_log || file_number == prev_log)
                 {
@@ -701,6 +740,20 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
         }
     }
 
+    // Schedule a WriteBatch to close batch processing thread for gracefully shutting down db
+    fn schedule_close_batch(&self) {
+        let (send, _) = crossbeam_channel::bounded(0);
+        let task = BatchTask {
+            stop_process: true,
+            force_mem_compaction: false,
+            batch: WriteBatch::default(),
+            signal: send,
+            options: WriteOptions::default(),
+        };
+        self.batch_queue.lock().unwrap().push_back(task);
+        self.process_batch_sem.notify_all();
+    }
+
     // Schedule the WriteBatch and wait for the result from the receiver.
     // This function wakes up the thread in `process_batch`.
     // An empty `WriteBatch` will trigger a force memtable compaction.
@@ -718,6 +771,7 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
         }
         let (send, recv) = crossbeam_channel::bounded(0);
         let task = BatchTask {
+            stop_process: false,
             force_mem_compaction,
             batch,
             signal: send,
@@ -747,7 +801,8 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
         // Group several batches from queue
         while !queue.is_empty() {
             let current = queue.pop_front().unwrap();
-            if current.options.sync && !grouped.options.sync {
+            if current.stop_process || (current.options.sync && !grouped.options.sync) {
+                // Do not include a stop process batch
                 // Do not include a sync write into a batch handled by a non-sync write.
                 queue.push_front(current);
                 break;
@@ -1291,6 +1346,7 @@ impl<S: Storage + Clone + 'static> DBImpl<S> {
 
 // A wrapper struct for scheduling `WriteBatch`
 struct BatchTask {
+    stop_process: bool, // flag for shutdown the batch processing thread gracefully
     force_mem_compaction: bool,
     batch: WriteBatch,
     signal: Sender<Result<()>>,
@@ -1418,7 +1474,7 @@ mod tests {
     }
 
     fn new_test_options(o: TestOption) -> Options {
-        match o {
+        let opt = match o {
             TestOption::Default => Options::default(),
             TestOption::Reuse => {
                 let mut o = Options::default();
@@ -1436,10 +1492,12 @@ mod tests {
                 o.compression = CompressionType::NoCompression;
                 o
             }
-        }
+        };
+        opt
     }
     struct DBTest {
-        store: MemStorage, // With the same db's inner storage
+        store: MemStorage, // Used as the db's inner storage
+        opt: Options,      // Used as the db's options
         db: WickDB<MemStorage>,
     }
 
@@ -1477,8 +1535,16 @@ mod tests {
         fn new(opt: Options) -> Self {
             let store = MemStorage::default();
             let name = "db_test";
-            let db = WickDB::open_db(opt, name, store.clone()).unwrap();
-            DBTest { store, db }
+            let db = WickDB::open_db(opt.clone(), name, store.clone()).unwrap();
+            DBTest { store, opt, db }
+        }
+
+        // Close the inner db without destroy the contents and establish a new WickDB on same db path with same option
+        fn reopen(&mut self) -> Result<()> {
+            self.db.close()?;
+            let db = WickDB::open_db(self.opt.clone(), self.db.inner.db_name, self.store.clone())?;
+            self.db = db;
+            Ok(())
         }
 
         // Put entries with default `WriteOptions`
@@ -1641,8 +1707,8 @@ mod tests {
             let store = MemStorage::default();
             let name = "db_test";
             let opt = new_test_options(TestOption::Default);
-            let db = WickDB::open_db(opt, name, store.clone()).unwrap();
-            DBTest { store, db }
+            let db = WickDB::open_db(opt.clone(), name, store.clone()).unwrap();
+            DBTest { store, opt, db }
         }
     }
 
@@ -2064,4 +2130,30 @@ mod tests {
             assert_iter_entry(&iter, "a", "va");
         }
     }
+
+    #[test]
+    fn test_reopen_with_empty_db() {
+        for mut t in default_cases() {
+            t.reopen().unwrap();
+            t.reopen().unwrap();
+        }
+    }
+
+    #[test]
+    fn test_recover_with_entries() {
+        for mut t in default_cases() {
+            t.put_entries(vec![("foo", "v1"), ("baz", "v5")]);
+            t.reopen().unwrap();
+            assert_eq!(t.get("foo", None).unwrap(), "v1");
+            assert_eq!(t.get("baz", None).unwrap(), "v5");
+
+            t.put_entries(vec![("bar", "v2"), ("foo", "v3")]);
+            t.reopen().unwrap();
+            assert_eq!(t.get("foo", None).unwrap(), "v3");
+            t.put("foo", "v4").unwrap();
+            assert_eq!(t.get("bar", None).unwrap(), "v2");
+            assert_eq!(t.get("foo", None).unwrap(), "v4");
+            assert_eq!(t.get("baz", None).unwrap(), "v5");
+        }
+    }
 }

src/db/mod.rs

@@ -46,21 +46,22 @@ impl Logger {
         let inner = match inner {
             Some(l) => l,
             None => {
-                let drain = if cfg!(debug_assertions) {
+                if cfg!(debug_assertions) {
                     // Use std out
                     let decorator = slog_term::TermDecorator::new().build();
-                    let drain = slog_term::FullFormat::new(decorator).build().fuse();
-                    slog_async::Async::new(drain).build().fuse()
+                    let drain = Mutex::new(slog_term::FullFormat::new(decorator).build()).fuse();
+                    slog::Logger::root(drain, o!())
                 } else {
                     // Use a file `Log` to record all logs
+                    // TODO: add file rotation
                     let file = storage
                         .create(generate_filename(db_path, FileType::InfoLog, 0).as_str())
                         .unwrap();
-                    slog_async::Async::new(FileBasedDrain::new(file))
+                    let drain = slog_async::Async::new(FileBasedDrain::new(file))
                         .build()
-                        .fuse()
-                };
-                slog::Logger::root(drain, o!())
+                        .fuse();
+                    slog::Logger::root(drain, o!())
+                }
             }
         };
         Self { inner, level }
@@ -95,7 +96,16 @@ impl Log for Logger {
                 level,
                 tag: target,
             };
-            self.inner.log(&slog::Record::new(&s, args, slog::b!()))
+            if cfg!(debug_assertions) {
+                let meta_info = format!("{}:{}", file, line);
+                self.inner.log(&slog::Record::new(
+                    &s,
+                    args,
+                    slog::b!("[location]" => meta_info),
+                ))
+            } else {
+                self.inner.log(&slog::Record::new(&s, args, slog::b!()))
+            }
         }
     }
     fn flush(&self) {}

src/logger.rs

@@ -44,6 +44,7 @@ impl From<u8> for CompressionType {
 }
 
 /// Options to control the behavior of a database (passed to `DB::Open`)
+#[derive(Clone)]
 pub struct Options {
     // -------------------
     // Parameters that affect behavior: