feat(drive): personal drives have the user's quota in commun
This commit is contained in:
@@ -15,7 +15,7 @@ use crate::infrastructure::repositories::pg::FileBlobWriteRepository;
|
||||
use crate::infrastructure::services::dedup_service::DedupService;
|
||||
use crate::infrastructure::services::file_content_cache::FileContentCache;
|
||||
use crate::infrastructure::services::pg_acl_engine::PgAclEngine;
|
||||
use tracing::{info, warn};
|
||||
use tracing::{Instrument, info, warn};
|
||||
|
||||
/// Service for file upload operations.
|
||||
///
|
||||
@@ -326,21 +326,39 @@ impl FileUploadService {
|
||||
};
|
||||
let delta = file.size as i64;
|
||||
|
||||
// Per-user delta — unchanged from `b5b80549` / `fbbae541`.
|
||||
if let Some(owner) = file
|
||||
let owner = file
|
||||
.owner_id
|
||||
.as_deref()
|
||||
.and_then(|s| Uuid::parse_str(s).ok())
|
||||
{
|
||||
.and_then(|s| Uuid::parse_str(s).ok());
|
||||
let folder = file
|
||||
.folder_id
|
||||
.as_deref()
|
||||
.and_then(|s| Uuid::parse_str(s).ok());
|
||||
|
||||
// Per-user delta — only when the target drive is `kind='personal'`.
|
||||
// The user envelope (`auth.users.storage_quota_bytes`) caps the SUM
|
||||
// of `used_bytes` across the user's personal drives; shared-drive
|
||||
// uploads do NOT count against any user. See
|
||||
// `docs/plan/drive.md` §7.
|
||||
//
|
||||
// The discrimination happens in one SQL statement via an EXISTS
|
||||
// subquery on the folder's drive kind — no extra round-trip vs
|
||||
// the unconditional delta. Without a folder id (root-level
|
||||
// upload — folder service refuses these) the user-side delta is
|
||||
// simply skipped; the sweep reconciles regardless.
|
||||
if let (Some(owner), Some(folder)) = (owner, folder) {
|
||||
let service_clone = Arc::clone(storage_service);
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = service_clone
|
||||
.add_user_storage_usage_delta(owner, delta)
|
||||
.await
|
||||
{
|
||||
warn!("Failed to bump storage usage for {owner}: {e}");
|
||||
tokio::spawn(
|
||||
async move {
|
||||
if let Err(e) = service_clone
|
||||
.add_user_storage_usage_delta_if_personal(owner, folder, delta)
|
||||
.await
|
||||
{
|
||||
warn!("Failed to bump user storage for {owner} (folder {folder}): {e}");
|
||||
}
|
||||
}
|
||||
});
|
||||
.in_current_span(),
|
||||
);
|
||||
}
|
||||
|
||||
// Per-drive delta (D4) — same fire-and-forget shape, resolves
|
||||
@@ -349,20 +367,19 @@ impl FileUploadService {
|
||||
// quota check and the picker quota bar read; drift from
|
||||
// deletes / trash is reconciled by the same sweep that handles
|
||||
// user-side drift.
|
||||
if let Some(folder) = file
|
||||
.folder_id
|
||||
.as_deref()
|
||||
.and_then(|s| Uuid::parse_str(s).ok())
|
||||
{
|
||||
if let Some(folder) = folder {
|
||||
let service_clone = Arc::clone(storage_service);
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = service_clone
|
||||
.add_drive_storage_usage_delta_by_folder(folder, delta)
|
||||
.await
|
||||
{
|
||||
warn!("Failed to bump drive usage for folder {folder}: {e}");
|
||||
tokio::spawn(
|
||||
async move {
|
||||
if let Err(e) = service_clone
|
||||
.add_drive_storage_usage_delta_by_folder(folder, delta)
|
||||
.await
|
||||
{
|
||||
warn!("Failed to bump drive usage for folder {folder}: {e}");
|
||||
}
|
||||
}
|
||||
});
|
||||
.in_current_span(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,13 +39,22 @@ impl StorageUsageService {
|
||||
/// (was three: user lookup + SUM + UPDATE). NOT called on the request
|
||||
/// path — only by the per-upload background update and the sweep.
|
||||
pub async fn update_user_storage_usage(&self, user_id: Uuid) -> Result<i64, DomainError> {
|
||||
// User envelope = SUM of `drives.used_bytes` across personal
|
||||
// drives owned by the user (see `docs/plan/drive.md` §7). Shared
|
||||
// drives don't count. Ownership is canonical via `role_grants`.
|
||||
let total_usage: Option<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
UPDATE auth.users u
|
||||
SET storage_used_bytes = COALESCE((
|
||||
SELECT SUM(f.size)::bigint
|
||||
FROM storage.files f
|
||||
WHERE f.user_id = u.id AND NOT f.is_trashed), 0)
|
||||
SELECT SUM(d.used_bytes)::bigint
|
||||
FROM storage.drives d
|
||||
JOIN storage.role_grants g
|
||||
ON g.resource_type = 'drive'
|
||||
AND g.resource_id = d.id
|
||||
AND g.role = 'owner'
|
||||
AND g.subject_type = 'user'
|
||||
AND g.subject_id = u.id
|
||||
WHERE d.kind = 'personal'), 0)
|
||||
WHERE u.id = $1
|
||||
RETURNING u.storage_used_bytes
|
||||
"#,
|
||||
@@ -77,9 +86,15 @@ impl StorageUsageService {
|
||||
r#"
|
||||
UPDATE auth.users u
|
||||
SET storage_used_bytes = COALESCE((
|
||||
SELECT SUM(f.size)::bigint
|
||||
FROM storage.files f
|
||||
WHERE f.user_id = u.id AND NOT f.is_trashed), 0)
|
||||
SELECT SUM(d.used_bytes)::bigint
|
||||
FROM storage.drives d
|
||||
JOIN storage.role_grants g
|
||||
ON g.resource_type = 'drive'
|
||||
AND g.resource_id = d.id
|
||||
AND g.role = 'owner'
|
||||
AND g.subject_type = 'user'
|
||||
AND g.subject_id = u.id
|
||||
WHERE d.kind = 'personal'), 0)
|
||||
WHERE u.username = $1
|
||||
RETURNING u.storage_used_bytes
|
||||
"#,
|
||||
@@ -128,6 +143,50 @@ impl StorageUsageService {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Conditional user-side delta: only fires when the target folder's
|
||||
/// drive is `kind='personal'`. See `docs/plan/drive.md` §7.
|
||||
///
|
||||
/// The new quota model: `auth.users.storage_quota_bytes` is the cap on
|
||||
/// the SUM of `used_bytes` across the user's personal drives. Shared
|
||||
/// drives never count against any user envelope. The upload hot path
|
||||
/// reads `drives.kind` from the same JOIN that already runs for the
|
||||
/// drive cap check; firing this conditional delta instead of the
|
||||
/// unconditional [`Self::add_user_storage_usage_delta`] keeps the
|
||||
/// counter aligned with that envelope semantics. Idempotent + clamped
|
||||
/// at zero, same as the unconditional sibling.
|
||||
///
|
||||
/// Implementation note: the EXISTS subquery is two indexed PK probes
|
||||
/// (folder by id, drive by id) so the personal/shared discrimination
|
||||
/// adds no real cost vs. the unconditional update.
|
||||
pub async fn add_user_storage_usage_delta_if_personal(
|
||||
&self,
|
||||
user_id: Uuid,
|
||||
folder_id: Uuid,
|
||||
delta: i64,
|
||||
) -> Result<(), DomainError> {
|
||||
sqlx::query(
|
||||
"UPDATE auth.users u
|
||||
SET storage_used_bytes = GREATEST(0, u.storage_used_bytes + $2)
|
||||
WHERE u.id = $1
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM storage.folders f
|
||||
JOIN storage.drives d ON d.id = f.drive_id
|
||||
WHERE f.id = $3
|
||||
AND d.kind = 'personal'
|
||||
)",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(delta)
|
||||
.bind(folder_id)
|
||||
.execute(self.pool.as_ref())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
DomainError::internal_error("StorageUsage", format!("usage delta if personal: {e}"))
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Incrementally adjust one drive's cached `storage.drives.used_bytes`
|
||||
/// by `delta` bytes — same shape as
|
||||
/// [`Self::add_user_storage_usage_delta`]: single statement, no
|
||||
@@ -312,15 +371,17 @@ impl StorageUsageService {
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
debug!("Running scheduled storage-usage reconciliation");
|
||||
if let Err(e) = service.update_all_users_storage_usage().await {
|
||||
error!("Scheduled user storage-usage reconciliation failed: {}", e);
|
||||
}
|
||||
// Drive sweep runs alongside the user sweep — same
|
||||
// cadence, same maintenance pool. Failure is logged
|
||||
// but doesn't skip the next tick.
|
||||
// Drive sweep runs FIRST: the user-side sweep below
|
||||
// reads `drives.used_bytes` (the per-drive sum) to
|
||||
// compute its own counter, so the drive counter must
|
||||
// be honest first. Failure of one is logged but
|
||||
// doesn't skip the other or the next tick.
|
||||
if let Err(e) = service.update_all_drives_storage_usage().await {
|
||||
error!("Scheduled drive storage-usage reconciliation failed: {}", e);
|
||||
}
|
||||
if let Err(e) = service.update_all_users_storage_usage().await {
|
||||
error!("Scheduled user storage-usage reconciliation failed: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -358,16 +419,33 @@ impl StorageUsagePort for StorageUsageService {
|
||||
async fn update_all_users_storage_usage(&self) -> Result<(), DomainError> {
|
||||
debug!("Starting storage-usage reconciliation sweep");
|
||||
|
||||
// User envelope = SUM of `drives.used_bytes` across the user's
|
||||
// personal drives. Shared drives don't count against any user
|
||||
// (`docs/plan/drive.md` §7). The drive-side sweep runs FIRST
|
||||
// (`start_reconciliation_job`) so `drives.used_bytes` is
|
||||
// already honest by the time we read it here.
|
||||
//
|
||||
// Ownership lookup uses `role_grants` (canonical per §1) so
|
||||
// both the user's default personal AND any secondary
|
||||
// personals owned via Owner grants are summed. Secondaries
|
||||
// aren't user-creatable today, but a backfill or admin path
|
||||
// can produce them — covering that surface from day one.
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE auth.users u
|
||||
SET storage_used_bytes = COALESCE(t.total, 0)
|
||||
FROM auth.users u2
|
||||
LEFT JOIN (
|
||||
SELECT user_id, SUM(size)::bigint AS total
|
||||
FROM storage.files
|
||||
WHERE NOT is_trashed
|
||||
GROUP BY user_id
|
||||
SELECT g.subject_id AS user_id,
|
||||
SUM(d.used_bytes)::bigint AS total
|
||||
FROM storage.drives d
|
||||
JOIN storage.role_grants g
|
||||
ON g.resource_type = 'drive'
|
||||
AND g.resource_id = d.id
|
||||
AND g.role = 'owner'
|
||||
AND g.subject_type = 'user'
|
||||
WHERE d.kind = 'personal'
|
||||
GROUP BY g.subject_id
|
||||
) t ON t.user_id = u2.id
|
||||
WHERE u.id = u2.id
|
||||
AND NOT u2.is_external
|
||||
|
||||
Reference in New Issue
Block a user