lb_rs/service/
usage.rs

1use crate::model::api::{FileUsage, GetUsageRequest};
2use crate::model::errors::LbResult;
3use crate::model::file_like::FileLike;
4use crate::model::tree_like::TreeLike;
5use crate::model::usage::{bytes_to_human, get_usage};
6use crate::Lb;
7use serde::Serialize;
8use std::collections::HashMap;
9use uuid::Uuid;
10
11#[derive(Serialize, Debug)]
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)]
19pub struct UsageItemMetric {
20    pub exact: u64,
21    pub readable: String,
22}
23
24impl Lb {
25    #[instrument(level = "debug", skip(self))]
26    pub async fn get_usage(&self) -> LbResult<UsageMetrics> {
27        let acc = self.get_account()?;
28        let usage = self.client.request(acc, GetUsageRequest {}).await?;
29        Ok(get_usage(usage))
30    }
31
32    #[instrument(level = "debug", skip(self))]
33    pub async fn get_uncompressed_usage_breakdown(&self) -> LbResult<HashMap<Uuid, usize>> {
34        let tx = self.ro_tx().await;
35        let db = tx.db();
36
37        let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
38        let mut breakdown = HashMap::default();
39
40        for id in tree.ids() {
41            let is_file_deleted = tree.calculate_deleted(&id)?;
42            let file = tree.find(&id)?;
43
44            if !is_file_deleted && file.is_document() {
45                let doc = self.read_document_helper(id, &mut tree).await?;
46                let doc_size = doc.len();
47                breakdown.insert(id, doc_size);
48            }
49        }
50
51        Ok(breakdown)
52    }
53
54    // big async opportunity
55    #[instrument(level = "debug", skip(self))]
56    pub async fn get_uncompressed_usage(&self) -> LbResult<UsageItemMetric> {
57        let tx = self.ro_tx().await;
58        let db = tx.db();
59
60        let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
61
62        let mut local_usage: u64 = 0;
63        for id in tree.ids() {
64            let is_file_deleted = tree.calculate_deleted(&id)?;
65            let file = tree.find(&id)?;
66
67            if !is_file_deleted && file.is_document() {
68                let doc = self.read_document_helper(id, &mut tree).await?;
69                local_usage += doc.len() as u64
70            }
71        }
72
73        let readable = bytes_to_human(local_usage);
74        Ok(UsageItemMetric { exact: local_usage, readable })
75    }
76}