reifydb_store_transaction/stats/
query.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4//! Query methods for storage statistics.
5
6use std::collections::HashMap;
7
8use reifydb_core::key::KeyKind;
9
10use super::{
11	tracker::{StorageTracker, StorageTrackerConfig},
12	types::{ObjectId, StorageStats, Tier, TierStats},
13};
14
15impl StorageTracker {
16	/// Get total storage stats across all tiers.
17	pub fn total_stats(&self) -> TierStats {
18		let inner = self.inner.read().unwrap();
19		let mut result = TierStats::new();
20
21		// Use by_tier which tracks all keys, including those without recognized KeyKind
22		for (tier, stats) in &inner.by_tier {
23			*result.get_mut(*tier) = stats.clone();
24		}
25
26		result
27	}
28
29	/// Get stats aggregated by KeyKind for a specific tier.
30	pub fn stats_by_type(&self, tier: Tier) -> HashMap<KeyKind, StorageStats> {
31		let inner = self.inner.read().unwrap();
32		let mut result = HashMap::new();
33
34		for ((t, kind), stats) in &inner.by_type {
35			if *t == tier {
36				result.insert(*kind, stats.clone());
37			}
38		}
39
40		result
41	}
42
43	/// Get stats aggregated by KeyKind across all tiers.
44	pub fn stats_by_type_all_tiers(&self) -> HashMap<KeyKind, TierStats> {
45		let inner = self.inner.read().unwrap();
46		let mut result: HashMap<KeyKind, TierStats> = HashMap::new();
47
48		for ((tier, kind), stats) in &inner.by_type {
49			let tier_stats = result.entry(*kind).or_insert_with(TierStats::new);
50			*tier_stats.get_mut(*tier) += stats.clone();
51		}
52
53		result
54	}
55
56	/// Get stats for a specific object across all tiers.
57	pub fn stats_for_object(&self, object_id: ObjectId) -> Option<TierStats> {
58		let inner = self.inner.read().unwrap();
59
60		let mut result = TierStats::new();
61		let mut found = false;
62
63		for tier in [Tier::Hot, Tier::Warm, Tier::Cold] {
64			if let Some(stats) = inner.by_object.get(&(tier, object_id)) {
65				*result.get_mut(tier) = stats.clone();
66				found = true;
67			}
68		}
69
70		if found {
71			Some(result)
72		} else {
73			None
74		}
75	}
76
77	/// Get all objects for a specific tier, sorted by total bytes descending.
78	pub fn objects_by_tier(&self, tier: Tier) -> Vec<(ObjectId, StorageStats)> {
79		let inner = self.inner.read().unwrap();
80		let mut result: Vec<_> = inner
81			.by_object
82			.iter()
83			.filter(|((t, _), _)| *t == tier)
84			.map(|((_, obj_id), stats)| (*obj_id, stats.clone()))
85			.collect();
86
87		result.sort_by(|(_, a), (_, b)| b.total_bytes().cmp(&a.total_bytes()));
88		result
89	}
90
91	/// Get top N objects by total storage consumption across all tiers.
92	pub fn top_objects_by_size(&self, n: usize) -> Vec<(ObjectId, TierStats)> {
93		let inner = self.inner.read().unwrap();
94
95		// Aggregate stats by object across all tiers
96		let mut by_object: HashMap<ObjectId, TierStats> = HashMap::new();
97		for ((tier, obj_id), stats) in &inner.by_object {
98			let tier_stats = by_object.entry(*obj_id).or_insert_with(TierStats::new);
99			*tier_stats.get_mut(*tier) = stats.clone();
100		}
101
102		// Sort by total bytes
103		let mut result: Vec<_> = by_object.into_iter().collect();
104		result.sort_by(|(_, a), (_, b)| b.total_bytes().cmp(&a.total_bytes()));
105		result.truncate(n);
106		result
107	}
108
109	/// Get configuration.
110	pub fn config(&self) -> StorageTrackerConfig {
111		self.inner.read().unwrap().config.clone()
112	}
113}