[PATCH proxmox-backup v2 06/14] datastore: implement index map for compact chunk usage tracking

[Christian Ebner <c.ebner@proxmox.com>]

In order to collect usage information on per-namespace/group/snapshot
level, a compact mapping from the chunks to the individual entities
referencing them needs to be established during phase 1 of garbage
collection in order to allow for reverse lookup and accounting during
phase 2, when the actual chunk raw size information is processed.
This needs to be done while keeping the memory footprint small.

IndexSet and IndexMap are therefore implemented, which store the value
or (key, value) in a compact vector, with fast item access by index,
while keeping a hashtable for fast lookup via key/value. Hash
collisions are taken into account by item comparison. The usable
index is kept generic to allow for different type checked index
implementations.

This follows along the implementation of hashed_index create,
implementing and exposing required functionality only.

Signed-off-by: Christian Ebner <c.ebner@proxmox.com>
---
 Cargo.toml                     |   1 +
 pbs-datastore/Cargo.toml       |   1 +
 pbs-datastore/src/index_map.rs | 101 +++++++++++++++++++++++++++++++++
 pbs-datastore/src/lib.rs       |   1 +
 4 files changed, 104 insertions(+)
 create mode 100644 pbs-datastore/src/index_map.rs

diff --git a/Cargo.toml b/Cargo.toml
index 6d82326c0..7d838f110 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -136,6 +136,7 @@ foreign-types = "0.3"
 futures = "0.3"
 h2 = { version = "0.4", features = [ "stream" ] }
 handlebars = "3.0"
+hashbrown = "0.14"
 hex = "0.4.3"
 hickory-resolver = { version = "0.24.1", default-features = false, features = [ "system-config", "tokio-runtime" ] }
 http-body-util = "0.1"
diff --git a/pbs-datastore/Cargo.toml b/pbs-datastore/Cargo.toml
index b51438dc6..c18e5b6b0 100644
--- a/pbs-datastore/Cargo.toml
+++ b/pbs-datastore/Cargo.toml
@@ -12,6 +12,7 @@ const_format.workspace = true
 crc32fast.workspace = true
 endian_trait.workspace = true
 futures.workspace = true
+hashbrown.workspace = true
 hex = { workspace = true, features = [ "serde" ] }
 http-body-util.workspace = true
 hyper.workspace = true
diff --git a/pbs-datastore/src/index_map.rs b/pbs-datastore/src/index_map.rs
new file mode 100644
index 000000000..f8affaf76
--- /dev/null
+++ b/pbs-datastore/src/index_map.rs
@@ -0,0 +1,101 @@
+use std::hash::{BuildHasher, Hash};
+
+use hashbrown::HashTable;
+use hashbrown::hash_map::DefaultHashBuilder;
+
+pub struct IndexSet<E, I = u32, H = DefaultHashBuilder> {
+    // Stores the hash of and entry and its index in the entries.
+    // Hash is not guaranteed to be collision free, so entries must be compared.
+    table: HashTable<I>,
+    entries: Vec<E>,
+    hasher: H,
+}
+
+impl<E, I> IndexSet<E, I> {
+    pub fn new() -> Self {
+        Self {
+            table: HashTable::<I>::new(),
+            entries: Vec::new(),
+            hasher: DefaultHashBuilder::default(),
+        }
+    }
+}
+
+impl<E: Hash + Eq, I: Copy + From<usize> + Into<usize>> IndexSet<E, I> {
+    pub fn insert(&mut self, entry: E) -> I {
+        let hash = self.hasher.hash_one(&entry);
+        if let Some(idx) = self
+            .table
+            .find(hash, |idx| self.entries[I::into(*idx)] == entry)
+        {
+            return *idx;
+        }
+        let idx = I::from(self.entries.len());
+        self.entries.push(entry);
+        self.table.insert_unique(hash, idx, |idx| {
+            self.hasher.hash_one(&self.entries[I::into(*idx)])
+        });
+        idx
+    }
+
+    pub fn lookup(&self, idx: I) -> Option<&E> {
+        self.entries.get(I::into(idx))
+    }
+}
+
+pub struct Entry<K, V> {
+    pub key: K,
+    pub value: V,
+}
+
+pub struct IndexMap<K, V, I = u32, H = DefaultHashBuilder> {
+    // Stores the hash of and entry and its index in the entries.
+    // Hash is not guaranteed to be collision free, so entry keys must be compared.
+    table: HashTable<I>,
+    entries: Vec<Entry<K, V>>,
+    hasher: H,
+}
+
+impl<K, V, I> IndexMap<K, V, I> {
+    pub fn new() -> Self {
+        Self {
+            table: HashTable::new(),
+            entries: Vec::new(),
+            hasher: DefaultHashBuilder::default(),
+        }
+    }
+}
+
+impl<K: Hash + Eq, V, I: Copy + From<usize> + Into<usize>> IndexMap<K, V, I> {
+    pub fn insert(&mut self, entry: Entry<K, V>) -> I {
+        let hash = self.hasher.hash_one(&entry.key);
+        if let Some(idx) = self.get_index_of(hash, &entry.key) {
+            self.entries[I::into(idx)].value = entry.value;
+            return idx;
+        }
+        let idx = I::from(self.entries.len());
+        self.entries.push(entry);
+        self.table.insert_unique(hash, idx, |idx| {
+            self.hasher.hash_one(&self.entries[I::into(*idx)].key)
+        });
+        idx
+    }
+
+    pub fn lookup(&self, idx: I) -> Option<&Entry<K, V>> {
+        self.entries.get(I::into(idx))
+    }
+
+    /// Get mutable reference to entry value with given key
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut Entry<K, V>> {
+        let hash = self.hasher.hash_one(key);
+        self.get_index_of(hash, key)
+            .and_then(|idx| self.entries.get_mut(I::into(idx)))
+    }
+
+    /// Find index for entry by its hash, comparig with key to detect hash collisions
+    fn get_index_of(&self, hash: u64, key: &K) -> Option<I> {
+        self.table
+            .find(hash, |idx| &self.entries[I::into(*idx)].key == key)
+            .copied()
+    }
+}
diff --git a/pbs-datastore/src/lib.rs b/pbs-datastore/src/lib.rs
index d9d7a4cb6..4e1222f9d 100644
--- a/pbs-datastore/src/lib.rs
+++ b/pbs-datastore/src/lib.rs
@@ -191,6 +191,7 @@ pub mod data_blob;
 pub mod data_blob_reader;
 pub mod file_formats;
 pub mod index;
+pub(crate) mod index_map;
 pub mod manifest;
 pub mod move_journal;
 pub mod paperkey;
-- 
2.47.3