[PATCH proxmox-backup v2 07/14] datastore: implement chunk map for collecting chunk usage statistics

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

The chunk map uses the previously introduced IndexSet and IndexMap to
keep track of chunk associations to snapshots, their groups and
namespaces during phase 1 of garbage collection and allow for
accounting after determining the raw chunk size for each referenced
chunk during phase 2.

Per entity indexes are defined for additional type checks, the
accumulated hierarchy statistics finally persisted to the datastore.

Signed-off-by: Christian Ebner <c.ebner@proxmox.com>
---
 pbs-datastore/src/chunk_map.rs | 315 +++++++++++++++++++++++++++++++++
 pbs-datastore/src/datastore.rs |  11 +-
 pbs-datastore/src/lib.rs       |   1 +
 3 files changed, 326 insertions(+), 1 deletion(-)
 create mode 100644 pbs-datastore/src/chunk_map.rs

diff --git a/pbs-datastore/src/chunk_map.rs b/pbs-datastore/src/chunk_map.rs
new file mode 100644
index 000000000..adff27628
--- /dev/null
+++ b/pbs-datastore/src/chunk_map.rs
@@ -0,0 +1,315 @@
+use std::collections::HashMap;
+use std::hash::Hash;
+
+use anyhow::{Context, Error};
+
+use pbs_api_types::{
+    BackupDir, BackupGroup, BackupNamespace, StatisticsListItem, StoreUsageStatistics,
+};
+
+use crate::DataStore;
+use crate::index_map::{Entry, IndexMap, IndexSet};
+
+// Limits index to u64::MAX chunks
+#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
+struct ChunkIndex(u64);
+
+impl From<ChunkIndex> for usize {
+    fn from(val: ChunkIndex) -> Self {
+        val.0 as usize
+    }
+}
+
+impl From<usize> for ChunkIndex {
+    fn from(id: usize) -> Self {
+        Self(id as u64)
+    }
+}
+
+#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
+struct SnapshotIndex(u32);
+
+impl From<SnapshotIndex> for usize {
+    fn from(val: SnapshotIndex) -> Self {
+        val.0 as usize
+    }
+}
+
+impl From<usize> for SnapshotIndex {
+    fn from(id: usize) -> Self {
+        Self(id as u32)
+    }
+}
+
+#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
+struct GroupIndex(u32);
+
+impl From<GroupIndex> for usize {
+    fn from(val: GroupIndex) -> Self {
+        val.0 as usize
+    }
+}
+
+impl From<usize> for GroupIndex {
+    fn from(id: usize) -> Self {
+        Self(id as u32)
+    }
+}
+
+#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd)]
+struct NamespaceIndex(u32);
+
+impl From<NamespaceIndex> for usize {
+    fn from(val: NamespaceIndex) -> Self {
+        val.0 as usize
+    }
+}
+
+impl From<usize> for NamespaceIndex {
+    fn from(id: usize) -> Self {
+        Self(id as u32)
+    }
+}
+
+#[derive(Eq, Hash, PartialEq)]
+struct BackupTime(i64);
+
+impl From<i64> for BackupTime {
+    fn from(time: i64) -> Self {
+        Self(time)
+    }
+}
+
+struct ChunkAssociation {
+    chunk_idx: ChunkIndex,
+    snapshot_idx: SnapshotIndex,
+    logical_size: u64,
+}
+
+pub struct ChunkMap<'a> {
+    /// Datastore associated with instance
+    datastore: &'a DataStore,
+    /// Compact list of digest and size
+    chunks: IndexMap<[u8; 32], u64, ChunkIndex>,
+    /// Compact list of namespaces
+    namespaces: IndexSet<BackupNamespace, NamespaceIndex>,
+    /// Compact list of groups
+    groups: IndexSet<(NamespaceIndex, BackupGroup), GroupIndex>,
+    /// Set keeping track of snapshots by index and their relation to group and namespace
+    snapshots: IndexSet<(GroupIndex, BackupTime), SnapshotIndex>,
+    /// Compact list keeping track of chunk <-> snapshot relation
+    associations: Vec<ChunkAssociation>,
+}
+
+impl<'a> ChunkMap<'a> {
+    /// Create new mapping instance for given datastore.
+    pub fn with_datastore(datastore: &'a DataStore) -> Self {
+        Self {
+            datastore,
+            chunks: IndexMap::new(),
+            namespaces: IndexSet::new(),
+            groups: IndexSet::new(),
+            snapshots: IndexSet::new(),
+            associations: Vec::new(),
+        }
+    }
+
+    /// Inserts the logical chunk information for an associated digest,
+    /// without accountig for the raw chunk size.
+    pub fn insert(
+        &mut self,
+        digest: &[u8; 32],
+        namespace: &BackupNamespace,
+        snapshot: &BackupDir,
+        logical_size: u64,
+    ) {
+        let chunk_entry = Entry {
+            key: *digest,
+            value: 0,
+        };
+        let chunk_idx = self.chunks.insert(chunk_entry);
+
+        let namespace_idx = self.namespaces.insert(namespace.clone());
+        let group_idx = self.groups.insert((namespace_idx, snapshot.group.clone()));
+        let snapshot_idx = self
+            .snapshots
+            .insert((group_idx, BackupTime::from(snapshot.time)));
+
+        self.associations.push(ChunkAssociation {
+            chunk_idx,
+            snapshot_idx,
+            logical_size,
+        });
+    }
+
+    /// Set the raw chunk size metadata information for an associated digest.
+    pub fn set_raw_chunk_size(&mut self, digest: &[u8; 32], raw_size: u64) {
+        if let Some(chunk_entry) = self.chunks.get_mut(digest) {
+            chunk_entry.value = raw_size;
+        } else {
+            let chunk_entry = Entry {
+                key: *digest,
+                value: raw_size,
+            };
+            let _idx = self.chunks.insert(chunk_entry);
+        }
+    }
+
+    /// Accumulate collected statistics on per-hierarchy level and store them
+    pub fn accumulate_and_persist(mut self) -> Result<(), Error> {
+        let mut datastore_stats = StoreUsageStatistics::default();
+        let mut namespace_stats: HashMap<NamespaceIndex, StoreUsageStatistics> = HashMap::new();
+        let mut group_stats: HashMap<GroupIndex, StoreUsageStatistics> = HashMap::new();
+        let mut snapshot_stats: HashMap<SnapshotIndex, StoreUsageStatistics> = HashMap::new();
+
+        // Buffer to handle individual chunk associations not only to
+        // snapshot index but also to its hierarchy
+        let mut association_buffer: Vec<(NamespaceIndex, GroupIndex, SnapshotIndex, u64)> =
+            Vec::new();
+
+        // Group associations by identical chunks in order to easily distinguish unique,
+        // duplicate and per-snapshot duplicate chunks.
+        self.associations
+            .sort_unstable_by_key(|association| (association.chunk_idx, association.snapshot_idx));
+        for association in self
+            .associations
+            .chunk_by(|a, b| a.chunk_idx == b.chunk_idx)
+        {
+            // All chunks within the same association are identical, can use first index
+            let chunk_idx = association[0].chunk_idx;
+            let raw_size = self
+                .chunks
+                .lookup(chunk_idx)
+                .map(|entry| entry.value)
+                .unwrap_or(0);
+
+            association_buffer.clear();
+            for record in association {
+                let (group_idx, _) = self
+                    .snapshots
+                    .lookup(record.snapshot_idx)
+                    .context("missing snapshot for index")?;
+                let (namespace_idx, _) = self
+                    .groups
+                    .lookup(*group_idx)
+                    .context("missing namespace for index")?;
+                association_buffer.push((
+                    *namespace_idx,
+                    *group_idx,
+                    record.snapshot_idx,
+                    record.logical_size,
+                ));
+            }
+
+            association_buffer.sort_unstable_by_key(
+                |&(namespace_idx, group_idx, snapshot_idx, _)| {
+                    (namespace_idx, group_idx, snapshot_idx)
+                },
+            );
+            let unique = association_buffer.len() == 1;
+            let mut datastore_sum = 0;
+
+            for namespace_associations in association_buffer.chunk_by(|a, b| a.0 == b.0) {
+                let namespace_idx = namespace_associations[0].0;
+                let mut namespace_sum = 0;
+
+                for group_associations in namespace_associations.chunk_by(|a, b| a.1 == b.1) {
+                    let group_idx = group_associations[0].1;
+                    let mut group_sum = 0;
+
+                    for snapshot_associations in group_associations.chunk_by(|a, b| a.2 == b.2) {
+                        let snapshot_idx = snapshot_associations[0].2;
+                        // Logical size must take into account if the same snapshot references the
+                        // same chunk multiple times, raw size is accounted once only.
+                        let snapshot_sum = snapshot_associations
+                            .iter()
+                            .map(|&(_, _, _, logical_size)| logical_size)
+                            .sum();
+                        group_sum += snapshot_sum;
+                        snapshot_stats.entry(snapshot_idx).or_default().add(
+                            raw_size,
+                            snapshot_sum,
+                            unique,
+                        );
+                    }
+
+                    namespace_sum += group_sum;
+                    group_stats
+                        .entry(group_idx)
+                        .or_default()
+                        .add(raw_size, group_sum, unique);
+                }
+
+                datastore_sum += namespace_sum;
+                namespace_stats.entry(namespace_idx).or_default().add(
+                    raw_size,
+                    namespace_sum,
+                    unique,
+                );
+            }
+
+            datastore_stats.add(raw_size, datastore_sum, unique);
+        }
+
+        let mut output = Vec::new();
+        for (snapshot_idx, stats) in snapshot_stats {
+            let (group_idx, snapshot_time) = self
+                .snapshots
+                .lookup(snapshot_idx)
+                .context("missing snapshot for index")?;
+            let (namespace_idx, group) = self
+                .groups
+                .lookup(*group_idx)
+                .context("missing group for index")?;
+            let namespace = self
+                .namespaces
+                .lookup(*namespace_idx)
+                .context("missing namespace for index")?;
+            output.push(StatisticsListItem {
+                ns: Some(namespace.clone()),
+                ty: Some(group.ty),
+                id: Some(group.id.clone()),
+                time: Some(snapshot_time.0),
+                stats,
+            });
+        }
+
+        for (group_idx, stats) in group_stats {
+            let (namespace_idx, group) = self
+                .groups
+                .lookup(group_idx)
+                .context("missing group for index")?;
+            let namespace = self
+                .namespaces
+                .lookup(*namespace_idx)
+                .context("missing namespace for index")?;
+            output.push(StatisticsListItem {
+                ns: Some(namespace.clone()),
+                ty: Some(group.ty),
+                id: Some(group.id.clone()),
+                time: None,
+                stats,
+            });
+        }
+
+        for (namespace_idx, stats) in namespace_stats {
+            let namespace = self
+                .namespaces
+                .lookup(namespace_idx)
+                .context("missing namespace for index")?;
+            output.push(StatisticsListItem {
+                ns: Some(namespace.clone()),
+                ty: None,
+                id: None,
+                time: None,
+                stats,
+            });
+        }
+
+        self.datastore
+            .write_statistics(&output)
+            .context("failed to write statistics")?;
+
+        Ok(())
+    }
+}
diff --git a/pbs-datastore/src/datastore.rs b/pbs-datastore/src/datastore.rs
index d6dcce77f..420221585 100644
--- a/pbs-datastore/src/datastore.rs
+++ b/pbs-datastore/src/datastore.rs
@@ -34,7 +34,7 @@ use pbs_api_types::{
     ArchiveType, Authid, BackupGroupDeleteStats, BackupNamespace, BackupType, ChunkOrder,
     DataStoreConfig, DatastoreBackendConfig, DatastoreBackendType, GarbageCollectionCacheStats,
     GarbageCollectionStatus, MAX_NAMESPACE_DEPTH, MaintenanceMode, MaintenanceType, Operation,
-    S3Statistics, UPID,
+    S3Statistics, StatisticsListItem, UPID,
 };
 use pbs_config::s3::S3_CFG_TYPE_ID;
 use pbs_config::{BackupLockGuard, ConfigVersionCache};
