From a330f0a4ef8946b0ddab296a923479fe64633773 Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Tue, 8 Mar 2022 14:49:41 +0100
Subject: [PATCH] Add api to get db size to content db (#987)

* Add size proc to content_db

* Add way to reclaim used space
---
 fluffy/content_db.nim            | 33 ++++++++++++++++++++++++-
 fluffy/tests/test_content_db.nim | 42 ++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/fluffy/content_db.nim b/fluffy/content_db.nim
index 3fc2c4d10..1de7e1900 100644
--- a/fluffy/content_db.nim
+++ b/fluffy/content_db.nim
@@ -32,6 +32,8 @@ export kvstore_sqlite3
 type
   ContentDB* = ref object
     kv: KvStoreRef
+    sizeStmt: SqliteStmt[NoParams, int64]
+    vacStmt: SqliteStmt[NoParams, void]
 
 template expectDb(x: auto): untyped =
   # There's no meaningful error handling implemented for a corrupt database or
@@ -46,7 +48,36 @@ proc new*(T: type ContentDB, path: string, inMemory = false): ContentDB =
     else:
       SqStoreRef.init(path, "fluffy").expectDb()
 
-  ContentDB(kv: kvStore db.openKvStore().expectDb())
+  let getSizeStmt = db.prepareStmt(
+    "SELECT page_count * page_size as size FROM pragma_page_count(), pragma_page_size();",
+    NoParams, int64).get()
+
+  let vacStmt = db.prepareStmt(
+    "VACUUM;",
+    NoParams, void).get()
+
+  ContentDB(kv: kvStore db.openKvStore().expectDb(), sizeStmt: getSizeStmt, vacStmt: vacStmt)
+
+proc reclaimSpace*(db: ContentDB): void =
+  ## Runs sqlie VACUMM commands which rebuilds db, repacking it into a minimal amount of disk space
+  ## Ideal mode of operation, is to run it after several deletes.
+  ## Another options would be to run 'PRAGMA auto_vacuum = FULL;' statement at the start of 
+  ## db to leave it in sqlite power to clean up
+  db.vacStmt.exec().expectDb()
+
+proc size*(db: ContentDB): int64 =
+  ## Retrun current size of DB as product of sqlite page_count and page_size
+  ## https://www.sqlite.org/pragma.html#pragma_page_count
+  ## https://www.sqlite.org/pragma.html#pragma_page_size
+  ## It returns total size of db i.e both data and metadata used to store content
+  ## also it is worth noting that when deleting content, size may lags behind due
+  ## to the way how deleting works in sqlite.
+  ## Good description can be found in: https://www.sqlite.org/lang_vacuum.html
+
+  var size: int64 = 0
+  discard (db.sizeStmt.exec do(res: int64):
+    size = res).expectDb()
+  return size
 
 proc get*(db: ContentDB, key: openArray[byte]): Option[seq[byte]] =
   var res: Option[seq[byte]]
diff --git a/fluffy/tests/test_content_db.nim b/fluffy/tests/test_content_db.nim
index 33c8965c4..423bc48d7 100644
--- a/fluffy/tests/test_content_db.nim
+++ b/fluffy/tests/test_content_db.nim
@@ -12,6 +12,14 @@ import
   ../network/state/state_content,
   ../content_db
 
+proc genByteSeq(length: int): seq[byte] = 
+  var i = 0
+  var resultSeq = newSeq[byte](length)
+  while i < length:
+    resultSeq[i] = byte(i)
+    inc i
+  return resultSeq
+
 suite "Content Database":
   # Note: We are currently not really testing something new here just basic
   # underlying kvstore.
@@ -43,3 +51,37 @@ suite "Content Database":
       check:
         val.isNone()
         db.contains(key) == false
+
+  test "ContentDB size":
+    let
+      db = ContentDB.new("", inMemory = true)
+
+    let numBytes = 10000
+    let size1 = db.size()
+    db.put(@[1'u8], genByteSeq(numBytes))
+    let size2 = db.size()
+    db.put(@[2'u8], genByteSeq(numBytes))
+    let size3 = db.size()
+    db.put(@[2'u8], genByteSeq(numBytes))
+    let size4 = db.size()
+
+    check:
+      size2 > size1
+      size3 > size2
+      size3 == size4
+
+    db.del(@[2'u8])
+    db.del(@[1'u8])
+    
+    let size5 = db.size()
+    
+    check:
+      size4 == size5
+
+    db.reclaimSpace()
+
+    let size6 = db.size()
+
+    check:
+      # After space reclamation size of db should be equal to initial size
+      size6 == size1