Skip to main content

shardline_server/postgres_backend/
stats.rs

1use shardline_index::RecordTraversal;
2use shardline_storage::ObjectPrefix;
3
4use crate::{
5    ServerError,
6    chunk_store::chunk_hash_from_chunk_object_key_if_present,
7    model::ServerStatsResponse,
8    object_store::visit_object_prefix,
9    overflow::{checked_add, checked_increment},
10};
11
12impl super::PostgresBackend {
13    /// Returns backend storage stats.
14    ///
15    /// # Errors
16    ///
17    /// Returns [`ServerError`] when metadata inventory cannot be loaded.
18    pub async fn stats(&self) -> Result<ServerStatsResponse, ServerError> {
19        let object_store = self.object_store();
20        let prefix = ObjectPrefix::parse("").map_err(|_error| ServerError::InvalidContentHash)?;
21        let mut chunks = 0_u64;
22        let mut chunk_bytes = 0_u64;
23        visit_object_prefix(&object_store, &prefix, |metadata| {
24            let is_chunk = chunk_hash_from_chunk_object_key_if_present(metadata.key())?.is_some();
25            if is_chunk {
26                chunks = checked_increment(chunks)?;
27                chunk_bytes = checked_add(chunk_bytes, metadata.length())?;
28            }
29
30            Ok(())
31        })?;
32        let mut files = 0_u64;
33        RecordTraversal::visit_latest_record_locators(&self.record_store, |_locator| {
34            files = checked_increment(files)?;
35            Ok::<(), ServerError>(())
36        })
37        .await?;
38
39        Ok(ServerStatsResponse {
40            chunks,
41            chunk_bytes,
42            files,
43        })
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use std::num::NonZeroUsize;
50
51    use super::*;
52    use crate::object_store::ServerObjectStore;
53
54    const TEST_PG_URL: &str = "postgres://localhost:5432/test";
55
56    async fn make_backend() -> (super::super::PostgresBackend, tempfile::TempDir) {
57        let root = tempfile::tempdir().expect("temp dir");
58        let object_store =
59            ServerObjectStore::local(root.path().join("chunks")).expect("local store");
60        let backend = super::super::PostgresBackend::new_with_object_store_and_upload_parallelism(
61            root.path().to_path_buf(),
62            "http://127.0.0.1:8080".to_owned(),
63            NonZeroUsize::new(65536).unwrap(),
64            NonZeroUsize::new(64).unwrap(),
65            TEST_PG_URL,
66            object_store,
67        )
68        .await
69        .expect("constructor");
70        (backend, root)
71    }
72
73    #[tokio::test]
74    async fn stats_runs_without_panic_when_empty() {
75        let (backend, _root) = make_backend().await;
76        // stats() calls RecordTraversal which will fail without a real PG DB.
77        // The object_store part works fine; the record_store part may error.
78        // We only verify that the function runs without panicking.
79        let _response = backend.stats().await;
80    }
81
82    #[test]
83    fn server_stats_response_default_fields() {
84        let response = ServerStatsResponse {
85            chunks: 0,
86            chunk_bytes: 0,
87            files: 0,
88        };
89        assert_eq!(response.chunks, 0);
90        assert_eq!(response.chunk_bytes, 0);
91        assert_eq!(response.files, 0);
92    }
93
94    #[test]
95    fn server_stats_response_arbitrary_values() {
96        let response = ServerStatsResponse {
97            chunks: 42,
98            chunk_bytes: 1_000_000,
99            files: 7,
100        };
101        assert_eq!(response.chunks, 42);
102        assert_eq!(response.chunk_bytes, 1_000_000);
103        assert_eq!(response.files, 7);
104    }
105}