lb_rs/service/
usage.rs

1use crate::Lb;
2use crate::model::api::{FileUsage, GetUsageRequest};
3use crate::model::errors::LbResult;
4use crate::model::file_like::FileLike;
5use crate::model::tree_like::TreeLike;
6use crate::model::usage::{bytes_to_human, get_usage};
7use serde::Serialize;
8use std::collections::HashMap;
9use uuid::Uuid;
10
11#[derive(Serialize, Debug, Clone)]
12pub struct UsageMetrics {
13    pub usages: Vec<FileUsage>,
14    pub server_usage: UsageItemMetric,
15    pub data_cap: UsageItemMetric,
16}
17
18#[derive(Serialize, PartialEq, Eq, Debug, Clone)]
19pub struct UsageItemMetric {
20    pub exact: u64,
21    pub readable: String,
22}
23
24impl Lb {
25    /// fetches data footprint on server along with data cap information
26    /// compares this to local changes to estimate net data increase
27    ///
28    /// callers of this function should be prepared to handle:
29    /// - [crate::LbErrKind::AccountNonexistent]
30    /// - [crate::LbErrKind::ClientUpdateRequired]
31    /// - [crate::LbErrKind::ServerUnreachable]
32    #[instrument(level = "debug", skip(self))]
33    pub async fn get_usage(&self) -> LbResult<UsageMetrics> {
34        let acc = self.get_account()?;
35        let usage = self.client.request(acc, GetUsageRequest {}).await?;
36        Ok(get_usage(usage))
37    }
38
39    #[instrument(level = "debug", skip(self))]
40    pub async fn get_uncompressed_usage_breakdown(&self) -> LbResult<HashMap<Uuid, usize>> {
41        let tx = self.ro_tx().await;
42        let db = tx.db();
43
44        let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
45        let mut breakdown = HashMap::default();
46
47        for id in tree.ids() {
48            let is_file_deleted = tree.calculate_deleted(&id)?;
49            let file = tree.find(&id)?;
50
51            if !is_file_deleted && file.is_document() {
52                let doc = self.read_document_helper(id, &mut tree).await?;
53                let doc_size = doc.len();
54                breakdown.insert(id, doc_size);
55            }
56        }
57
58        Ok(breakdown)
59    }
60
61    // big async opportunity
62    #[instrument(level = "debug", skip(self))]
63    pub async fn get_uncompressed_usage(&self) -> LbResult<UsageItemMetric> {
64        let tx = self.ro_tx().await;
65        let db = tx.db();
66
67        let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
68
69        let mut local_usage: u64 = 0;
70        for id in tree.ids() {
71            let is_file_deleted = tree.calculate_deleted(&id)?;
72            let file = tree.find(&id)?;
73
74            if !is_file_deleted && file.is_document() {
75                let doc = self.read_document_helper(id, &mut tree).await?;
76                local_usage += doc.len() as u64
77            }
78        }
79
80        let readable = bytes_to_human(local_usage);
81        Ok(UsageItemMetric { exact: local_usage, readable })
82    }
83}