From 9f712e690b8c28a637ea58dbd6d65c0ef9e170ff Mon Sep 17 00:00:00 2001 From: DarkSky Date: Mon, 22 Jun 2026 00:51:47 +0800 Subject: [PATCH] fix: test & lint --- Cargo.lock | 2 + Cargo.toml | 1 + packages/backend/native/Cargo.toml | 2 + packages/backend/native/index.d.ts | 52 +++- .../src/backend_runtime/blob_complete.rs | 252 ++++++++++++++++++ .../native/src/backend_runtime/config.rs | 8 +- .../native/src/backend_runtime/constants.rs | 2 + .../src/backend_runtime/coordination_lease.rs | 9 +- .../src/backend_runtime/doc_compactor.rs | 15 +- .../native/src/backend_runtime/doc_storage.rs | 158 +++++++++++ .../backend/native/src/backend_runtime/mod.rs | 8 +- .../backend_runtime/object_storage/client.rs | 6 +- .../backend_runtime/object_storage/config.rs | 12 +- .../src/backend_runtime/object_storage/mod.rs | 7 +- .../runtime_state/byok_local_lease.rs | 136 ++++++++++ .../src/backend_runtime/runtime_state/dto.rs | 28 +- .../runtime_state/invite_link.rs | 31 ++- .../runtime_state/magic_link_otp.rs | 3 +- .../src/backend_runtime/runtime_state/mod.rs | 26 +- .../backend_runtime/runtime_state/store.rs | 20 +- .../runtime_state/verification_token.rs | 3 +- .../native/src/backend_runtime/tests.rs | 99 ++++--- .../native/src/backend_runtime/types.rs | 37 +++ .../src/backend_runtime/workspace_stats.rs | 158 ++++++++++- packages/backend/native/src/utils.rs | 7 +- .../migration.sql | 39 +++ packages/common/native/src/hashcash.rs | 2 +- 27 files changed, 1022 insertions(+), 101 deletions(-) create mode 100644 packages/backend/native/src/backend_runtime/blob_complete.rs create mode 100644 packages/backend/native/src/backend_runtime/doc_storage.rs create mode 100644 packages/backend/native/src/backend_runtime/runtime_state/byok_local_lease.rs create mode 100644 packages/backend/server/migrations/20260618000000_backend_runtime_tables/migration.sql diff --git a/Cargo.lock b/Cargo.lock index 9731f5ee42..5a021d9a80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -197,10 +197,12 @@ dependencies = [ "affine_common", "anyhow", "aws-sdk-s3", + "base64", "chrono", "doc_extractor", "file-format", "hex", + "homedir", "image", "infer", "jsonschema", diff --git a/Cargo.toml b/Cargo.toml index 4fc617d680..6825ce9386 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ resolver = "3" anyhow = "1" arbitrary = { version = "1.3", features = ["derive"] } assert-json-diff = "2.0" + base64 = "0.22.1" base64-simd = "0.8" bitvec = "1.0" block2 = "0.6" diff --git a/packages/backend/native/Cargo.toml b/packages/backend/native/Cargo.toml index 42b6a8da97..d53d121c90 100644 --- a/packages/backend/native/Cargo.toml +++ b/packages/backend/native/Cargo.toml @@ -17,10 +17,12 @@ affine_common = { workspace = true, features = [ ] } anyhow = { workspace = true } aws-sdk-s3 = "1.115" +base64 = { workspace = true } chrono = { workspace = true } doc_extractor = { workspace = true } file-format = { workspace = true } hex = { workspace = true } +homedir = { workspace = true } image = { workspace = true } infer = { workspace = true } jsonschema = "0.46" diff --git a/packages/backend/native/index.d.ts b/packages/backend/native/index.d.ts index 4b8c7b13cf..a17ad5b553 100644 --- a/packages/backend/native/index.d.ts +++ b/packages/backend/native/index.d.ts @@ -1,12 +1,23 @@ /* auto-generated by NAPI-RS */ /* eslint-disable */ export declare class BackendRuntime { + completeBlobUpload(workspaceId: string, key: string, expectedSize: number, expectedMime: string): Promise + completeFsBlobUpload(root: string, bucket: string, workspaceId: string, key: string, expectedSize: number, expectedMime: string): Promise cleanupExpiredPendingBlobs(cutoffMs: number, limit: number): Promise releaseDeletedBlobs(workspaceId: string, limit: number): Promise acquireCoordinationLease(key: string, owner: string, ttlMs: number): Promise - releaseCoordinationLease(key: string, owner: string, fencingToken: number): Promise - renewCoordinationLease(key: string, owner: string, fencingToken: number, ttlMs: number): Promise + releaseCoordinationLease(key: string, owner: string, fencingToken: bigint | number): Promise + renewCoordinationLease(key: string, owner: string, fencingToken: bigint | number, ttlMs: number): Promise + /** + * Merge pending doc updates with y-octo and persist the merged snapshot. + * + * Do not use this for snapshots that will be sent back to yjs clients until + * the y-octo/yjs round-trip compatibility issue is resolved. + */ compactPendingDocUpdates(workspaceId: string, docId: string, batchLimit: number, historyMinIntervalMs: number, owner: string, leaseTtlMs: number): Promise + upsertDocSnapshot(workspaceId: string, docId: string, blob: Buffer, timestampMs: number, editorId?: string | undefined | null): Promise + createDocHistory(input: RuntimeDocHistoryInput): Promise + deleteDocStorage(workspaceId: string, docId: string): Promise putRuntimeGateIfAbsent(key: string, ttlMs: number): Promise cleanupExpiredRuntimeGates(limit: number): Promise cleanupExpiredUserSessions(limit: number): Promise @@ -36,10 +47,13 @@ export declare class BackendRuntime { getWorkspaceInviteLink(workspaceId: string): Promise getWorkspaceInviteLinkById(inviteId: string): Promise revokeWorkspaceInviteLink(workspaceId: string): Promise + createByokLocalLease(activeKey: string, leaseId: string, payload: any, ttlMs: number): Promise + getByokLocalLease(leaseId: string): Promise cleanupExpiredRuntimeStates(limit: number): Promise refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise + recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise constructor() start(): Promise stop(): Promise @@ -220,7 +234,7 @@ export interface CommandResponse { export interface CoordinationLeaseGrant { key: string owner: string - fencingToken: number + fencingToken: bigint | number } /** @@ -804,6 +818,20 @@ export interface RuntimeBlobCleanupResult { workspaceIds: Array } +export interface RuntimeBlobCompleteResult { + ok: boolean + reason?: string + contentType?: string + contentLength?: number + lastModifiedMs?: number +} + +export interface RuntimeByokLocalLeaseRecord { + leaseId: string + payload: any + expiresAtMs: number +} + export interface RuntimeDocCompactionResult { leaseAcquired: boolean merged: boolean @@ -813,6 +841,17 @@ export interface RuntimeDocCompactionResult { historyCreated: boolean } +export interface RuntimeDocHistoryInput { + workspaceId: string + docId: string + blob: Buffer + timestampMs: number + editorId?: string + force: boolean + historyMinIntervalMs: number + historyMaxAgeMs: number +} + export interface RuntimeMagicLinkOtpConsumeResult { ok: boolean token?: string @@ -889,6 +928,13 @@ export interface RuntimeWorkspaceInviteLinkRecord { expiresAtMs: number } +export interface RuntimeWorkspaceStatsDailyRecalibrationResult { + processed: number + lastSid: number + snapshotted: number + skipped: boolean +} + export interface RuntimeWorkspaceStatsRecalibrationResult { processed: number lastSid: number diff --git a/packages/backend/native/src/backend_runtime/blob_complete.rs b/packages/backend/native/src/backend_runtime/blob_complete.rs new file mode 100644 index 0000000000..140561b19d --- /dev/null +++ b/packages/backend/native/src/backend_runtime/blob_complete.rs @@ -0,0 +1,252 @@ +use std::{ + fs, + path::{Path, PathBuf}, +}; + +use base64::{Engine as _, engine::general_purpose::URL_SAFE}; +use napi::Result; +use serde::Deserialize; +use sha2::{Digest, Sha256}; + +use super::{BackendRuntime, error::napi_error, types::RuntimeBlobCompleteResult}; + +fn object_missing_error(err: &napi::Error) -> bool { + let message = err.to_string(); + message.contains("NoSuchKey") || message.contains("NotFound") || message.contains("not found") +} + +fn blob_complete_failure(reason: &str) -> RuntimeBlobCompleteResult { + RuntimeBlobCompleteResult { + ok: false, + reason: Some(reason.to_string()), + content_type: None, + content_length: None, + last_modified_ms: None, + } +} + +fn blob_complete_success( + content_type: String, + content_length: i64, + last_modified_ms: i64, +) -> RuntimeBlobCompleteResult { + RuntimeBlobCompleteResult { + ok: true, + reason: None, + content_type: Some(content_type), + content_length: Some(content_length), + last_modified_ms: Some(last_modified_ms), + } +} + +fn sha256_base64_url_with_padding(body: &[u8]) -> String { + URL_SAFE.encode(Sha256::digest(body)) +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct FsBlobMetadata { + content_type: String, + content_length: i64, + last_modified: i64, +} + +fn normalize_storage_key(key: &str) -> Result> { + let normalized = key.replace('\\', "/"); + let segments = normalized.split('/').map(ToString::to_string).collect::>(); + + if normalized.is_empty() + || normalized.starts_with('/') + || segments + .iter() + .any(|segment| segment.is_empty() || segment == "." || segment == "..") + { + return Err(napi_error(format!("Invalid storage key: {key}"))); + } + + Ok(segments) +} + +fn fs_bucket_path(root: &str, bucket: &str) -> PathBuf { + if let Some(stripped) = root.strip_prefix("~/") + && let Ok(Some(home)) = homedir::my_home() + { + return home.join(stripped).join(bucket); + } + + Path::new(root).join(bucket) +} + +fn fs_object_path(root: &str, bucket: &str, key: &str) -> Result { + let mut path = fs_bucket_path(root, bucket); + for segment in normalize_storage_key(key)? { + path.push(segment); + } + Ok(path) +} + +fn read_fs_metadata(path: &Path) -> Result> { + let metadata_path = PathBuf::from(format!("{}.metadata.json", path.display())); + let raw = match fs::read_to_string(metadata_path) { + Ok(raw) => raw, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(err) => { + return Err(napi_error(format!("BlobComplete read fs metadata failed: {err}"))); + } + }; + + serde_json::from_str(&raw).map(Some).map_err(|err| { + napi_error(format!( + "BlobComplete parse fs metadata failed for {}: {err}", + path.display() + )) + }) +} + +async fn upsert_completed_blob( + runtime: &BackendRuntime, + workspace_id: &str, + key: &str, + mime: &str, + size: i64, +) -> Result<()> { + sqlx::query( + r#" + INSERT INTO blobs (workspace_id, key, mime, size, status, upload_id) + VALUES ($1, $2, $3, $4, 'completed', NULL) + ON CONFLICT (workspace_id, key) + DO UPDATE SET + mime = EXCLUDED.mime, + size = EXCLUDED.size, + status = EXCLUDED.status, + upload_id = NULL + "#, + ) + .bind(workspace_id) + .bind(key) + .bind(mime) + .bind(size as i32) + .execute(&runtime.pool().await?) + .await + .map_err(|err| napi_error(format!("BlobComplete upsert metadata failed: {err}")))?; + + Ok(()) +} + +#[napi_derive::napi] +impl BackendRuntime { + #[napi] + pub async fn complete_blob_upload( + &self, + workspace_id: String, + key: String, + expected_size: i64, + expected_mime: String, + ) -> Result { + let object_key = format!("{workspace_id}/{key}"); + let object = match self.object_storage_get(object_key.clone()).await { + Ok(Some(object)) => object, + Ok(None) => return Ok(blob_complete_failure("not_found")), + Err(err) if object_missing_error(&err) => return Ok(blob_complete_failure("not_found")), + Err(err) => return Err(err), + }; + + if object.metadata.content_length != expected_size { + return Ok(blob_complete_failure("size_mismatch")); + } + + if !expected_mime.is_empty() && object.metadata.content_type != expected_mime { + return Ok(blob_complete_failure("mime_mismatch")); + } + + if sha256_base64_url_with_padding(&object.body) != key { + match self.object_storage_delete(object_key).await { + Ok(()) => {} + Err(err) if object_missing_error(&err) => {} + Err(err) => return Err(err), + } + return Ok(blob_complete_failure("checksum_mismatch")); + } + + upsert_completed_blob( + self, + &workspace_id, + &key, + &object.metadata.content_type, + object.metadata.content_length, + ) + .await?; + + Ok(blob_complete_success( + object.metadata.content_type, + object.metadata.content_length, + object.metadata.last_modified_ms, + )) + } + + #[napi] + pub async fn complete_fs_blob_upload( + &self, + root: String, + bucket: String, + workspace_id: String, + key: String, + expected_size: i64, + expected_mime: String, + ) -> Result { + let storage_key = format!("{workspace_id}/{key}"); + let path = fs_object_path(&root, &bucket, &storage_key)?; + let metadata = match read_fs_metadata(&path)? { + Some(metadata) => metadata, + None => return Ok(blob_complete_failure("not_found")), + }; + + if metadata.content_length != expected_size { + return Ok(blob_complete_failure("size_mismatch")); + } + + if !expected_mime.is_empty() && metadata.content_type != expected_mime { + return Ok(blob_complete_failure("mime_mismatch")); + } + + let body = match fs::read(&path) { + Ok(body) => body, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(blob_complete_failure("not_found")), + Err(err) => return Err(napi_error(format!("BlobComplete read fs object failed: {err}"))), + }; + + if sha256_base64_url_with_padding(&body) != key { + let _ = fs::remove_file(&path); + let _ = fs::remove_file(PathBuf::from(format!("{}.metadata.json", path.display()))); + return Ok(blob_complete_failure("checksum_mismatch")); + } + + upsert_completed_blob( + self, + &workspace_id, + &key, + &metadata.content_type, + metadata.content_length, + ) + .await?; + + Ok(blob_complete_success( + metadata.content_type, + metadata.content_length, + metadata.last_modified, + )) + } +} + +#[cfg(test)] +mod tests { + use super::sha256_base64_url_with_padding; + + #[test] + fn sha256_base64_url_keeps_padding() { + assert_eq!( + sha256_base64_url_with_padding(b"hello"), + "LPJNul-wow4m6DsqxbninhsWHlwfp0JecwQzYpOLmCQ=" + ); + } +} diff --git a/packages/backend/native/src/backend_runtime/config.rs b/packages/backend/native/src/backend_runtime/config.rs index f133e3d6ef..9543f72c8e 100644 --- a/packages/backend/native/src/backend_runtime/config.rs +++ b/packages/backend/native/src/backend_runtime/config.rs @@ -49,10 +49,10 @@ fn database_url_from_config_files() -> Result> { .map_err(|err| napi_error(format!("failed to read config file {}: {err}", path.display())))?; let config: AppConfigFile = serde_json::from_str(&raw) .map_err(|err| napi_error(format!("failed to parse config file {}: {err}", path.display())))?; - if let Some(next) = config.db.and_then(|db| db.datasource_url) { - if !next.trim().is_empty() { - database_url = Some(next); - } + if let Some(next) = config.db.and_then(|db| db.datasource_url) + && !next.trim().is_empty() + { + database_url = Some(next); } } diff --git a/packages/backend/native/src/backend_runtime/constants.rs b/packages/backend/native/src/backend_runtime/constants.rs index a332a51150..fc3138696a 100644 --- a/packages/backend/native/src/backend_runtime/constants.rs +++ b/packages/backend/native/src/backend_runtime/constants.rs @@ -1,4 +1,6 @@ pub(super) const DEFAULT_HISTORY_PERIOD_SECONDS: i32 = 7 * 24 * 60 * 60; +pub(super) const BYOK_LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active"; +pub(super) const BYOK_LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease"; pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp"; pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10; pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id"; diff --git a/packages/backend/native/src/backend_runtime/coordination_lease.rs b/packages/backend/native/src/backend_runtime/coordination_lease.rs index 14891691d2..b048a05e92 100644 --- a/packages/backend/native/src/backend_runtime/coordination_lease.rs +++ b/packages/backend/native/src/backend_runtime/coordination_lease.rs @@ -108,7 +108,12 @@ impl BackendRuntime { } #[napi] - pub async fn release_coordination_lease(&self, key: String, owner: String, fencing_token: i64) -> Result { + pub async fn release_coordination_lease( + &self, + key: String, + owner: String, + #[napi(ts_arg_type = "bigint | number")] fencing_token: i64, + ) -> Result { CoordinationLeaseStore::new(self.pool().await?) .release(&key, &owner, fencing_token) .await @@ -119,7 +124,7 @@ impl BackendRuntime { &self, key: String, owner: String, - fencing_token: i64, + #[napi(ts_arg_type = "bigint | number")] fencing_token: i64, ttl_ms: i64, ) -> Result { if ttl_ms <= 0 { diff --git a/packages/backend/native/src/backend_runtime/doc_compactor.rs b/packages/backend/native/src/backend_runtime/doc_compactor.rs index 1abd17000b..644c42217e 100644 --- a/packages/backend/native/src/backend_runtime/doc_compactor.rs +++ b/packages/backend/native/src/backend_runtime/doc_compactor.rs @@ -313,12 +313,11 @@ async fn compact_doc( .await?; let mut history_created = false; - if snapshot_updated { - if let Some(snapshot) = &snapshot { - if should_create_history(&mut tx, snapshot, workspace_id, doc_id, history_min_interval_ms).await? { - history_created = create_history(&mut tx, workspace_id, doc_id, snapshot).await?; - } - } + if snapshot_updated + && let Some(snapshot) = &snapshot + && should_create_history(&mut tx, snapshot, workspace_id, doc_id, history_min_interval_ms).await? + { + history_created = create_history(&mut tx, workspace_id, doc_id, snapshot).await?; } let timestamps = updates.iter().map(|update| update.created_at).collect::>(); @@ -333,6 +332,10 @@ async fn compact_doc( #[napi_derive::napi] impl BackendRuntime { + /// Merge pending doc updates with y-octo and persist the merged snapshot. + /// + /// Do not use this for snapshots that will be sent back to yjs clients until + /// the y-octo/yjs round-trip compatibility issue is resolved. #[napi] pub async fn compact_pending_doc_updates( &self, diff --git a/packages/backend/native/src/backend_runtime/doc_storage.rs b/packages/backend/native/src/backend_runtime/doc_storage.rs new file mode 100644 index 0000000000..0f2f61a8e6 --- /dev/null +++ b/packages/backend/native/src/backend_runtime/doc_storage.rs @@ -0,0 +1,158 @@ +use chrono::{DateTime, Duration, Utc}; +use napi::{Result, bindgen_prelude::Buffer}; +use sqlx::{PgPool, Row}; + +use super::{BackendRuntime, error::napi_error, types::RuntimeDocHistoryInput}; + +fn is_empty_doc(bin: &[u8]) -> bool { + bin.is_empty() || (bin.len() == 1 && bin[0] == 0) || (bin.len() == 2 && bin[0] == 0 && bin[1] == 0) +} + +async fn latest_history_timestamp(pool: &PgPool, workspace_id: &str, doc_id: &str) -> Result>> { + sqlx::query( + r#" + SELECT timestamp + FROM snapshot_histories + WHERE workspace_id = $1 AND guid = $2 + ORDER BY timestamp DESC + LIMIT 1 + "#, + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_optional(pool) + .await + .map(|row| row.map(|row| row.get("timestamp"))) + .map_err(|err| napi_error(format!("DocStorage load latest history failed: {err}"))) +} + +#[napi_derive::napi] +impl BackendRuntime { + #[napi] + pub async fn upsert_doc_snapshot( + &self, + workspace_id: String, + doc_id: String, + blob: Buffer, + timestamp_ms: i64, + editor_id: Option, + ) -> Result { + if is_empty_doc(blob.as_ref()) { + return Ok(false); + } + + let timestamp = DateTime::::from_timestamp_millis(timestamp_ms) + .ok_or_else(|| napi_error(format!("Invalid doc snapshot timestamp: {timestamp_ms}")))?; + let pool = self.pool().await?; + let row = sqlx::query( + r#" + INSERT INTO snapshots + (workspace_id, guid, blob, size, created_at, updated_at, created_by, updated_by) + VALUES + ($1, $2, $3, $4, $5, $5, $6, $6) + ON CONFLICT (workspace_id, guid) + DO UPDATE SET + blob = $3, + size = $4, + updated_at = $5, + updated_by = $6 + WHERE snapshots.workspace_id = $1 + AND snapshots.guid = $2 + AND snapshots.updated_at <= $5 + RETURNING updated_at + "#, + ) + .bind(&workspace_id) + .bind(&doc_id) + .bind(blob.as_ref()) + .bind(blob.len() as i64) + .bind(timestamp) + .bind(editor_id.as_deref()) + .fetch_optional(&pool) + .await + .map_err(|err| napi_error(format!("DocStorage upsert snapshot failed: {err}")))?; + + Ok(row.is_some()) + } + + #[napi] + pub async fn create_doc_history(&self, input: RuntimeDocHistoryInput) -> Result { + if input.history_min_interval_ms < 0 { + return Err(napi_error("doc history interval must be non-negative")); + } + if input.history_max_age_ms <= 0 || is_empty_doc(input.blob.as_ref()) { + return Ok(false); + } + + let timestamp = DateTime::::from_timestamp_millis(input.timestamp_ms) + .ok_or_else(|| napi_error(format!("Invalid doc history timestamp: {}", input.timestamp_ms)))?; + let pool = self.pool().await?; + let should_create = match latest_history_timestamp(&pool, &input.workspace_id, &input.doc_id).await? { + None => true, + Some(last_timestamp) if last_timestamp == timestamp => false, + Some(last_timestamp) => { + input.force || last_timestamp < timestamp - Duration::milliseconds(input.history_min_interval_ms) + } + }; + + if !should_create { + return Ok(false); + } + + let expired_at = Utc::now() + Duration::milliseconds(input.history_max_age_ms); + sqlx::query( + r#" + INSERT INTO snapshot_histories + (workspace_id, guid, timestamp, blob, expired_at, created_by) + VALUES + ($1, $2, $3, $4, $5, $6) + ON CONFLICT (workspace_id, guid, timestamp) + DO UPDATE SET expired_at = EXCLUDED.expired_at + "#, + ) + .bind(&input.workspace_id) + .bind(&input.doc_id) + .bind(timestamp) + .bind(input.blob.as_ref()) + .bind(expired_at) + .bind(input.editor_id.as_deref()) + .execute(&pool) + .await + .map_err(|err| napi_error(format!("DocStorage create history failed: {err}")))?; + + Ok(true) + } + + #[napi] + pub async fn delete_doc_storage(&self, workspace_id: String, doc_id: String) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool + .begin() + .await + .map_err(|err| napi_error(format!("DocStorage delete begin transaction failed: {err}")))?; + + sqlx::query("DELETE FROM snapshots WHERE workspace_id = $1 AND guid = $2") + .bind(&workspace_id) + .bind(&doc_id) + .execute(&mut *tx) + .await + .map_err(|err| napi_error(format!("DocStorage delete snapshot failed: {err}")))?; + sqlx::query("DELETE FROM updates WHERE workspace_id = $1 AND guid = $2") + .bind(&workspace_id) + .bind(&doc_id) + .execute(&mut *tx) + .await + .map_err(|err| napi_error(format!("DocStorage delete updates failed: {err}")))?; + sqlx::query("DELETE FROM snapshot_histories WHERE workspace_id = $1 AND guid = $2") + .bind(&workspace_id) + .bind(&doc_id) + .execute(&mut *tx) + .await + .map_err(|err| napi_error(format!("DocStorage delete histories failed: {err}")))?; + + tx.commit() + .await + .map_err(|err| napi_error(format!("DocStorage delete commit failed: {err}")))?; + Ok(()) + } +} diff --git a/packages/backend/native/src/backend_runtime/mod.rs b/packages/backend/native/src/backend_runtime/mod.rs index 9fd3389c63..ec547f0916 100644 --- a/packages/backend/native/src/backend_runtime/mod.rs +++ b/packages/backend/native/src/backend_runtime/mod.rs @@ -1,8 +1,10 @@ +mod blob_complete; mod blob_reclaimer; mod config; mod constants; mod coordination_lease; mod doc_compactor; +mod doc_storage; mod error; mod gate; mod housekeeping; @@ -76,8 +78,8 @@ impl BackendRuntime { #[napi] pub async fn health(&self) -> Result { - let guard = self.pool.lock().await; - let database_connected = match guard.as_ref() { + let pool = self.pool.lock().await.as_ref().cloned(); + let database_connected = match pool.as_ref() { Some(pool) => sqlx::query("SELECT 1") .fetch_one(pool) .await @@ -87,7 +89,7 @@ impl BackendRuntime { }; Ok(BackendRuntimeHealth { - started: guard.is_some(), + started: pool.is_some(), database_connected, object_storage_configured: self.config.storage.is_some(), }) diff --git a/packages/backend/native/src/backend_runtime/object_storage/client.rs b/packages/backend/native/src/backend_runtime/object_storage/client.rs index 1b751a9faa..a92c3e15f7 100644 --- a/packages/backend/native/src/backend_runtime/object_storage/client.rs +++ b/packages/backend/native/src/backend_runtime/object_storage/client.rs @@ -8,12 +8,11 @@ use aws_sdk_s3::{ }; use napi::Result; -use crate::backend_runtime::error::napi_error; - use super::types::{ MultipartUploadInitResult, MultipartUploadPart, ObjectGetResult, ObjectListEntry, ObjectMetadata, ObjectPutMetadata, PresignedObjectRequest, completed_multipart_parts, trim_etag, }; +use crate::backend_runtime::error::napi_error; #[derive(Clone)] pub(super) struct ObjectStorageClient { @@ -121,9 +120,10 @@ impl ObjectStorageClient { )) })?; + let expires_at_ms = expires_at_ms(self.presign_expires_in_seconds)?; Ok(result.upload_id.map(|upload_id| MultipartUploadInitResult { upload_id, - expires_at_ms: expires_at_ms(self.presign_expires_in_seconds).unwrap_or(0), + expires_at_ms, })) } diff --git a/packages/backend/native/src/backend_runtime/object_storage/config.rs b/packages/backend/native/src/backend_runtime/object_storage/config.rs index fd82c76b75..c29a8af97a 100644 --- a/packages/backend/native/src/backend_runtime/object_storage/config.rs +++ b/packages/backend/native/src/backend_runtime/object_storage/config.rs @@ -1,15 +1,14 @@ use aws_sdk_s3::config::{ - BehaviorVersion, Credentials, Region, RequestChecksumCalculation, ResponseChecksumValidation, + BehaviorVersion, Credentials, Region, RequestChecksumCalculation, ResponseChecksumValidation, timeout::TimeoutConfig, }; use napi::Result; use serde::Deserialize; +use super::{client::ObjectStorageClient, types::StorageProviderConfig}; use crate::backend_runtime::{ config::blob_storage_config_from_config_files, error::napi_error, types::RuntimeObjectStorageHealth, }; -use super::{client::ObjectStorageClient, types::StorageProviderConfig}; - #[derive(Clone, Debug)] pub(in crate::backend_runtime) struct ObjectStorageConfig { pub(super) provider: String, @@ -175,6 +174,13 @@ impl ObjectStorageConfig { if let Some(endpoint) = &self.endpoint { builder = builder.endpoint_url(endpoint); } + if let Some(request_timeout_ms) = self.request_timeout_ms { + builder = builder.timeout_config( + TimeoutConfig::builder() + .operation_timeout(std::time::Duration::from_millis(request_timeout_ms)) + .build(), + ); + } Ok(ObjectStorageClient::new( builder.build(), diff --git a/packages/backend/native/src/backend_runtime/object_storage/mod.rs b/packages/backend/native/src/backend_runtime/object_storage/mod.rs index 1b8e7345ae..9016f5fa1b 100644 --- a/packages/backend/native/src/backend_runtime/object_storage/mod.rs +++ b/packages/backend/native/src/backend_runtime/object_storage/mod.rs @@ -4,7 +4,10 @@ mod config; mod tests; mod types; +use client::ObjectStorageClient; +pub(super) use config::ObjectStorageConfig; use napi::{Result, bindgen_prelude::Buffer}; +pub(super) use types::StorageProviderConfig; use super::{ BackendRuntime, @@ -14,10 +17,6 @@ use super::{ }, }; -use client::ObjectStorageClient; -pub(super) use config::ObjectStorageConfig; -pub(super) use types::StorageProviderConfig; - #[napi_derive::napi] impl BackendRuntime { fn object_storage_client(&self) -> Result { diff --git a/packages/backend/native/src/backend_runtime/runtime_state/byok_local_lease.rs b/packages/backend/native/src/backend_runtime/runtime_state/byok_local_lease.rs new file mode 100644 index 0000000000..57c00fca43 --- /dev/null +++ b/packages/backend/native/src/backend_runtime/runtime_state/byok_local_lease.rs @@ -0,0 +1,136 @@ +use napi::Result; + +use super::dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows}; +use crate::backend_runtime::{ + constants::{BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE}, + error::napi_error, + types::RuntimeByokLocalLeaseRecord, +}; + +pub(super) async fn get(rows: &RuntimeStateRows, lease_id: String) -> Result> { + get_lease_by_id(rows, &lease_id).await +} + +pub(super) async fn create( + rows: &RuntimeStateRows, + active_key: String, + lease_id: String, + payload: serde_json::Value, + ttl_ms: i64, +) -> Result { + if ttl_ms <= 0 { + return Err(napi_error("BYOK local lease ttl must be positive")); + } + + let mut tx = rows.begin("RuntimeState BYOK local lease").await?; + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(&active_key) + .execute(&mut *tx) + .await + .map_err(|err| napi_error(format!("RuntimeState BYOK local lease active lock failed: {err}")))?; + + if let Some(active) = rows + .active_payload_with_expires_for_update_in_tx( + &mut tx, + BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, + &active_key, + "RuntimeState BYOK local lease active get", + ) + .await? + { + let existing_lease = match active.payload.get("leaseId").and_then(serde_json::Value::as_str) { + Some(existing_lease_id) => get_lease_by_id_in_tx(rows, &mut tx, existing_lease_id).await?, + None => None, + }; + if let Some(lease) = existing_lease { + tx.commit().await.map_err(|err| { + napi_error(format!( + "RuntimeState BYOK local lease transaction commit failed: {err}" + )) + })?; + return Ok(lease); + } + + rows + .delete_by_key_in_tx( + &mut tx, + BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, + &active_key, + "RuntimeState BYOK local lease stale active delete", + ) + .await?; + } + + let expires_at_ms = rows + .insert_payload_returning_expires_in_tx( + &mut tx, + RuntimeStateInsertPayload { + purpose: BYOK_LOCAL_LEASE_PURPOSE, + token: &lease_id, + lookup_key: &active_key, + payload: &payload, + ttl_ms, + context: "RuntimeState BYOK local lease create", + }, + ) + .await?; + let active_payload = serde_json::json!({ "leaseId": lease_id }); + rows + .insert_payload_returning_expires_in_tx( + &mut tx, + RuntimeStateInsertPayload { + purpose: BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, + token: &active_key, + lookup_key: &active_key, + payload: &active_payload, + ttl_ms, + context: "RuntimeState BYOK local lease active create", + }, + ) + .await?; + + tx.commit().await.map_err(|err| { + napi_error(format!( + "RuntimeState BYOK local lease transaction commit failed: {err}" + )) + })?; + + Ok(RuntimeByokLocalLeaseRecord { + lease_id, + payload, + expires_at_ms, + }) +} + +async fn get_lease_by_id(rows: &RuntimeStateRows, lease_id: &str) -> Result> { + rows + .active_payload_with_expires(BYOK_LOCAL_LEASE_PURPOSE, lease_id, "RuntimeState BYOK local lease get") + .await? + .map(|row| record_from_row(lease_id, row)) + .transpose() +} + +async fn get_lease_by_id_in_tx( + rows: &RuntimeStateRows, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + lease_id: &str, +) -> Result> { + rows + .active_payload_with_expires_for_update_in_tx( + tx, + BYOK_LOCAL_LEASE_PURPOSE, + lease_id, + "RuntimeState BYOK local lease get", + ) + .await? + .map(|row| record_from_row(lease_id, row)) + .transpose() +} + +fn record_from_row(lease_id: &str, row: RuntimeStatePayloadRow) -> Result { + Ok(RuntimeByokLocalLeaseRecord { + lease_id: lease_id.to_string(), + payload: row.payload, + expires_at_ms: row.expires_at_ms, + }) +} diff --git a/packages/backend/native/src/backend_runtime/runtime_state/dto.rs b/packages/backend/native/src/backend_runtime/runtime_state/dto.rs index df9ddac45a..8bdea9bf44 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/dto.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/dto.rs @@ -14,6 +14,15 @@ pub(super) struct RuntimeStateLockedRow { pub(super) expires_at: chrono::DateTime, } +pub(super) struct RuntimeStateInsertPayload<'a> { + pub(super) purpose: &'a str, + pub(super) token: &'a str, + pub(super) lookup_key: &'a str, + pub(super) payload: &'a serde_json::Value, + pub(super) ttl_ms: i64, + pub(super) context: &'a str, +} + #[derive(Clone)] pub(super) struct RuntimeStateRows { pub(super) pool: PgPool, @@ -291,12 +300,7 @@ impl RuntimeStateRows { pub(super) async fn insert_payload_returning_expires_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - purpose: &str, - token: &str, - lookup_key: &str, - payload: &serde_json::Value, - ttl_ms: i64, - context: &str, + input: RuntimeStateInsertPayload<'_>, ) -> Result { let row = sqlx::query( r#" @@ -305,14 +309,14 @@ impl RuntimeStateRows { RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms "#, ) - .bind(purpose) - .bind(token_hash(token)) - .bind(lookup_key) - .bind(payload) - .bind(ttl_ms as f64) + .bind(input.purpose) + .bind(token_hash(input.token)) + .bind(input.lookup_key) + .bind(input.payload) + .bind(input.ttl_ms as f64) .fetch_one(&mut **tx) .await - .map_err(|err| napi_error(format!("{context} failed: {err}")))?; + .map_err(|err| napi_error(format!("{} failed: {err}", input.context)))?; Ok(row.get::("expires_at_ms")) } diff --git a/packages/backend/native/src/backend_runtime/runtime_state/invite_link.rs b/packages/backend/native/src/backend_runtime/runtime_state/invite_link.rs index d25a01a49c..02109703d0 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/invite_link.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/invite_link.rs @@ -1,13 +1,12 @@ use napi::Result; +use super::dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows}; use crate::backend_runtime::{ constants::{WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE}, error::napi_error, types::RuntimeWorkspaceInviteLinkRecord, }; -use super::dto::{RuntimeStatePayloadRow, RuntimeStateRows}; - pub(super) async fn get_by_workspace( rows: &RuntimeStateRows, workspace_id: String, @@ -55,23 +54,27 @@ pub(super) async fn create( let expires_at_ms = rows .insert_payload_returning_expires_in_tx( &mut tx, - WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, - &workspace_id, - &workspace_id, - &payload, - ttl_ms, - "RuntimeState workspace invite link create", + RuntimeStateInsertPayload { + purpose: WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, + token: &workspace_id, + lookup_key: &workspace_id, + payload: &payload, + ttl_ms, + context: "RuntimeState workspace invite link create", + }, ) .await?; rows .insert_payload_returning_expires_in_tx( &mut tx, - WORKSPACE_INVITE_LINK_ID_PURPOSE, - &invite_id, - &invite_id, - &payload, - ttl_ms, - "RuntimeState workspace invite link create", + RuntimeStateInsertPayload { + purpose: WORKSPACE_INVITE_LINK_ID_PURPOSE, + token: &invite_id, + lookup_key: &invite_id, + payload: &payload, + ttl_ms, + context: "RuntimeState workspace invite link create", + }, ) .await?; diff --git a/packages/backend/native/src/backend_runtime/runtime_state/magic_link_otp.rs b/packages/backend/native/src/backend_runtime/runtime_state/magic_link_otp.rs index 5b74c31412..a26106edfe 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/magic_link_otp.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/magic_link_otp.rs @@ -1,13 +1,12 @@ use napi::Result; +use super::dto::RuntimeStateRows; use crate::backend_runtime::{ constants::{MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS}, error::napi_error, types::RuntimeMagicLinkOtpConsumeResult, }; -use super::dto::RuntimeStateRows; - impl RuntimeMagicLinkOtpConsumeResult { fn ok(token: String) -> Self { Self { diff --git a/packages/backend/native/src/backend_runtime/runtime_state/mod.rs b/packages/backend/native/src/backend_runtime/runtime_state/mod.rs index 28cda3aa0c..39ed1f7610 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/mod.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/mod.rs @@ -3,10 +3,14 @@ use napi::Result; use super::{ BackendRuntime, error::napi_error, - types::{RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord}, + types::{ + RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, + RuntimeWorkspaceInviteLinkRecord, + }, }; mod auth_challenge; +mod byok_local_lease; mod dto; mod invite_link; mod magic_link_otp; @@ -172,6 +176,26 @@ impl BackendRuntime { .await } + #[napi] + pub async fn create_byok_local_lease( + &self, + active_key: String, + lease_id: String, + payload: serde_json::Value, + ttl_ms: i64, + ) -> Result { + RuntimeStateStore::new(self.pool().await?) + .create_byok_local_lease(active_key, lease_id, payload, ttl_ms) + .await + } + + #[napi] + pub async fn get_byok_local_lease(&self, lease_id: String) -> Result> { + RuntimeStateStore::new(self.pool().await?) + .get_byok_local_lease(lease_id) + .await + } + #[napi] pub async fn cleanup_expired_runtime_states(&self, limit: i64) -> Result { if limit <= 0 { diff --git a/packages/backend/native/src/backend_runtime/runtime_state/store.rs b/packages/backend/native/src/backend_runtime/runtime_state/store.rs index 5157c7ee82..4d12a356b2 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/store.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/store.rs @@ -1,12 +1,12 @@ use napi::Result; use sqlx::PgPool; +use super::{auth_challenge, byok_local_lease, dto::RuntimeStateRows, invite_link, magic_link_otp, verification_token}; use crate::backend_runtime::types::{ - RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord, + RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, + RuntimeWorkspaceInviteLinkRecord, }; -use super::{auth_challenge, dto::RuntimeStateRows, invite_link, magic_link_otp, verification_token}; - pub(super) struct RuntimeStateStore { rows: RuntimeStateRows, } @@ -122,4 +122,18 @@ impl RuntimeStateStore { pub(super) async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result { invite_link::revoke(&self.rows, workspace_id).await } + + pub(super) async fn create_byok_local_lease( + &self, + active_key: String, + lease_id: String, + payload: serde_json::Value, + ttl_ms: i64, + ) -> Result { + byok_local_lease::create(&self.rows, active_key, lease_id, payload, ttl_ms).await + } + + pub(super) async fn get_byok_local_lease(&self, lease_id: String) -> Result> { + byok_local_lease::get(&self.rows, lease_id).await + } } diff --git a/packages/backend/native/src/backend_runtime/runtime_state/verification_token.rs b/packages/backend/native/src/backend_runtime/runtime_state/verification_token.rs index e9ba14518d..6de836870b 100644 --- a/packages/backend/native/src/backend_runtime/runtime_state/verification_token.rs +++ b/packages/backend/native/src/backend_runtime/runtime_state/verification_token.rs @@ -2,12 +2,11 @@ use napi::Result; use sqlx::{PgPool, Row}; use uuid::Uuid; -use crate::backend_runtime::{error::napi_error, token_hash, types::RuntimeVerificationTokenRecord}; - use super::{ dto::{RuntimeStatePayloadRow, RuntimeStateRows}, verification_token_purpose, }; +use crate::backend_runtime::{error::napi_error, token_hash, types::RuntimeVerificationTokenRecord}; pub(super) async fn create( rows: &RuntimeStateRows, diff --git a/packages/backend/native/src/backend_runtime/tests.rs b/packages/backend/native/src/backend_runtime/tests.rs index f5e793a624..86a692b95b 100644 --- a/packages/backend/native/src/backend_runtime/tests.rs +++ b/packages/backend/native/src/backend_runtime/tests.rs @@ -1,6 +1,9 @@ use super::{runtime_state::*, *}; +use anyhow::{Context, Result as AnyResult, anyhow}; + static PG_TEST_LOCK: std::sync::OnceLock> = std::sync::OnceLock::new(); +const TEST_VERIFICATION_TOKEN_TYPE: i32 = 99_999; fn pg_test_lock() -> &'static tokio::sync::Mutex<()> { PG_TEST_LOCK.get_or_init(|| tokio::sync::Mutex::new(())) @@ -30,48 +33,52 @@ fn verification_token_state_uses_typed_purpose_and_token_hash() { assert_ne!(token_hash("verification-token"), token_hash("other-token")); } -async fn runtime_from_database_url() -> Option { - let database_url = std::env::var("DATABASE_URL").ok()?; +async fn runtime_from_database_url() -> AnyResult> { + let Ok(database_url) = std::env::var("DATABASE_URL") else { + return Ok(None); + }; let pool = PgPoolOptions::new() .max_connections(5) .connect(&database_url) .await - .ok()?; - migrate_runtime_tables(&pool).await.ok()?; + .context("connect postgres for backend runtime tests")?; + migrate_runtime_tables(&pool) + .await + .map_err(|err| anyhow!(err.to_string()))?; sqlx::query( r#" DELETE FROM runtime_states WHERE purpose LIKE 'rust_test:%' OR purpose LIKE 'auth_challenge:rust_test:%' - OR purpose LIKE 'verification_token:%' + OR purpose = 'verification_token:99999' "#, ) .execute(&pool) .await - .ok()?; + .context("cleanup runtime_states for backend runtime tests")?; sqlx::query("DELETE FROM runtime_gates WHERE key LIKE 'rust-test:%'") .execute(&pool) .await - .ok()?; + .context("cleanup runtime_gates for backend runtime tests")?; sqlx::query("DELETE FROM runtime_leases WHERE key LIKE 'rust-test:%'") .execute(&pool) .await - .ok()?; + .context("cleanup runtime_leases for backend runtime tests")?; - Some(BackendRuntime { + Ok(Some(BackendRuntime { config: RuntimeConfig { database_url, storage: None, }, pool: Mutex::new(Some(pool)), - }) + })) } #[tokio::test] async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() { let _guard = pg_test_lock().lock().await; - let Some(runtime) = runtime_from_database_url().await else { - eprintln!("skipping postgres integration test: DATABASE_URL is not set or unavailable"); + let Some(runtime) = runtime_from_database_url().await.unwrap() else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); return; }; @@ -151,8 +158,8 @@ async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() { #[tokio::test] async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() { let _guard = pg_test_lock().lock().await; - let Some(runtime) = runtime_from_database_url().await else { - eprintln!("skipping postgres integration test: DATABASE_URL is not set or unavailable"); + let Some(runtime) = runtime_from_database_url().await.unwrap() else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); return; }; @@ -248,8 +255,8 @@ async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() { #[tokio::test] async fn runtime_state_cleanup_deletes_expired_and_consumed_rows() { let _guard = pg_test_lock().lock().await; - let Some(runtime) = runtime_from_database_url().await else { - eprintln!("skipping postgres integration test: DATABASE_URL is not set or unavailable"); + let Some(runtime) = runtime_from_database_url().await.unwrap() else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); return; }; @@ -291,65 +298,92 @@ async fn runtime_state_cleanup_deletes_expired_and_consumed_rows() { #[tokio::test] async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup() { let _guard = pg_test_lock().lock().await; - let Some(runtime) = runtime_from_database_url().await else { - eprintln!("skipping postgres integration test: DATABASE_URL is not set or unavailable"); + let Some(runtime) = runtime_from_database_url().await.unwrap() else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); return; }; let mismatch_token = runtime - .create_verification_token(0, Some("user@affine.test".to_string()), 30_000) + .create_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + Some("user@affine.test".to_string()), + 30_000, + ) .await .unwrap(); assert!( runtime - .verify_verification_token(0, mismatch_token.clone(), Some("wrong@affine.test".to_string()), None) + .verify_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + mismatch_token.clone(), + Some("wrong@affine.test".to_string()), + None, + ) .await .unwrap() .is_none() ); assert!( runtime - .verify_verification_token(0, mismatch_token.clone(), Some("user@affine.test".to_string()), None) + .verify_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + mismatch_token.clone(), + Some("user@affine.test".to_string()), + None, + ) .await .unwrap() .is_some() ); assert!( runtime - .verify_verification_token(0, mismatch_token.clone(), Some("user@affine.test".to_string()), None) + .verify_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + mismatch_token.clone(), + Some("user@affine.test".to_string()), + None, + ) .await .unwrap() .is_none() ); let keep_token = runtime - .create_verification_token(0, Some("keep@affine.test".to_string()), 30_000) + .create_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + Some("keep@affine.test".to_string()), + 30_000, + ) .await .unwrap(); assert!( runtime - .get_verification_token(0, keep_token.clone(), Some(true)) + .get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), Some(true)) .await .unwrap() .is_some() ); assert!( runtime - .get_verification_token(0, keep_token.clone(), None) + .get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None) .await .unwrap() .is_some() ); assert!( runtime - .get_verification_token(0, keep_token.clone(), None) + .get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None) .await .unwrap() .is_none() ); let concurrent_token = runtime - .create_verification_token(0, Some("concurrent@affine.test".to_string()), 30_000) + .create_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + Some("concurrent@affine.test".to_string()), + 30_000, + ) .await .unwrap(); let mut tasks = Vec::new(); @@ -361,7 +395,12 @@ async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup() let token = concurrent_token.clone(); tasks.push(tokio::spawn(async move { runtime - .verify_verification_token(0, token, Some("concurrent@affine.test".to_string()), None) + .verify_verification_token( + TEST_VERIFICATION_TOKEN_TYPE, + token, + Some("concurrent@affine.test".to_string()), + None, + ) .await .unwrap() .is_some() @@ -376,13 +415,13 @@ async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup() assert_eq!(successful, 1); let expired_token = runtime - .create_verification_token(0, Some("expired@affine.test".to_string()), 1) + .create_verification_token(TEST_VERIFICATION_TOKEN_TYPE, Some("expired@affine.test".to_string()), 1) .await .unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; assert!( runtime - .get_verification_token(0, expired_token.clone(), None) + .get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, expired_token.clone(), None) .await .unwrap() .is_none() diff --git a/packages/backend/native/src/backend_runtime/types.rs b/packages/backend/native/src/backend_runtime/types.rs index 40622385d6..513859a8da 100644 --- a/packages/backend/native/src/backend_runtime/types.rs +++ b/packages/backend/native/src/backend_runtime/types.rs @@ -36,6 +36,7 @@ pub struct RuntimeObjectStorageHealth { pub struct CoordinationLeaseGrant { pub key: String, pub owner: String, + #[napi(ts_type = "bigint | number")] pub fencing_token: i64, } @@ -54,6 +55,25 @@ pub struct RuntimeWorkspaceInviteLinkRecord { pub expires_at_ms: i64, } +#[napi_derive::napi(object)] +pub struct RuntimeByokLocalLeaseRecord { + pub lease_id: String, + pub payload: serde_json::Value, + pub expires_at_ms: i64, +} + +#[napi_derive::napi(object)] +pub struct RuntimeDocHistoryInput { + pub workspace_id: String, + pub doc_id: String, + pub blob: Buffer, + pub timestamp_ms: i64, + pub editor_id: Option, + pub force: bool, + pub history_min_interval_ms: i64, + pub history_max_age_ms: i64, +} + #[napi_derive::napi(object)] pub struct RuntimeObjectStoragePutOptions { pub content_type: Option, @@ -109,6 +129,15 @@ pub struct RuntimeBlobCleanupResult { pub workspace_ids: Vec, } +#[napi_derive::napi(object)] +pub struct RuntimeBlobCompleteResult { + pub ok: bool, + pub reason: Option, + pub content_type: Option, + pub content_length: Option, + pub last_modified_ms: Option, +} + #[napi_derive::napi(object)] pub struct RuntimeDocCompactionResult { pub lease_acquired: bool, @@ -138,3 +167,11 @@ pub struct RuntimeWorkspaceStatsSnapshotResult { pub snapshotted: i64, pub skipped: bool, } + +#[napi_derive::napi(object)] +pub struct RuntimeWorkspaceStatsDailyRecalibrationResult { + pub processed: i64, + pub last_sid: i64, + pub snapshotted: i64, + pub skipped: bool, +} diff --git a/packages/backend/native/src/backend_runtime/workspace_stats.rs b/packages/backend/native/src/backend_runtime/workspace_stats.rs index 598751e27c..70e749b00d 100644 --- a/packages/backend/native/src/backend_runtime/workspace_stats.rs +++ b/packages/backend/native/src/backend_runtime/workspace_stats.rs @@ -1,13 +1,14 @@ use napi::Result; use sqlx::{FromRow, PgPool, Postgres, Row, Transaction}; +use tokio::time::{Duration as TokioDuration, sleep}; use super::{ BackendRuntime, constants::{WORKSPACE_STATS_LEASE_KEY, WORKSPACE_STATS_LOCK_NAMESPACE, WORKSPACE_STATS_REFRESH_LOCK_KEY}, error::napi_error, types::{ - CoordinationLeaseGrant, RuntimeWorkspaceStatsRecalibrationResult, RuntimeWorkspaceStatsRefreshResult, - RuntimeWorkspaceStatsSnapshotResult, + CoordinationLeaseGrant, RuntimeWorkspaceStatsDailyRecalibrationResult, RuntimeWorkspaceStatsRecalibrationResult, + RuntimeWorkspaceStatsRefreshResult, RuntimeWorkspaceStatsSnapshotResult, }, }; @@ -108,6 +109,94 @@ impl BackendRuntime { release_workspace_stats_lease(self, lease).await?; result } + + #[napi] + pub async fn recalibrate_workspace_admin_stats_daily( + &self, + batch_limit: i64, + owner: String, + lease_ttl_ms: i64, + lock_retry_times: i64, + lock_retry_delay_ms: i64, + ) -> Result { + if batch_limit <= 0 { + return Err(napi_error("workspace stats daily recalibration limit must be positive")); + } + if lock_retry_times <= 0 { + return Err(napi_error( + "workspace stats daily recalibration retry times must be positive", + )); + } + if lock_retry_delay_ms < 0 { + return Err(napi_error( + "workspace stats daily recalibration retry delay must be non-negative", + )); + } + + let Some(lease) = acquire_workspace_stats_lease_with_retry( + self, + owner.clone(), + lease_ttl_ms, + lock_retry_times, + lock_retry_delay_ms, + ) + .await? + else { + return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult { + processed: 0, + last_sid: 0, + snapshotted: 0, + skipped: true, + }); + }; + + let result = async { + let store = WorkspaceStatsStore::new(self.pool().await?); + let mut processed = 0; + let mut last_sid = 0; + + loop { + let batch = retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || { + store.recalibrate(last_sid, batch_limit) + }) + .await?; + + if batch.skipped { + return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult { + processed, + last_sid, + snapshotted: 0, + skipped: true, + }); + } + + if batch.processed == 0 { + break; + } + + processed += batch.processed; + last_sid = batch.last_sid; + + if batch.processed < batch_limit { + break; + } + } + + let snapshot = + retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || store.write_daily_snapshot()).await?; + + Ok(RuntimeWorkspaceStatsDailyRecalibrationResult { + processed, + last_sid, + snapshotted: snapshot.snapshotted, + skipped: snapshot.skipped, + }) + } + .await; + + release_workspace_stats_lease(self, lease).await?; + result + } } #[derive(FromRow)] @@ -251,6 +340,69 @@ async fn release_workspace_stats_lease(runtime: &BackendRuntime, lease: Coordina Ok(()) } +async fn acquire_workspace_stats_lease_with_retry( + runtime: &BackendRuntime, + owner: String, + lease_ttl_ms: i64, + retry_times: i64, + retry_delay_ms: i64, +) -> Result> { + for attempt in 0..retry_times { + let lease = runtime + .acquire_coordination_lease(WORKSPACE_STATS_LEASE_KEY.to_string(), owner.clone(), lease_ttl_ms) + .await?; + if lease.is_some() { + return Ok(lease); + } + + if attempt < retry_times - 1 && retry_delay_ms > 0 { + sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await; + } + } + + Ok(None) +} + +async fn retry_workspace_stats_operation( + retry_times: i64, + retry_delay_ms: i64, + mut operation: F, +) -> Result +where + T: WorkspaceStatsSkippable, + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + for attempt in 0..retry_times { + let result = operation().await?; + if !result.skipped() || attempt == retry_times - 1 { + return Ok(result); + } + + if retry_delay_ms > 0 { + sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await; + } + } + + unreachable!("workspace stats retry loop validates retry_times > 0") +} + +trait WorkspaceStatsSkippable { + fn skipped(&self) -> bool; +} + +impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsRecalibrationResult { + fn skipped(&self) -> bool { + self.skipped + } +} + +impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsSnapshotResult { + fn skipped(&self) -> bool { + self.skipped + } +} + async fn try_transaction_lock(tx: &mut Transaction<'_, Postgres>) -> Result { let row = sqlx::query( r#" @@ -333,7 +485,7 @@ async fn fetch_workspace_batch( LIMIT $2 "#, ) - .bind(last_sid as i32) + .bind(last_sid) .bind(limit) .fetch_all(&mut **tx) .await diff --git a/packages/backend/native/src/utils.rs b/packages/backend/native/src/utils.rs index 143f1ff6bc..6940aa6253 100644 --- a/packages/backend/native/src/utils.rs +++ b/packages/backend/native/src/utils.rs @@ -26,11 +26,8 @@ fn try_remove_label(s: &str, i: usize) -> Option { return None; } - if let Some(ch) = s[next_idx..].chars().next() { - if !ch.is_whitespace() { - return None; - } - } else { + let ch = s[next_idx..].chars().next()?; + if !ch.is_whitespace() { return None; } diff --git a/packages/backend/server/migrations/20260618000000_backend_runtime_tables/migration.sql b/packages/backend/server/migrations/20260618000000_backend_runtime_tables/migration.sql new file mode 100644 index 0000000000..24a86dc91d --- /dev/null +++ b/packages/backend/server/migrations/20260618000000_backend_runtime_tables/migration.sql @@ -0,0 +1,39 @@ +CREATE TABLE "runtime_states" ( + "purpose" TEXT NOT NULL, + "token_hash" TEXT NOT NULL, + "lookup_key" TEXT, + "payload" JSONB NOT NULL, + "attempts" INTEGER NOT NULL DEFAULT 0, + "consumed_at" TIMESTAMPTZ(3), + "expires_at" TIMESTAMPTZ(3) NOT NULL, + "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "runtime_states_pkey" PRIMARY KEY ("purpose", "token_hash") +); + +CREATE INDEX "runtime_states_lookup_idx" ON "runtime_states"("purpose", "lookup_key") WHERE "lookup_key" IS NOT NULL AND "consumed_at" IS NULL; +CREATE INDEX "runtime_states_expires_at_idx" ON "runtime_states"("expires_at"); + +CREATE TABLE "runtime_gates" ( + "key" TEXT NOT NULL, + "expires_at" TIMESTAMPTZ(3) NOT NULL, + "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "runtime_gates_pkey" PRIMARY KEY ("key") +); + +CREATE INDEX "runtime_gates_expires_at_idx" ON "runtime_gates"("expires_at"); + +CREATE TABLE "runtime_leases" ( + "key" TEXT NOT NULL, + "owner" TEXT NOT NULL, + "fencing_token" BIGINT NOT NULL, + "expires_at" TIMESTAMPTZ(3) NOT NULL, + "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "runtime_leases_pkey" PRIMARY KEY ("key") +); + +CREATE INDEX "runtime_leases_expires_at_idx" ON "runtime_leases"("expires_at"); diff --git a/packages/common/native/src/hashcash.rs b/packages/common/native/src/hashcash.rs index a01c4ada8a..1c2354483c 100644 --- a/packages/common/native/src/hashcash.rs +++ b/packages/common/native/src/hashcash.rs @@ -61,7 +61,7 @@ impl Stamp { let ts = now.format("%Y%m%d%H%M%S"); let bits = bits.unwrap_or(20); let rand = String::from_iter(Alphanumeric.sample_iter(rng()).take(SALT_LENGTH).map(char::from)); - let challenge = format!("{}:{}:{}:{}:{}:{}", version, bits, ts, &resource, "", rand); + let challenge = format!("{}:{}:{}:{}:{}:{}", version, bits, ts, resource, "", rand); Stamp { version: version.to_string(),