diff --git a/Cargo.lock b/Cargo.lock index cdee14874631..70edc68b5347 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1312,6 +1312,12 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.28" @@ -4913,6 +4919,7 @@ dependencies = [ "base64 0.22.0", "bincode", "chrono", + "fs_extra", "futures", "hash-ids", "hyper 0.14.27", diff --git a/ee/tabby-webserver/Cargo.toml b/ee/tabby-webserver/Cargo.toml index 6eaa1a361070..842d4e3dad75 100644 --- a/ee/tabby-webserver/Cargo.toml +++ b/ee/tabby-webserver/Cargo.toml @@ -51,10 +51,12 @@ validator = { version = "0.16.1", features = ["derive"] } regex.workspace = true tabby-search = { path = "../tabby-search" } octocrab = "0.38.0" +fs_extra = "1.3.0" [dev-dependencies] assert_matches = "1.5.0" tokio = { workspace = true, features = ["macros"] } tabby-db = { path = "../../ee/tabby-db", features = ["testutils"] } +tabby-common = { path = "../../crates/tabby-common", features = [ "testutils" ] } serial_test = { workspace = true } temp_testdir = { workspace = true } diff --git a/ee/tabby-webserver/graphql/schema.graphql b/ee/tabby-webserver/graphql/schema.graphql index cc1bf57ccd3e..c52e2e6a3bc0 100644 --- a/ee/tabby-webserver/graphql/schema.graphql +++ b/ee/tabby-webserver/graphql/schema.graphql @@ -91,6 +91,13 @@ type UserConnection { pageInfo: PageInfo! } +type DiskUsageStats { + events: DiskUsage! + indexedRepositories: DiskUsage! + database: DiskUsage! + models: DiskUsage! +} + input UpdateOAuthCredentialInput { provider: OAuthProvider! clientId: String! @@ -258,6 +265,7 @@ type Query { dailyStatsInPastYear(users: [ID!]): [CompletionStats!]! dailyStats(start: DateTimeUtc!, end: DateTimeUtc!, users: [ID!], languages: [Language!]): [CompletionStats!]! userEvents(after: String, before: String, first: Int, last: Int, users: [ID!], start: DateTimeUtc!, end: DateTimeUtc!): UserEventConnection! + diskUsageStats: DiskUsageStats! } input NetworkSettingInput { @@ -275,6 +283,11 @@ type UserEdge { cursor: String! } +type DiskUsage { + filePaths: [String!]! + size: Float! +} + type JobRunConnection { edges: [JobRunEdge!]! pageInfo: PageInfo! diff --git a/ee/tabby-webserver/src/path.rs b/ee/tabby-webserver/src/path.rs index cbc9da76254b..708fda4692e1 100644 --- a/ee/tabby-webserver/src/path.rs +++ b/ee/tabby-webserver/src/path.rs @@ -2,7 +2,7 @@ use std::path::PathBuf; use tabby_common::path::tabby_root; -fn tabby_ee_root() -> PathBuf { +pub fn tabby_ee_root() -> PathBuf { tabby_root().join("ee") } diff --git a/ee/tabby-webserver/src/schema/analytic.rs b/ee/tabby-webserver/src/schema/analytic.rs index 0b17f4ec9f88..72cf903df09a 100644 --- a/ee/tabby-webserver/src/schema/analytic.rs +++ b/ee/tabby-webserver/src/schema/analytic.rs @@ -8,6 +8,33 @@ use strum::{EnumIter, IntoEnumIterator}; use crate::schema::Result; +#[derive(GraphQLObject)] +pub struct DiskUsageStats { + pub events: DiskUsage, + pub indexed_repositories: DiskUsage, + pub database: DiskUsage, + pub models: DiskUsage, +} + +#[derive(GraphQLObject)] +pub struct DiskUsage { + pub file_paths: Vec, + pub size: f64, +} + +impl DiskUsage { + pub fn combine(self, other: Self) -> Self { + DiskUsage { + size: self.size + other.size, + file_paths: self + .file_paths + .into_iter() + .chain(other.file_paths) + .collect(), + } + } +} + #[derive(GraphQLObject, Debug, Clone)] pub struct CompletionStats { pub start: DateTime, @@ -101,4 +128,6 @@ pub trait AnalyticService: Send + Sync { users: Vec, languages: Vec, ) -> Result>; + + async fn disk_usage_stats(&self) -> Result; } diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs index a9311d9c2232..efe69d5935da 100644 --- a/ee/tabby-webserver/src/schema/mod.rs +++ b/ee/tabby-webserver/src/schema/mod.rs @@ -29,7 +29,7 @@ use validator::{Validate, ValidationErrors}; use worker::{Worker, WorkerService}; use self::{ - analytic::{AnalyticService, CompletionStats}, + analytic::{AnalyticService, CompletionStats, DiskUsageStats}, auth::{ JWTPayload, OAuthCredential, OAuthProvider, PasswordChangeInput, PasswordResetInput, RequestInvitationInput, RequestPasswordResetEmailInput, UpdateOAuthCredentialInput, @@ -457,6 +457,12 @@ impl Query { ) .await } + + async fn disk_usage_stats(ctx: &Context) -> Result { + check_admin(ctx).await?; + let storage_stats = ctx.locator.analytic().disk_usage_stats().await?; + Ok(storage_stats) + } } #[derive(GraphQLObject)] diff --git a/ee/tabby-webserver/src/service/analytic.rs b/ee/tabby-webserver/src/service/analytic.rs index 9e08b9ecb39f..4157b992910d 100644 --- a/ee/tabby-webserver/src/service/analytic.rs +++ b/ee/tabby-webserver/src/service/analytic.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{path::PathBuf, sync::Arc}; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -8,7 +8,7 @@ use tracing::warn; use super::AsRowid; use crate::schema::{ - analytic::{AnalyticService, CompletionStats, Language}, + analytic::{AnalyticService, CompletionStats, DiskUsage, DiskUsageStats, Language}, Result, }; @@ -70,6 +70,35 @@ impl AnalyticService for AnalyticServiceImpl { .collect(); Ok(stats) } + + async fn disk_usage_stats(&self) -> Result { + Ok(DiskUsageStats { + events: recursive_dir_size(tabby_common::path::events_dir()).await?, + indexed_repositories: recursive_dir_size(tabby_common::path::dataset_dir()) + .await? + .combine(recursive_dir_size(tabby_common::path::index_dir()).await?), + database: recursive_dir_size(crate::path::tabby_ee_root()).await?, + models: recursive_dir_size(tabby_common::path::models_dir()).await?, + }) + } +} + +/// Calculate the size of a directory in kilobytes recursively +async fn recursive_dir_size(path: PathBuf) -> Result { + let path_str = path.to_string_lossy().to_string(); + + let size = if path.exists() { + tokio::task::spawn_blocking(|| async { fs_extra::dir::get_size(path) }) + .await? + .await? + } else { + 0 + }; + + Ok(DiskUsage { + file_paths: vec![path_str], + size: size as f64 / 1024.0, + }) } fn convert_ids(ids: Vec) -> Vec { @@ -91,6 +120,8 @@ pub fn new_analytic_service(db: DbConn) -> Arc { #[cfg(test)] mod tests { use chrono::{Days, Duration}; + use tabby_common::path::set_tabby_root; + use temp_testdir::TempDir; use super::*; use crate::service::AsID; @@ -297,4 +328,31 @@ mod tests { .unwrap() .is_empty()); } + + #[tokio::test] + async fn test_disk_usage() { + let tmp_dir = TempDir::default(); + set_tabby_root(tmp_dir.to_path_buf()); + + tokio::fs::create_dir_all(tabby_common::path::models_dir()) + .await + .unwrap(); + + tokio::fs::write( + tabby_common::path::models_dir().join("testfile"), + "0".repeat(1024).as_bytes(), + ) + .await + .unwrap(); + + let db = DbConn::new_in_memory().await.unwrap(); + let service = new_analytic_service(db); + + let disk_usage = service.disk_usage_stats().await.unwrap(); + + assert_eq!(disk_usage.events.size, 0.0); + assert_eq!(disk_usage.indexed_repositories.size, 0.0); + assert_eq!(disk_usage.database.size, 0.0); + assert_eq!(disk_usage.models.size, 1.0); + } }