@@ -90,6 +90,8 @@ const S3_DELETE_BATCH_LIMIT: usize = 100;
 const S3_DELETE_DEFER_LIMIT_SECONDS: Duration = Duration::from_secs(60 * 5);
 // upper bound on concurrent snapshot locks during move operations to avoid exhausting FDs
 const MAX_SNAPSHOT_LOCKS: usize = 512;
+// Filename for the datastore statistics
+const DATASTORE_STATISTICS_FILENAME: &str = "statistics.json";
 
 /// Error from a locked per-group operation. Soft errors (e.g. lock conflicts)
 /// are retryable, hard errors are not.
@@ -3489,6 +3491,13 @@ impl DataStore {
 
         result
     }
+
+    pub(crate) fn write_statistics(&self, list: &[StatisticsListItem]) -> Result<(), Error> {
+        let path = self.base_path().join(DATASTORE_STATISTICS_FILENAME);
+        let contents = serde_json::to_vec(list)?;
+        replace_file(&path, &contents, CreateOptions::new(), true)?;
+        Ok(())
+    }
 }
 
 /// Track S3 object keys to be deleted by garbage collection while holding their file lock.
diff --git a/pbs-datastore/src/lib.rs b/pbs-datastore/src/lib.rs
index 4e1222f9d..8b98735d0 100644
--- a/pbs-datastore/src/lib.rs
+++ b/pbs-datastore/src/lib.rs
@@ -182,6 +182,7 @@ pub mod cached_chunk_reader;
 pub mod catalog;
 pub mod checksum_reader;
 pub mod checksum_writer;
+mod chunk_map;
 pub mod chunk_stat;
 pub mod chunk_store;
 pub mod chunker;
-- 
2.47.3