From d62da60df53ec37a4d73eaafbcb626103c2f42e4 Mon Sep 17 00:00:00 2001
From: Nan Jiang <najiang@mozilla.com>
Date: Thu, 4 Apr 2019 10:41:33 -0400
Subject: [PATCH] Add freelist to environment (#41)

---
 Cargo.toml         |  2 +-
 src/database.rs    |  6 ++++
 src/environment.rs | 77 +++++++++++++++++++++++++++++++++++++++++++++-
 src/lib.rs         |  5 ++-
 4 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c54694c..df6311e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -35,10 +35,10 @@ members = [
 
 [dependencies]
 bitflags = "1"
+byteorder = "1.0"
 libc = "0.2"
 lmdb-rkv-sys = "0.8.2"
 
 [dev-dependencies]
 rand = "0.4"
 tempdir = "0.3"
-byteorder = "1.0"
diff --git a/src/database.rs b/src/database.rs
index 61656dc..e7973ba 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -31,6 +31,12 @@ impl Database {
         Ok(Database { dbi: dbi })
     }
 
+    pub(crate) fn freelist_db() -> Database {
+        Database {
+            dbi: 0,
+        }
+    }
+
     /// Returns the underlying LMDB database handle.
     ///
     /// The caller **must** ensure that the handle is not used after the lifetime of the
diff --git a/src/environment.rs b/src/environment.rs
index ffb39d7..dc0f96a 100644
--- a/src/environment.rs
+++ b/src/environment.rs
@@ -10,7 +10,10 @@ use std::sync::Mutex;
 
 use ffi;
 
-use error::{Result, lmdb_result};
+use byteorder::{ByteOrder, NativeEndian};
+
+use cursor::Cursor;
+use error::{Error, Result, lmdb_result};
 use database::Database;
 use transaction::{RoTransaction, RwTransaction, Transaction};
 use flags::{DatabaseFlags, EnvironmentFlags};
@@ -173,6 +176,52 @@ impl Environment {
         }
     }
 
+    /// Retrieves the total number of pages on the freelist.
+    ///
+    /// Along with `Environment::info()`, this can be used to calculate the exact number
+    /// of used pages as well as free pages in this environment.
+    ///
+    /// ```ignore
+    /// let env = Environment::new().open("/tmp/test").unwrap();
+    /// let info = env.info().unwrap();
+    /// let stat = env.stat().unwrap();
+    /// let freelist = env.freelist().unwrap();
+    /// let last_pgno = info.last_pgno() + 1; // pgno is 0 based.
+    /// let total_pgs = info.map_size() / stat.page_size() as usize;
+    /// let pgs_in_use = last_pgno - freelist;
+    /// let pgs_free = total_pgs - pgs_in_use;
+    /// ```
+    ///
+    /// Note:
+    ///
+    /// * LMDB stores all the freelists in the designated database 0 in each environment,
+    ///   and the freelist count is stored at the beginning of the value as `libc::size_t`
+    ///   in the native byte order.
+    ///
+    /// * It will create a read transaction to traverse the freelist database.
+    pub fn freelist(&self) -> Result<size_t> {
+        let mut freelist: size_t = 0;
+        let db = Database::freelist_db();
+        let txn = self.begin_ro_txn()?;
+        let mut cursor = txn.open_ro_cursor(db)?;
+
+        for result in cursor.iter() {
+            let (_key, value) = result?;
+            if value.len() < mem::size_of::<size_t>() {
+                return Err(Error::Corrupted);
+            }
+
+            let s = &value[..mem::size_of::<size_t>()];
+            if cfg!(target_pointer_width = "64") {
+                freelist += NativeEndian::read_u64(s) as size_t;
+            } else {
+                freelist += NativeEndian::read_u32(s) as size_t;
+            }
+        }
+
+        Ok(freelist)
+    }
+
     /// Sets the size of the memory map to use for the environment.
     ///
     /// This could be used to resize the map when the environment is already open.
@@ -549,6 +598,32 @@ mod test {
         assert_eq!(info.num_readers(), 0);
     }
 
+    #[test]
+    fn test_freelist() {
+        let dir = TempDir::new("test").unwrap();
+        let env = Environment::new().open(dir.path()).unwrap();
+
+        let db = env.open_db(None).unwrap();
+        let mut freelist = env.freelist().unwrap();
+        assert_eq!(freelist, 0);
+
+        // Write a few small values.
+        for i in 0..64 {
+            let mut value = [0u8; 8];
+            LittleEndian::write_u64(&mut value, i);
+            let mut tx = env.begin_rw_txn().expect("begin_rw_txn");
+            tx.put(db, &value, &value, WriteFlags::default()).expect("tx.put");
+            tx.commit().expect("tx.commit")
+        }
+        let mut tx = env.begin_rw_txn().expect("begin_rw_txn");
+        tx.clear_db(db).expect("clear");
+        tx.commit().expect("tx.commit");
+
+        // Freelist should not be empty after clear_db.
+        freelist = env.freelist().unwrap();
+        assert!(freelist > 0);
+    }
+
     #[test]
     fn test_set_map_size() {
         let dir = TempDir::new("test").unwrap();
diff --git a/src/lib.rs b/src/lib.rs
index 5e09823..127e422 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,6 +7,7 @@
 
 extern crate libc;
 extern crate lmdb_rkv_sys as ffi;
+extern crate byteorder;
 
 #[cfg(test)] extern crate rand;
 #[cfg(test)] extern crate tempdir;
@@ -67,9 +68,7 @@ mod transaction;
 #[cfg(test)]
 mod test_utils {
 
-    extern crate byteorder;
-
-    use self::byteorder::{ByteOrder, LittleEndian};
+    use byteorder::{ByteOrder, LittleEndian};
     use tempdir::TempDir;
 
     use super::*;