laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::{
    Ident,
    database::{
      Database,
      GenerationEpoch,
      storage::Partitions,
    },
  },
  serde::{
    Deserialize,
    Serialize,
  },
};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DbStats {
  pub total_index_entries: usize,
  pub current_epoch:       GenerationEpoch,
  pub partition_stats:     Vec<PartitionStats>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionStats {
  pub partition_key: String,
  pub entry_count:   usize,
}

pub fn collect_stats<P: Partitions>(db: &Database<P>) -> DbStats {
  let current_epoch = db.get_current_epoch();

  let mut partition_counts: std::collections::HashMap<Ident, usize> =
    std::collections::HashMap::new();
  let mut total_index_entries = 0;

  for partition_key in P::partition_keys() {
    let entries = db.index_range(*partition_key, "");
    let count = entries.len();
    total_index_entries += count;
    partition_counts.insert(*partition_key, count);
  }

  let partition_stats: Vec<PartitionStats> = partition_counts
    .into_iter()
    .map(|(partition_key, entry_count)| PartitionStats {
      partition_key: partition_key.to_string(),
      entry_count,
    })
    .collect();

  DbStats {
    total_index_entries,
    current_epoch,
    partition_stats,
  }
}