feat(server): impl storage runtime (#15181)

#### PR Dependency Tree


* **PR #15181** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an additional storage backend option: asset-pack based storage
(provider for avatar, blob, and copilot).
* Introduced a dedicated storage runtime with provider capability
reporting and expanded object operations (put/head/get/list/delete),
including presigned and multipart flows where supported.
* Cloudflare R2 `jurisdiction` now uses an explicit default when
omitted.
* **Bug Fixes**
  * Broadened avatar access to allow both fs and asset-pack providers.
* Improved workspace blob upload completion validation and handling when
stored objects are missing or mismatched.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-07-01 22:24:10 +08:00
committed by GitHub
parent da7d438377
commit 8ebdb7452f
102 changed files with 6487 additions and 4508 deletions

View File

@@ -0,0 +1,9 @@
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";
pub(super) const WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE: &str = "workspace_invite_link:workspace";
pub(super) const WORKSPACE_STATS_LEASE_KEY: &str = "workspace:admin-stats:refresh";
pub(super) const WORKSPACE_STATS_LOCK_NAMESPACE: i64 = 97_301;
pub(super) const WORKSPACE_STATS_REFRESH_LOCK_KEY: i64 = 1;

View File

@@ -0,0 +1,163 @@
use napi::Result;
use sqlx::{FromRow, PgPool};
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error, types::CoordinationLeaseGrant};
#[derive(FromRow)]
struct LeaseGrantRow {
fencing_token: i64,
}
struct CoordinationLeaseStore {
pool: PgPool,
}
impl CoordinationLeaseStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn acquire(&self, key: String, owner: String, ttl_ms: i64) -> RuntimeResult<Option<CoordinationLeaseGrant>> {
let row = sqlx::query_as::<_, LeaseGrantRow>(
r#"
INSERT INTO runtime_leases (key, owner, fencing_token, expires_at)
VALUES ($1, $2, 1, CURRENT_TIMESTAMP + ($3 * INTERVAL '1 millisecond'))
ON CONFLICT (key) DO UPDATE
SET owner = EXCLUDED.owner,
fencing_token = runtime_leases.fencing_token + 1,
expires_at = EXCLUDED.expires_at,
updated_at = CURRENT_TIMESTAMP
WHERE runtime_leases.expires_at <= CURRENT_TIMESTAMP
RETURNING fencing_token
"#,
)
.bind(&key)
.bind(&owner)
.bind(ttl_ms as f64)
.fetch_optional(&self.pool)
.await
.map_err(|err| RuntimeError::database("CoordinationLease acquire failed", err))?;
Ok(row.map(|row| CoordinationLeaseGrant {
key,
owner,
fencing_token: row.fencing_token,
}))
}
async fn release(&self, key: &str, owner: &str, fencing_token: i64) -> RuntimeResult<bool> {
let result = sqlx::query(
r#"
DELETE FROM runtime_leases
WHERE key = $1 AND owner = $2 AND fencing_token = $3
"#,
)
.bind(key)
.bind(owner)
.bind(fencing_token)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database("CoordinationLease release failed", err))?;
Ok(result.rows_affected() == 1)
}
async fn renew(&self, key: &str, owner: &str, fencing_token: i64, ttl_ms: i64) -> RuntimeResult<bool> {
let result = sqlx::query(
r#"
UPDATE runtime_leases
SET expires_at = CURRENT_TIMESTAMP + ($4 * INTERVAL '1 millisecond'),
updated_at = CURRENT_TIMESTAMP
WHERE key = $1
AND owner = $2
AND fencing_token = $3
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(key)
.bind(owner)
.bind(fencing_token)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database("CoordinationLease renew failed", err))?;
Ok(result.rows_affected() == 1)
}
}
#[napi_derive::napi]
impl BackendRuntime {
pub(crate) async fn acquire_coordination_lease_inner(
&self,
key: String,
owner: String,
ttl_ms: i64,
) -> RuntimeResult<Option<CoordinationLeaseGrant>> {
if ttl_ms <= 0 {
return Err(RuntimeError::invalid_input("coordination lease ttl must be positive"));
}
if owner.is_empty() {
return Err(RuntimeError::invalid_input("coordination lease owner is required"));
}
CoordinationLeaseStore::new(self.pool().await?)
.acquire(key, owner, ttl_ms)
.await
}
pub(crate) async fn release_coordination_lease_inner(
&self,
key: String,
owner: String,
fencing_token: i64,
) -> RuntimeResult<bool> {
CoordinationLeaseStore::new(self.pool().await?)
.release(&key, &owner, fencing_token)
.await
}
#[napi]
pub async fn acquire_coordination_lease(
&self,
key: String,
owner: String,
ttl_ms: i64,
) -> Result<Option<CoordinationLeaseGrant>> {
self
.acquire_coordination_lease_inner(key, owner, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn release_coordination_lease(
&self,
key: String,
owner: String,
#[napi(ts_arg_type = "bigint | number")] fencing_token: i64,
) -> Result<bool> {
self
.release_coordination_lease_inner(key, owner, fencing_token)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn renew_coordination_lease(
&self,
key: String,
owner: String,
#[napi(ts_arg_type = "bigint | number")] fencing_token: i64,
ttl_ms: i64,
) -> Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("coordination lease ttl must be positive"));
}
CoordinationLeaseStore::new(self.pool().await?)
.renew(&key, &owner, fencing_token, ttl_ms)
.await
.map_err(napi::Error::from)
}
}

View File

@@ -0,0 +1,410 @@
use chrono::{DateTime, Duration, Utc};
use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use y_octo::Doc;
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error, types::RuntimeDocCompactionResult};
#[derive(FromRow)]
struct SnapshotRow {
blob: Vec<u8>,
updated_at: DateTime<Utc>,
updated_by: Option<String>,
}
#[derive(FromRow)]
struct UpdateRow {
blob: Vec<u8>,
created_at: DateTime<Utc>,
created_by: Option<String>,
}
struct DocCompactorStore {
pool: PgPool,
}
impl DocCompactorStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn compact_doc(
&self,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
history_min_interval_ms: i64,
history_max_age_seconds: i64,
) -> RuntimeResult<(i64, bool)> {
compact_doc(
self.pool.clone(),
workspace_id,
doc_id,
batch_limit,
history_min_interval_ms,
history_max_age_seconds,
)
.await
}
}
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)
}
fn apply_updates(updates: impl IntoIterator<Item = Vec<u8>>) -> RuntimeResult<Vec<u8>> {
let mut doc = Doc::default();
for update in updates {
doc
.apply_update_from_binary_v1(&update)
.map_err(|err| RuntimeError::invalid_state(format!("DocCompactor merge failed: {err}")))?;
}
doc
.encode_update_v1()
.map_err(|err| RuntimeError::invalid_state(format!("DocCompactor encode failed: {err}")))
}
fn checked_milliseconds(value: i64, field: &str) -> RuntimeResult<Duration> {
Duration::try_milliseconds(value)
.ok_or_else(|| RuntimeError::invalid_input(format!("DocCompactor {field} is too large")))
}
fn checked_seconds(value: i64, field: &str) -> RuntimeResult<Duration> {
Duration::try_seconds(value).ok_or_else(|| RuntimeError::invalid_input(format!("DocCompactor {field} is too large")))
}
async fn load_snapshot(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<SnapshotRow>> {
sqlx::query_as::<_, SnapshotRow>(
r#"
SELECT blob, updated_at, updated_by
FROM snapshots
WHERE workspace_id = $1 AND guid = $2
FOR UPDATE
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor load snapshot failed", err))
}
async fn load_updates(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
) -> RuntimeResult<Vec<UpdateRow>> {
sqlx::query_as::<_, UpdateRow>(
r#"
SELECT blob, created_at, created_by
FROM updates
WHERE workspace_id = $1 AND guid = $2
ORDER BY created_at ASC
LIMIT $3
FOR UPDATE
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(batch_limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor load updates failed", err))
}
async fn upsert_snapshot(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
blob: &[u8],
timestamp: DateTime<Utc>,
editor: Option<&str>,
) -> RuntimeResult<bool> {
if is_empty_doc(blob) {
return Ok(false);
}
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)
.bind(blob.len() as i64)
.bind(timestamp)
.bind(editor)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor upsert snapshot failed", err))?;
Ok(row.is_some())
}
async fn should_create_history(
tx: &mut Transaction<'_, Postgres>,
snapshot: &SnapshotRow,
workspace_id: &str,
doc_id: &str,
history_min_interval_ms: i64,
) -> RuntimeResult<bool> {
if is_empty_doc(&snapshot.blob) {
return Ok(false);
}
let row = 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(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor load latest history failed", err))?;
let Some(row) = row else {
return Ok(true);
};
let last_timestamp: DateTime<Utc> = row.get("timestamp");
if last_timestamp == snapshot.updated_at {
return Ok(false);
}
let min_interval = checked_milliseconds(history_min_interval_ms, "history interval")?;
let threshold = snapshot
.updated_at
.checked_sub_signed(min_interval)
.ok_or_else(|| RuntimeError::invalid_input("DocCompactor history interval is out of range"))?;
Ok(last_timestamp < threshold)
}
async fn create_history(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
snapshot: &SnapshotRow,
max_age_seconds: i64,
) -> RuntimeResult<bool> {
if max_age_seconds <= 0 {
return Ok(false);
}
let max_age = checked_seconds(max_age_seconds, "history max age")?;
let expired_at = Utc::now()
.checked_add_signed(max_age)
.ok_or_else(|| RuntimeError::invalid_input("DocCompactor history max age is out of range"))?;
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(workspace_id)
.bind(doc_id)
.bind(snapshot.updated_at)
.bind(&snapshot.blob)
.bind(expired_at)
.bind(snapshot.updated_by.as_deref())
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor create history failed", err))?;
Ok(true)
}
async fn delete_updates(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
timestamps: &[DateTime<Utc>],
) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
DELETE FROM updates
WHERE workspace_id = $1
AND guid = $2
AND created_at = ANY($3)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(timestamps)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("DocCompactor delete updates failed", err))?;
Ok(result.rows_affected() as i64)
}
async fn compact_doc(
pool: PgPool,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
history_min_interval_ms: i64,
history_max_age_seconds: i64,
) -> RuntimeResult<(i64, bool)> {
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("DocCompactor begin transaction failed", err))?;
let snapshot = load_snapshot(&mut tx, workspace_id, doc_id).await?;
let updates = load_updates(&mut tx, workspace_id, doc_id, batch_limit).await?;
if updates.is_empty() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("DocCompactor commit transaction failed", err))?;
return Ok((0, false));
}
let last = updates.last().expect("updates is not empty");
let mut merge_inputs = Vec::with_capacity(updates.len() + usize::from(snapshot.is_some()));
if let Some(snapshot) = &snapshot {
merge_inputs.push(snapshot.blob.clone());
}
merge_inputs.extend(updates.iter().map(|update| update.blob.clone()));
let final_blob = if merge_inputs.len() == 1 {
merge_inputs.remove(0)
} else {
apply_updates(merge_inputs)?
};
let snapshot_updated = upsert_snapshot(
&mut tx,
workspace_id,
doc_id,
&final_blob,
last.created_at,
last.created_by.as_deref(),
)
.await?;
let mut history_created = false;
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, history_max_age_seconds).await?;
}
let timestamps = updates.iter().map(|update| update.created_at).collect::<Vec<_>>();
let deleted = delete_updates(&mut tx, workspace_id, doc_id, &timestamps).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("DocCompactor commit transaction failed", err))?;
Ok((deleted, history_created))
}
#[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.
///
/// The caller owns quota reconciliation and must pass a fresh
/// history_max_age_seconds value. The compactor intentionally does not read
/// effective_workspace_quota_states; if a future caller cannot provide a
/// fresh quota state, fail and retry after Node reconciles it.
#[napi]
#[allow(clippy::too_many_arguments)]
pub async fn compact_pending_doc_updates(
&self,
workspace_id: String,
doc_id: String,
batch_limit: i64,
history_min_interval_ms: i64,
history_max_age_seconds: i64,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeDocCompactionResult> {
if batch_limit <= 0 {
return Err(napi_error("doc compactor batch limit must be positive"));
}
if history_min_interval_ms < 0 {
return Err(napi_error("doc compactor history interval must be non-negative"));
}
if history_max_age_seconds < 0 {
return Err(napi_error("doc compactor history max age must be non-negative"));
}
checked_milliseconds(history_min_interval_ms, "history interval")?;
if history_max_age_seconds > 0 {
let max_age = checked_seconds(history_max_age_seconds, "history max age")?;
Utc::now()
.checked_add_signed(max_age)
.ok_or_else(|| RuntimeError::invalid_input("DocCompactor history max age is out of range"))?;
}
let lease_key = format!("doc:update:{workspace_id}:{doc_id}");
let Some(lease) = self.acquire_coordination_lease(lease_key, owner, lease_ttl_ms).await? else {
return Ok(RuntimeDocCompactionResult {
lease_acquired: false,
merged: false,
workspace_id,
doc_id,
updates_merged: 0,
history_created: false,
});
};
let result = DocCompactorStore::new(self.pool().await?)
.compact_doc(
&workspace_id,
&doc_id,
batch_limit,
history_min_interval_ms,
history_max_age_seconds,
)
.await;
let released = self
.release_coordination_lease(lease.key, lease.owner, lease.fencing_token)
.await?;
if !released {
return Err(RuntimeError::invalid_state("DocCompactor failed to release coordination lease").into());
}
let (updates_merged, history_created) = result?;
Ok(RuntimeDocCompactionResult {
lease_acquired: true,
merged: updates_merged > 0,
workspace_id,
doc_id,
updates_merged,
history_created,
})
}
}

View File

@@ -0,0 +1,162 @@
use chrono::{DateTime, Duration, Utc};
use napi::bindgen_prelude::Buffer;
use sqlx::{PgPool, Row};
use super::{BackendRuntime, RuntimeError, RuntimeResult, 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,
) -> RuntimeResult<Option<DateTime<Utc>>> {
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| RuntimeError::database("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<String>,
) -> napi::Result<bool> {
if is_empty_doc(blob.as_ref()) {
return Ok(false);
}
let timestamp = DateTime::<Utc>::from_timestamp_millis(timestamp_ms)
.ok_or_else(|| RuntimeError::invalid_input(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| RuntimeError::database("DocStorage upsert snapshot failed", err))?;
Ok(row.is_some())
}
#[napi]
pub async fn create_doc_history(&self, input: RuntimeDocHistoryInput) -> napi::Result<bool> {
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::<Utc>::from_timestamp_millis(input.timestamp_ms)
.ok_or_else(|| RuntimeError::invalid_input(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| RuntimeError::database("DocStorage create history failed", err))?;
Ok(true)
}
#[napi]
pub async fn delete_doc_storage(&self, workspace_id: String, doc_id: String) -> napi::Result<()> {
let pool = self.pool().await?;
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("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| RuntimeError::database("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| RuntimeError::database("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| RuntimeError::database("DocStorage delete histories failed", err))?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("DocStorage delete commit failed", err))?;
Ok(())
}
}

View File

@@ -0,0 +1,94 @@
use napi::Result;
use sqlx::PgPool;
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error};
struct RuntimeGateStore {
pool: PgPool,
}
impl RuntimeGateStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn put_if_absent(&self, key: &str, ttl_ms: i64) -> RuntimeResult<bool> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("RuntimeGate transaction failed", err))?;
sqlx::query("DELETE FROM runtime_gates WHERE key = $1 AND expires_at <= CURRENT_TIMESTAMP")
.bind(key)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("RuntimeGate expired cleanup failed", err))?;
let inserted = sqlx::query(
r#"
INSERT INTO runtime_gates (key, expires_at)
VALUES ($1, CURRENT_TIMESTAMP + ($2 * INTERVAL '1 millisecond'))
ON CONFLICT (key) DO NOTHING
"#,
)
.bind(key)
.bind(ttl_ms as f64)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("RuntimeGate put_if_absent failed", err))?
.rows_affected()
== 1;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeGate transaction commit failed", err))?;
Ok(inserted)
}
async fn cleanup_expired(&self, limit: i64) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_gates
WHERE key IN (
SELECT key FROM runtime_gates
WHERE expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database("RuntimeGate cleanup failed", err))?;
Ok(result.rows_affected() as i64)
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn put_runtime_gate_if_absent(&self, key: String, ttl_ms: i64) -> Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("runtime gate ttl must be positive"));
}
RuntimeGateStore::new(self.pool().await?)
.put_if_absent(&key, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn cleanup_expired_runtime_gates(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("runtime gate cleanup limit must be positive"));
}
RuntimeGateStore::new(self.pool().await?)
.cleanup_expired(limit)
.await
.map_err(napi::Error::from)
}
}

View File

@@ -0,0 +1,82 @@
use napi::Result;
use sqlx::PgPool;
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error};
struct HousekeepingStore {
pool: PgPool,
}
impl HousekeepingStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn cleanup_expired_user_sessions(&self, limit: i64) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
DELETE FROM user_sessions
WHERE id IN (
SELECT id FROM user_sessions
WHERE expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database("Housekeeping user sessions cleanup failed", err))?;
Ok(result.rows_affected() as i64)
}
async fn cleanup_expired_snapshot_histories(&self, limit: i64) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
DELETE FROM snapshot_histories
WHERE (workspace_id, guid, timestamp) IN (
SELECT workspace_id, guid, timestamp
FROM snapshot_histories
WHERE expired_at <= CURRENT_TIMESTAMP
ORDER BY expired_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database("Housekeeping snapshot histories cleanup failed", err))?;
Ok(result.rows_affected() as i64)
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn cleanup_expired_user_sessions(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("user sessions cleanup limit must be positive"));
}
HousekeepingStore::new(self.pool().await?)
.cleanup_expired_user_sessions(limit)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn cleanup_expired_snapshot_histories(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("snapshot histories cleanup limit must be positive"));
}
HousekeepingStore::new(self.pool().await?)
.cleanup_expired_snapshot_histories(limit)
.await
.map_err(napi::Error::from)
}
}

View File

@@ -0,0 +1,133 @@
mod constants;
mod coordination_lease;
mod doc_compactor;
mod doc_storage;
mod gate;
mod housekeeping;
mod runtime_state;
#[cfg(test)]
mod tests;
mod workspace_stats;
use std::{sync::RwLock, time::Duration};
use napi::Result;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
use tokio::sync::Mutex;
use self::types::BackendRuntimeHealth;
pub(crate) use super::types;
use super::{
BackendRuntimeConfig, RuntimeError, RuntimeResult, migrations::migrate_runtime_tables, napi_error, to_napi_error,
};
pub(super) fn token_hash(token: &str) -> String {
hex::encode(Sha256::digest(token.as_bytes()))
}
#[napi_derive::napi]
pub struct BackendRuntime {
config: RwLock<BackendRuntimeConfig>,
pool: Mutex<Option<PgPool>>,
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi(constructor)]
pub fn new() -> Result<Self> {
Ok(Self {
config: RwLock::new(BackendRuntimeConfig::from_config_files().map_err(to_napi_error)?),
pool: Mutex::new(None),
})
}
#[napi]
pub async fn start(&self) -> Result<()> {
self.start_inner().await.map_err(to_napi_error)
}
async fn start_inner(&self) -> RuntimeResult<()> {
let mut guard = self.pool.lock().await;
if guard.is_some() {
return Ok(());
}
let database_url = self.config()?.database_url;
let pool = PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(5))
.connect(&database_url)
.await
.map_err(|err| RuntimeError::database("BackendRuntime failed to connect postgres", err))?;
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(|err| RuntimeError::database("BackendRuntime postgres health check failed", err))?;
let config = self.config()?.with_db_overrides(&pool).await?;
self.update_config(config)?;
*guard = Some(pool);
Ok(())
}
#[napi]
pub async fn stop(&self) -> Result<()> {
let pool = self.pool.lock().await.take();
if let Some(pool) = pool {
pool.close().await;
}
Ok(())
}
#[napi]
pub async fn health(&self) -> Result<BackendRuntimeHealth> {
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
.map(|row| row.try_get::<i32, _>(0).unwrap_or(0) == 1)
.unwrap_or(false),
None => false,
};
Ok(BackendRuntimeHealth {
started: pool.is_some(),
database_connected,
})
}
#[napi]
pub async fn run_migrations(&self) -> Result<()> {
let pool = self.pool().await?;
migrate_runtime_tables(&pool).await.map_err(to_napi_error)
}
pub(crate) async fn pool(&self) -> RuntimeResult<PgPool> {
self
.pool
.lock()
.await
.as_ref()
.cloned()
.ok_or_else(|| RuntimeError::invalid_state("BackendRuntime must be started before using postgres operations"))
}
pub(crate) fn config(&self) -> RuntimeResult<BackendRuntimeConfig> {
self
.config
.read()
.map(|config| config.clone())
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))
}
fn update_config(&self, config: BackendRuntimeConfig) -> RuntimeResult<()> {
*self
.config
.write()
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))? = config;
Ok(())
}
}

View File

@@ -0,0 +1,40 @@
use super::{Result, auth_challenge_purpose, dto::RuntimeStateRows};
pub(super) async fn create(
rows: &RuntimeStateRows,
purpose: &str,
token: &str,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<bool> {
rows
.insert_payload_if_absent(
&auth_challenge_purpose(purpose),
token,
None,
payload,
ttl_ms,
"RuntimeState auth challenge create",
)
.await
}
pub(super) async fn get(rows: &RuntimeStateRows, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
rows
.active_payload(
&auth_challenge_purpose(purpose),
token,
"RuntimeState auth challenge get",
)
.await
}
pub(super) async fn consume(rows: &RuntimeStateRows, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
rows
.consume_payload(
&auth_challenge_purpose(purpose),
token,
"RuntimeState auth challenge consume",
)
.await
}

View File

@@ -0,0 +1,128 @@
use super::{
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, Result, RuntimeByokLocalLeaseRecord, RuntimeError,
dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows},
};
pub(super) async fn get(rows: &RuntimeStateRows, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
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<RuntimeByokLocalLeaseRecord> {
if ttl_ms <= 0 {
return Err(RuntimeError::invalid_input("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| RuntimeError::database("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| RuntimeError::database("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| RuntimeError::database("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<Option<RuntimeByokLocalLeaseRecord>> {
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<Option<RuntimeByokLocalLeaseRecord>> {
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<RuntimeByokLocalLeaseRecord> {
Ok(RuntimeByokLocalLeaseRecord {
lease_id: lease_id.to_string(),
payload: row.payload,
expires_at_ms: row.expires_at_ms,
})
}

View File

@@ -0,0 +1,455 @@
use sqlx::{PgPool, Row};
use super::{RuntimeError, RuntimeResult, token_hash};
type Result<T> = RuntimeResult<T>;
pub(super) struct RuntimeStatePayloadRow {
pub(super) payload: serde_json::Value,
pub(super) expires_at_ms: i64,
}
pub(super) struct RuntimeStateLockedRow {
pub(super) payload: serde_json::Value,
pub(super) attempts: i32,
pub(super) expires_at: chrono::DateTime<chrono::Utc>,
}
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,
}
impl RuntimeStateRows {
pub(super) fn new(pool: PgPool) -> Self {
Self { pool }
}
pub(super) fn pool(&self) -> &PgPool {
&self.pool
}
pub(super) async fn begin(&self, context: &str) -> Result<sqlx::Transaction<'_, sqlx::Postgres>> {
self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database(format!("{context} transaction failed"), err))
}
pub(super) async fn insert_payload(
&self,
purpose: &str,
token: &str,
lookup_key: Option<&str>,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(())
}
pub(super) async fn insert_payload_if_absent(
&self,
purpose: &str,
token: &str,
lookup_key: Option<&str>,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<bool> {
let inserted = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO NOTHING
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?
.rows_affected()
== 1;
Ok(inserted)
}
pub(super) async fn upsert_payload_reset_attempts(
&self,
purpose: &str,
token: &str,
lookup_key: &str,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, attempts, consumed_at, expires_at)
VALUES ($1, $2, $3, $4, 0, NULL, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO UPDATE
SET lookup_key = EXCLUDED.lookup_key,
payload = EXCLUDED.payload,
attempts = 0,
consumed_at = NULL,
expires_at = EXCLUDED.expires_at,
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(())
}
pub(super) async fn active_payload(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<serde_json::Value>> {
let row = sqlx::query(
r#"
SELECT payload
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(|row| row.get::<serde_json::Value, _>("payload")))
}
pub(super) async fn active_payload_with_expires(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(payload_row))
}
pub(super) async fn consume_payload(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<serde_json::Value>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
RETURNING payload
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(|row| row.get::<serde_json::Value, _>("payload")))
}
pub(super) async fn consume_payload_with_expires(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
RETURNING payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(payload_row))
}
pub(super) async fn active_payload_with_expires_for_update_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > clock_timestamp()
FOR UPDATE
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(payload_row))
}
pub(super) async fn unconsumed_row_for_update_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStateLockedRow>> {
let row = sqlx::query(
r#"
SELECT payload, attempts, expires_at
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
FOR UPDATE
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(row.map(|row| RuntimeStateLockedRow {
payload: row.get("payload"),
attempts: row.get("attempts"),
expires_at: row.get("expires_at"),
}))
}
pub(super) async fn insert_payload_returning_expires_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
input: RuntimeStateInsertPayload<'_>,
) -> Result<i64> {
let row = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.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| RuntimeError::database(input.context, err))?;
Ok(row.get::<i64, _>("expires_at_ms"))
}
pub(super) async fn upsert_expired_or_consumed_payload_returning_expires_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
input: RuntimeStateInsertPayload<'_>,
) -> Result<Option<i64>> {
let row = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO UPDATE
SET lookup_key = EXCLUDED.lookup_key,
payload = EXCLUDED.payload,
attempts = 0,
consumed_at = NULL,
expires_at = clock_timestamp() + ($5 * INTERVAL '1 millisecond')
WHERE runtime_states.consumed_at IS NOT NULL
OR runtime_states.expires_at <= clock_timestamp()
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(input.purpose)
.bind(token_hash(input.token))
.bind(input.lookup_key)
.bind(input.payload)
.bind(input.ttl_ms as f64)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database(input.context, err))?;
Ok(row.map(|row| row.get::<i64, _>("expires_at_ms")))
}
pub(super) async fn update_attempts_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
attempts: i32,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
UPDATE runtime_states
SET attempts = $3,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(attempts)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(())
}
pub(super) async fn delete_by_key_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<()> {
sqlx::query("DELETE FROM runtime_states WHERE purpose = $1 AND token_hash = $2")
.bind(purpose)
.bind(token_hash(token))
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(())
}
pub(super) async fn cleanup_expired_or_consumed(&self, limit: i64, context: &str) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_states
WHERE (purpose, token_hash) IN (
SELECT purpose, token_hash FROM runtime_states
WHERE expires_at <= CURRENT_TIMESTAMP
OR consumed_at IS NOT NULL
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(result.rows_affected() as i64)
}
pub(super) async fn cleanup_expired_by_purpose_prefix(
&self,
purpose_prefix: &str,
limit: i64,
context: &str,
) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_states
WHERE (purpose, token_hash) IN (
SELECT purpose, token_hash FROM runtime_states
WHERE purpose LIKE $1
AND expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $2
)
"#,
)
.bind(format!("{purpose_prefix}%"))
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| RuntimeError::database(context, err))?;
Ok(result.rows_affected() as i64)
}
}
fn payload_row(row: sqlx::postgres::PgRow) -> RuntimeStatePayloadRow {
RuntimeStatePayloadRow {
payload: row.get("payload"),
expires_at_ms: row.get("expires_at_ms"),
}
}

View File

@@ -0,0 +1,184 @@
use super::{
Result, RuntimeError, RuntimeWorkspaceInviteLinkRecord, WORKSPACE_INVITE_LINK_ID_PURPOSE,
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows},
};
pub(super) async fn get_by_workspace(
rows: &RuntimeStateRows,
workspace_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
get_by_key(rows, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await
}
pub(super) async fn get_by_invite_id(
rows: &RuntimeStateRows,
invite_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
get_by_key(rows, WORKSPACE_INVITE_LINK_ID_PURPOSE, &invite_id).await
}
pub(super) async fn create(
rows: &RuntimeStateRows,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> Result<RuntimeWorkspaceInviteLinkRecord> {
if ttl_ms <= 0 {
return Err(RuntimeError::invalid_input(
"workspace invite link ttl must be positive",
));
}
let mut tx = rows.begin("RuntimeState workspace invite link").await?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&workspace_id)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link active lock failed", err))?;
if let Some(existing) =
get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?
{
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link transaction commit failed", err))?;
return Ok(existing);
}
let payload = serde_json::json!({
"workspaceId": workspace_id,
"inviteId": invite_id,
"inviterUserId": inviter_user_id,
});
let Some(expires_at_ms) = rows
.upsert_expired_or_consumed_payload_returning_expires_in_tx(
&mut tx,
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?
else {
let existing = get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link transaction commit failed", err))?;
return existing
.ok_or_else(|| RuntimeError::invalid_state("RuntimeState workspace invite link active conflict missing row"));
};
rows
.insert_payload_returning_expires_in_tx(
&mut tx,
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?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link transaction commit failed", err))?;
Ok(RuntimeWorkspaceInviteLinkRecord {
workspace_id,
invite_id,
inviter_user_id,
expires_at_ms,
})
}
pub(super) async fn revoke(rows: &RuntimeStateRows, workspace_id: String) -> Result<bool> {
let mut tx = rows.begin("RuntimeState workspace invite link").await?;
let existing = get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?;
let Some(existing) = existing else {
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link transaction commit failed", err))?;
return Ok(false);
};
rows
.delete_by_key_in_tx(
&mut tx,
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
&workspace_id,
"RuntimeState workspace invite link revoke",
)
.await?;
rows
.delete_by_key_in_tx(
&mut tx,
WORKSPACE_INVITE_LINK_ID_PURPOSE,
&existing.invite_id,
"RuntimeState workspace invite link revoke",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState workspace invite link transaction commit failed", err))?;
Ok(true)
}
async fn get_by_key(
rows: &RuntimeStateRows,
purpose: &str,
key: &str,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
rows
.active_payload_with_expires(purpose, key, "RuntimeState workspace invite link get")
.await?
.map(record_from_row)
.transpose()
}
async fn get_by_key_in_tx(
rows: &RuntimeStateRows,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
key: &str,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
rows
.active_payload_with_expires_for_update_in_tx(tx, purpose, key, "RuntimeState workspace invite link get")
.await?
.map(record_from_row)
.transpose()
}
fn record_from_row(row: RuntimeStatePayloadRow) -> Result<RuntimeWorkspaceInviteLinkRecord> {
Ok(RuntimeWorkspaceInviteLinkRecord {
workspace_id: row
.payload
.get("workspaceId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("RuntimeState workspace invite link payload missing workspaceId"))?
.to_string(),
invite_id: row
.payload
.get("inviteId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("RuntimeState workspace invite link payload missing inviteId"))?
.to_string(),
inviter_user_id: row
.payload
.get("inviterUserId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("RuntimeState workspace invite link payload missing inviterUserId"))?
.to_string(),
expires_at_ms: row.expires_at_ms,
})
}

View File

@@ -0,0 +1,172 @@
use super::{
MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS, Result, RuntimeError, RuntimeMagicLinkOtpConsumeResult,
dto::RuntimeStateRows,
};
impl RuntimeMagicLinkOtpConsumeResult {
fn ok(token: String) -> Self {
Self {
ok: true,
token: Some(token),
reason: None,
}
}
fn fail(reason: &'static str) -> Self {
Self {
ok: false,
token: None,
reason: Some(reason.to_string()),
}
}
}
pub(super) async fn upsert(
rows: &RuntimeStateRows,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> Result<()> {
if ttl_ms <= 0 {
return Err(RuntimeError::invalid_input("magic link otp ttl must be positive"));
}
let payload = serde_json::json!({
"otpHash": otp_hash,
"token": token,
"clientNonce": client_nonce,
});
rows
.upsert_payload_reset_attempts(
MAGIC_LINK_OTP_PURPOSE,
&email,
&email,
payload,
ttl_ms,
"RuntimeState magic link otp upsert",
)
.await
}
pub(super) async fn consume(
rows: &RuntimeStateRows,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> Result<RuntimeMagicLinkOtpConsumeResult> {
let mut tx = rows.begin("RuntimeState magic link otp").await?;
let row = rows
.unconsumed_row_for_update_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp lookup",
)
.await?;
let Some(row) = row else {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction rollback failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("not_found"));
};
let payload = row.payload;
let attempts = row.attempts;
let expires_at = row.expires_at;
if expires_at <= chrono::Utc::now() {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("expired"));
}
let stored_client_nonce = payload.get("clientNonce").and_then(serde_json::Value::as_str);
if stored_client_nonce.is_some() && stored_client_nonce != client_nonce.as_deref() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("nonce_mismatch"));
}
if attempts >= MAX_MAGIC_LINK_OTP_ATTEMPTS {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("locked"));
}
let stored_otp_hash = payload.get("otpHash").and_then(serde_json::Value::as_str);
if stored_otp_hash != Some(otp_hash.as_str()) {
let attempts = attempts + 1;
if attempts >= MAX_MAGIC_LINK_OTP_ATTEMPTS {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("locked"));
}
rows
.update_attempts_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
attempts,
"RuntimeState magic link otp attempts update",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("invalid_otp"));
}
let token = payload
.get("token")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("RuntimeState magic link otp payload missing token"))?
.to_string();
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState magic link otp transaction commit failed", err))?;
Ok(RuntimeMagicLinkOtpConsumeResult::ok(token))
}

View File

@@ -0,0 +1,258 @@
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error};
pub(super) use super::{
constants::{
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS,
WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
},
token_hash,
types::{
RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord,
},
};
mod auth_challenge;
mod byok_local_lease;
mod dto;
mod invite_link;
mod magic_link_otp;
mod store;
mod verification_token;
use store::RuntimeStateStore;
pub(super) type Result<T> = RuntimeResult<T>;
pub(super) fn auth_challenge_purpose(purpose: &str) -> String {
format!("auth_challenge:{purpose}")
}
pub(super) fn verification_token_purpose(token_type: i32) -> String {
format!("verification_token:{token_type}")
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn create_auth_challenge(
&self,
purpose: String,
token: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> napi::Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("auth challenge ttl must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.create_auth_challenge(&purpose, &token, payload, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_auth_challenge(&self, purpose: String, token: String) -> napi::Result<Option<serde_json::Value>> {
RuntimeStateStore::new(self.pool().await?)
.get_auth_challenge(&purpose, &token)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn consume_auth_challenge(
&self,
purpose: String,
token: String,
) -> napi::Result<Option<serde_json::Value>> {
RuntimeStateStore::new(self.pool().await?)
.consume_auth_challenge(&purpose, &token)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn create_verification_token(
&self,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> napi::Result<String> {
if ttl_ms <= 0 {
return Err(napi_error("verification token ttl must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.create_verification_token(token_type, credential, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_verification_token(
&self,
token_type: i32,
token: String,
keep: Option<bool>,
) -> napi::Result<Option<RuntimeVerificationTokenRecord>> {
let keep = keep.unwrap_or(false);
RuntimeStateStore::new(self.pool().await?)
.get_verification_token(token_type, token, keep)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn verify_verification_token(
&self,
token_type: i32,
token: String,
credential: Option<String>,
keep: Option<bool>,
) -> napi::Result<Option<RuntimeVerificationTokenRecord>> {
let keep = keep.unwrap_or(false);
RuntimeStateStore::new(self.pool().await?)
.verify_verification_token(token_type, token, credential, keep)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn cleanup_expired_verification_tokens(&self, limit: i64) -> napi::Result<i64> {
if limit <= 0 {
return Err(napi_error("verification token cleanup limit must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.cleanup_expired_verification_tokens(limit)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn upsert_magic_link_otp(
&self,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> napi::Result<()> {
RuntimeStateStore::new(self.pool().await?)
.upsert_magic_link_otp(email, otp_hash, token, client_nonce, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn consume_magic_link_otp(
&self,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> napi::Result<RuntimeMagicLinkOtpConsumeResult> {
RuntimeStateStore::new(self.pool().await?)
.consume_magic_link_otp(email, otp_hash, client_nonce)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn create_workspace_invite_link(
&self,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceInviteLinkRecord> {
RuntimeStateStore::new(self.pool().await?)
.create_workspace_invite_link(workspace_id, invite_id, inviter_user_id, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_workspace_invite_link(
&self,
workspace_id: String,
) -> napi::Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_workspace_invite_link(workspace_id)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_workspace_invite_link_by_id(
&self,
invite_id: String,
) -> napi::Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_workspace_invite_link_by_id(invite_id)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn revoke_workspace_invite_link(&self, workspace_id: String) -> napi::Result<bool> {
RuntimeStateStore::new(self.pool().await?)
.revoke_workspace_invite_link(workspace_id)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn create_byok_local_lease(
&self,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> napi::Result<RuntimeByokLocalLeaseRecord> {
RuntimeStateStore::new(self.pool().await?)
.create_byok_local_lease(active_key, lease_id, payload, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_byok_local_lease(&self, lease_id: String) -> napi::Result<Option<RuntimeByokLocalLeaseRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_byok_local_lease(lease_id)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn cleanup_expired_runtime_states(&self, limit: i64) -> napi::Result<i64> {
if limit <= 0 {
return Err(napi_error("runtime state cleanup limit must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.cleanup_expired_runtime_states(limit)
.await
.map_err(napi::Error::from)
}
}
#[cfg(test)]
mod tests {
use super::{
MAGIC_LINK_OTP_PURPOSE, WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, token_hash,
};
#[test]
fn magic_link_otp_uses_scoped_purpose_and_email_hash() {
assert_eq!(MAGIC_LINK_OTP_PURPOSE, "magic_link_otp");
assert_ne!(token_hash("user@affine.test"), "user@affine.test");
assert_eq!(token_hash("user@affine.test"), token_hash("user@affine.test"));
assert_ne!(token_hash("user@affine.test"), token_hash("other@affine.test"));
}
#[test]
fn workspace_invite_link_uses_scoped_purposes_and_hashes() {
assert_eq!(
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
"workspace_invite_link:workspace"
);
assert_eq!(WORKSPACE_INVITE_LINK_ID_PURPOSE, "workspace_invite_link:id");
assert_ne!(token_hash("workspace-id"), "workspace-id");
assert_ne!(token_hash("invite-id"), "invite-id");
}
}

View File

@@ -0,0 +1,138 @@
use sqlx::PgPool;
use super::{
Result, RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord, auth_challenge, byok_local_lease, dto::RuntimeStateRows, invite_link,
magic_link_otp, verification_token,
};
pub(super) struct RuntimeStateStore {
rows: RuntimeStateRows,
}
impl RuntimeStateStore {
pub(super) fn new(pool: PgPool) -> Self {
Self {
rows: RuntimeStateRows::new(pool),
}
}
pub(super) async fn create_auth_challenge(
&self,
purpose: &str,
token: &str,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<bool> {
auth_challenge::create(&self.rows, purpose, token, payload, ttl_ms).await
}
pub(super) async fn get_auth_challenge(&self, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
auth_challenge::get(&self.rows, purpose, token).await
}
pub(super) async fn consume_auth_challenge(&self, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
auth_challenge::consume(&self.rows, purpose, token).await
}
pub(super) async fn create_verification_token(
&self,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> Result<String> {
verification_token::create(&self.rows, token_type, credential, ttl_ms).await
}
pub(super) async fn get_verification_token(
&self,
token_type: i32,
token: String,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
verification_token::get(&self.rows, token_type, token, keep).await
}
pub(super) async fn verify_verification_token(
&self,
token_type: i32,
token: String,
credential: Option<String>,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
verification_token::verify(&self.rows, token_type, token, credential, keep).await
}
pub(super) async fn cleanup_expired_verification_tokens(&self, limit: i64) -> Result<i64> {
verification_token::cleanup_expired(&self.rows, limit).await
}
pub(super) async fn cleanup_expired_runtime_states(&self, limit: i64) -> Result<i64> {
self
.rows
.cleanup_expired_or_consumed(limit, "RuntimeState cleanup")
.await
}
pub(super) async fn upsert_magic_link_otp(
&self,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> Result<()> {
magic_link_otp::upsert(&self.rows, email, otp_hash, token, client_nonce, ttl_ms).await
}
pub(super) async fn consume_magic_link_otp(
&self,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> Result<RuntimeMagicLinkOtpConsumeResult> {
magic_link_otp::consume(&self.rows, email, otp_hash, client_nonce).await
}
pub(super) async fn create_workspace_invite_link(
&self,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> Result<RuntimeWorkspaceInviteLinkRecord> {
invite_link::create(&self.rows, workspace_id, invite_id, inviter_user_id, ttl_ms).await
}
pub(super) async fn get_workspace_invite_link(
&self,
workspace_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
invite_link::get_by_workspace(&self.rows, workspace_id).await
}
pub(super) async fn get_workspace_invite_link_by_id(
&self,
invite_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
invite_link::get_by_invite_id(&self.rows, invite_id).await
}
pub(super) async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result<bool> {
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<RuntimeByokLocalLeaseRecord> {
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<Option<RuntimeByokLocalLeaseRecord>> {
byok_local_lease::get(&self.rows, lease_id).await
}
}

View File

@@ -0,0 +1,149 @@
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{
Result, RuntimeError, RuntimeVerificationTokenRecord,
dto::{RuntimeStatePayloadRow, RuntimeStateRows},
token_hash, verification_token_purpose,
};
pub(super) async fn create(
rows: &RuntimeStateRows,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> Result<String> {
let token = Uuid::new_v4().to_string();
let payload = serde_json::json!({ "credential": credential });
rows
.insert_payload(
&verification_token_purpose(token_type),
&token,
credential.as_deref(),
payload,
ttl_ms,
"RuntimeState verification token create",
)
.await?;
Ok(token)
}
pub(super) async fn get(
rows: &RuntimeStateRows,
token_type: i32,
token: String,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let purpose = verification_token_purpose(token_type);
let row = if keep {
rows
.active_payload_with_expires(&purpose, &token, "RuntimeState verification token get")
.await?
} else {
rows
.consume_payload_with_expires(&purpose, &token, "RuntimeState verification token get")
.await?
};
Ok(row.map(|row| record_from_row(token_type, token, row)))
}
pub(super) async fn verify(
rows: &RuntimeStateRows,
token_type: i32,
token: String,
credential: Option<String>,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let purpose = verification_token_purpose(token_type);
let row = if keep {
active_payload_with_credential(rows.pool(), &purpose, &token, credential.as_deref()).await
} else {
consume_payload_with_credential(rows.pool(), &purpose, &token, credential.as_deref()).await
}
.map_err(|err| RuntimeError::database("RuntimeState verification token verify failed", err))?;
Ok(row.map(|row| record_from_row(token_type, token, row)))
}
pub(super) async fn cleanup_expired(rows: &RuntimeStateRows, limit: i64) -> Result<i64> {
rows
.cleanup_expired_by_purpose_prefix("verification_token:", limit, "RuntimeState verification token cleanup")
.await
}
async fn active_payload_with_credential(
pool: &PgPool,
purpose: &str,
token: &str,
credential: Option<&str>,
) -> sqlx::Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
AND (payload->>'credential' IS NULL OR payload->>'credential' = $3)
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(credential)
.fetch_optional(pool)
.await?;
Ok(row.map(payload_row))
}
async fn consume_payload_with_credential(
pool: &PgPool,
purpose: &str,
token: &str,
credential: Option<&str>,
) -> sqlx::Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
AND (payload->>'credential' IS NULL OR payload->>'credential' = $3)
RETURNING payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(credential)
.fetch_optional(pool)
.await?;
Ok(row.map(payload_row))
}
fn payload_row(row: sqlx::postgres::PgRow) -> RuntimeStatePayloadRow {
RuntimeStatePayloadRow {
payload: row.get("payload"),
expires_at_ms: row.get("expires_at_ms"),
}
}
fn record_from_row(token_type: i32, token: String, row: RuntimeStatePayloadRow) -> RuntimeVerificationTokenRecord {
RuntimeVerificationTokenRecord {
token_type,
token,
credential: row
.payload
.get("credential")
.and_then(serde_json::Value::as_str)
.map(ToString::to_string),
expires_at_ms: row.expires_at_ms,
}
}

View File

@@ -0,0 +1,85 @@
WITH targets AS (
SELECT UNNEST($1::varchar[]) AS workspace_id
),
snapshot_stats AS (
SELECT workspace_id,
COUNT(*) AS snapshot_count,
COALESCE(SUM(COALESCE(size, octet_length(blob))), 0) AS snapshot_size
FROM snapshots
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
blob_stats AS (
SELECT workspace_id,
COUNT(*) FILTER (WHERE deleted_at IS NULL AND status = 'completed') AS blob_count,
COALESCE(SUM(size) FILTER (WHERE deleted_at IS NULL AND status = 'completed'), 0) AS blob_size
FROM blobs
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
member_stats AS (
SELECT workspace_id, COUNT(*) AS member_count
FROM workspace_user_permissions
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
public_page_stats AS (
SELECT workspace_id, COUNT(*) AS public_page_count
FROM workspace_pages
WHERE public = TRUE AND workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
feature_stats AS (
SELECT workspace_id,
ARRAY_AGG(DISTINCT name ORDER BY name) FILTER (WHERE activated) AS features
FROM workspace_features
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
aggregated AS (
SELECT t.workspace_id,
COALESCE(ss.snapshot_count, 0) AS snapshot_count,
COALESCE(ss.snapshot_size, 0) AS snapshot_size,
COALESCE(bs.blob_count, 0) AS blob_count,
COALESCE(bs.blob_size, 0) AS blob_size,
COALESCE(ms.member_count, 0) AS member_count,
COALESCE(pp.public_page_count, 0) AS public_page_count,
COALESCE(fs.features, ARRAY[]::text[]) AS features
FROM targets t
LEFT JOIN snapshot_stats ss ON ss.workspace_id = t.workspace_id
LEFT JOIN blob_stats bs ON bs.workspace_id = t.workspace_id
LEFT JOIN member_stats ms ON ms.workspace_id = t.workspace_id
LEFT JOIN public_page_stats pp ON pp.workspace_id = t.workspace_id
LEFT JOIN feature_stats fs ON fs.workspace_id = t.workspace_id
)
INSERT INTO workspace_admin_stats (
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
updated_at
)
SELECT
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
NOW()
FROM aggregated
ON CONFLICT (workspace_id) DO UPDATE SET
snapshot_count = EXCLUDED.snapshot_count,
snapshot_size = EXCLUDED.snapshot_size,
blob_count = EXCLUDED.blob_count,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
public_page_count = EXCLUDED.public_page_count,
features = EXCLUDED.features,
updated_at = EXCLUDED.updated_at

View File

@@ -0,0 +1,436 @@
use anyhow::{Context, Result as AnyResult, anyhow};
use super::{
super::migrations::{RUNTIME_MIGRATIONS, migrate_runtime_tables},
runtime_state::*,
*,
};
static PG_TEST_LOCK: std::sync::OnceLock<tokio::sync::Mutex<()>> = 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(()))
}
#[test]
fn migrations_include_runtime_tables_without_worker_heartbeats() {
assert!(RUNTIME_MIGRATIONS.contains("runtime_states"));
assert!(RUNTIME_MIGRATIONS.contains("runtime_gates"));
assert!(RUNTIME_MIGRATIONS.contains("runtime_leases"));
assert!(RUNTIME_MIGRATIONS.contains("blob_reconciliation_runs"));
assert!(RUNTIME_MIGRATIONS.contains("blob_reconciliation_checkpoints"));
assert!(RUNTIME_MIGRATIONS.contains("doc_blob_refs"));
assert!(RUNTIME_MIGRATIONS.contains("blob_cleanup_candidates"));
assert!(!RUNTIME_MIGRATIONS.contains("runtime_worker_heartbeats"));
}
#[test]
fn auth_challenge_state_uses_scoped_purpose_and_token_hash() {
assert_eq!(auth_challenge_purpose("oauth_state"), "auth_challenge:oauth_state");
assert_ne!(token_hash("plain-token"), "plain-token");
assert_eq!(token_hash("plain-token"), token_hash("plain-token"));
assert_ne!(token_hash("plain-token"), token_hash("other-token"));
}
#[test]
fn verification_token_state_uses_typed_purpose_and_token_hash() {
assert_eq!(verification_token_purpose(0), "verification_token:0");
assert_ne!(token_hash("verification-token"), "verification-token");
assert_eq!(token_hash("verification-token"), token_hash("verification-token"));
assert_ne!(token_hash("verification-token"), token_hash("other-token"));
}
async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return Ok(None);
};
let pool = PgPoolOptions::new()
.max_connections(5)
.connect(&database_url)
.await
.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 = 'verification_token:99999'
"#,
)
.execute(&pool)
.await
.context("cleanup runtime_states for backend runtime tests")?;
sqlx::query("DELETE FROM runtime_gates WHERE key LIKE 'rust-test:%'")
.execute(&pool)
.await
.context("cleanup runtime_gates for backend runtime tests")?;
sqlx::query("DELETE FROM runtime_leases WHERE key LIKE 'rust-test:%'")
.execute(&pool)
.await
.context("cleanup runtime_leases for backend runtime tests")?;
Ok(Some(BackendRuntime {
config: std::sync::RwLock::new(BackendRuntimeConfig { database_url }),
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.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
struct Case {
key: &'static str,
first_ttl_ms: i64,
wait_ms: Option<u64>,
second_expected: bool,
}
for case in [
Case {
key: "rust-test:gate:same-key",
first_ttl_ms: 30_000,
wait_ms: None,
second_expected: false,
},
Case {
key: "rust-test:gate:expired-key",
first_ttl_ms: 1,
wait_ms: Some(20),
second_expected: true,
},
] {
assert!(
runtime
.put_runtime_gate_if_absent(case.key.to_string(), case.first_ttl_ms)
.await
.unwrap()
);
if let Some(wait_ms) = case.wait_ms {
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
}
assert_eq!(
runtime
.put_runtime_gate_if_absent(case.key.to_string(), 30_000)
.await
.unwrap(),
case.second_expected,
"{}",
case.key
);
}
let mut tasks = Vec::new();
for _ in 0..16 {
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
tasks.push(tokio::spawn(async move {
runtime
.put_runtime_gate_if_absent("rust-test:gate:concurrent".to_string(), 30_000)
.await
.unwrap()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
assert!(
runtime
.put_runtime_gate_if_absent("rust-test:gate:cleanup".to_string(), 1)
.await
.unwrap()
);
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(runtime.cleanup_expired_runtime_gates(100).await.unwrap(), 1);
assert_eq!(runtime.cleanup_expired_runtime_gates(100).await.unwrap(), 0);
}
#[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.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
let lease = runtime
.acquire_coordination_lease("rust-test:lease:basic".to_string(), "owner-1".to_string(), 30_000)
.await
.unwrap()
.expect("first owner should acquire lease");
assert_eq!(lease.fencing_token, 1);
assert!(
!runtime
.release_coordination_lease(lease.key.clone(), "owner-2".to_string(), lease.fencing_token)
.await
.unwrap()
);
assert!(
runtime
.release_coordination_lease(lease.key.clone(), lease.owner.clone(), lease.fencing_token)
.await
.unwrap()
);
let mut tasks = Vec::new();
for index in 0..16 {
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
tasks.push(tokio::spawn(async move {
runtime
.acquire_coordination_lease(
"rust-test:lease:concurrent".to_string(),
format!("owner-{index}"),
30_000,
)
.await
.unwrap()
.is_some()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
let stale = runtime
.acquire_coordination_lease("rust-test:lease:stale".to_string(), "owner-1".to_string(), 1)
.await
.unwrap()
.expect("stale lease owner should acquire");
tokio::time::sleep(Duration::from_millis(20)).await;
let takeover = runtime
.acquire_coordination_lease("rust-test:lease:stale".to_string(), "owner-2".to_string(), 30_000)
.await
.unwrap()
.expect("expired lease should be taken over");
assert_eq!(takeover.fencing_token, stale.fencing_token + 1);
assert!(
!runtime
.release_coordination_lease(stale.key.clone(), stale.owner.clone(), stale.fencing_token)
.await
.unwrap()
);
let renew = runtime
.acquire_coordination_lease("rust-test:lease:renew".to_string(), "owner-1".to_string(), 30_000)
.await
.unwrap()
.expect("renew lease owner should acquire");
assert!(
!runtime
.renew_coordination_lease(renew.key.clone(), "owner-2".to_string(), renew.fencing_token, 30_000)
.await
.unwrap()
);
assert!(
!runtime
.renew_coordination_lease(renew.key.clone(), renew.owner.clone(), renew.fencing_token + 1, 30_000)
.await
.unwrap()
);
assert!(
runtime
.renew_coordination_lease(renew.key.clone(), renew.owner.clone(), renew.fencing_token, 30_000)
.await
.unwrap()
);
}
#[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.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
assert!(
runtime
.create_auth_challenge(
"rust_test:cleanup".to_string(),
"expired".to_string(),
serde_json::json!({}),
1
)
.await
.unwrap()
);
assert!(
runtime
.create_auth_challenge(
"rust_test:cleanup".to_string(),
"consumed".to_string(),
serde_json::json!({}),
30_000,
)
.await
.unwrap()
);
assert!(
runtime
.consume_auth_challenge("rust_test:cleanup".to_string(), "consumed".to_string())
.await
.unwrap()
.is_some()
);
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(runtime.cleanup_expired_runtime_states(100).await.unwrap(), 2);
assert_eq!(runtime.cleanup_expired_runtime_states(100).await.unwrap(), 0);
}
#[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.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
let mismatch_token = runtime
.create_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
Some("user@affine.test".to_string()),
30_000,
)
.await
.unwrap();
assert!(
runtime
.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(
TEST_VERIFICATION_TOKEN_TYPE,
mismatch_token.clone(),
Some("user@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_some()
);
assert!(
runtime
.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(
TEST_VERIFICATION_TOKEN_TYPE,
Some("keep@affine.test".to_string()),
30_000,
)
.await
.unwrap();
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), Some(true))
.await
.unwrap()
.is_some()
);
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None)
.await
.unwrap()
.is_some()
);
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None)
.await
.unwrap()
.is_none()
);
let concurrent_token = runtime
.create_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
Some("concurrent@affine.test".to_string()),
30_000,
)
.await
.unwrap();
let mut tasks = Vec::new();
for _ in 0..16 {
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
let token = concurrent_token.clone();
tasks.push(tokio::spawn(async move {
runtime
.verify_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
token,
Some("concurrent@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_some()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
let expired_token = runtime
.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(TEST_VERIFICATION_TOKEN_TYPE, expired_token.clone(), None)
.await
.unwrap()
.is_none()
);
assert_eq!(runtime.cleanup_expired_verification_tokens(100).await.unwrap(), 1);
assert_eq!(runtime.cleanup_expired_verification_tokens(100).await.unwrap(), 0);
}

View File

@@ -0,0 +1,530 @@
use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use tokio::time::{Duration as TokioDuration, sleep};
use super::{
BackendRuntime, RuntimeError, RuntimeResult,
constants::{WORKSPACE_STATS_LEASE_KEY, WORKSPACE_STATS_LOCK_NAMESPACE, WORKSPACE_STATS_REFRESH_LOCK_KEY},
napi_error,
types::{
CoordinationLeaseGrant, RuntimeWorkspaceStatsDailyRecalibrationResult, RuntimeWorkspaceStatsRecalibrationResult,
RuntimeWorkspaceStatsRefreshResult, RuntimeWorkspaceStatsSnapshotResult,
},
};
const UPSERT_WORKSPACE_ADMIN_STATS_SQL: &str = include_str!("sql/upsert_workspace_admin_stats.sql");
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn refresh_workspace_admin_stats_dirty(
&self,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsRefreshResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats dirty refresh limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.refresh_dirty(batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
#[napi]
pub async fn recalibrate_workspace_admin_stats(
&self,
last_sid: i64,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsRecalibrationResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats recalibration limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.recalibrate(last_sid, batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
#[napi]
pub async fn write_workspace_admin_stats_daily_snapshot(
&self,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsSnapshotResult> {
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.write_daily_snapshot()
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(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,
) -> napi::Result<RuntimeWorkspaceStatsDailyRecalibrationResult> {
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: RuntimeResult<RuntimeWorkspaceStatsDailyRecalibrationResult> = 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?;
Ok(result?)
}
}
#[derive(FromRow)]
struct WorkspaceSid {
id: String,
sid: i32,
}
struct WorkspaceStatsStore {
pool: PgPool,
}
impl WorkspaceStatsStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn refresh_dirty(&self, batch_limit: i64) -> RuntimeResult<RuntimeWorkspaceStatsRefreshResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
}
let backlog = count_dirty(&mut tx).await?;
let dirty = load_dirty(&mut tx, batch_limit).await?;
if dirty.is_empty() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog,
skipped: false,
});
}
upsert_stats(&mut tx, &dirty).await?;
clear_dirty(&mut tx, &dirty).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
Ok(RuntimeWorkspaceStatsRefreshResult {
processed: dirty.len() as i64,
backlog,
skipped: false,
})
}
async fn recalibrate(
&self,
last_sid: i64,
batch_limit: i64,
) -> RuntimeResult<RuntimeWorkspaceStatsRecalibrationResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
}
let workspaces = fetch_workspace_batch(&mut tx, last_sid, batch_limit).await?;
if workspaces.is_empty() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: false,
});
}
let ids = workspaces
.iter()
.map(|workspace| workspace.id.clone())
.collect::<Vec<_>>();
let next_sid = workspaces
.last()
.map(|workspace| workspace.sid as i64)
.unwrap_or(last_sid);
upsert_stats(&mut tx, &ids).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: ids.len() as i64,
last_sid: next_sid,
skipped: false,
})
}
async fn write_daily_snapshot(&self) -> RuntimeResult<RuntimeWorkspaceStatsSnapshotResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot commit failed", err))?;
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
}
let snapshotted = write_daily_snapshot(&mut tx).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot commit failed", err))?;
Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted,
skipped: false,
})
}
}
async fn release_workspace_stats_lease(runtime: &BackendRuntime, lease: CoordinationLeaseGrant) -> RuntimeResult<()> {
let _ = runtime
.release_coordination_lease_inner(lease.key, lease.owner, lease.fencing_token)
.await?;
Ok(())
}
async fn acquire_workspace_stats_lease_with_retry(
runtime: &BackendRuntime,
owner: String,
lease_ttl_ms: i64,
retry_times: i64,
retry_delay_ms: i64,
) -> RuntimeResult<Option<CoordinationLeaseGrant>> {
for attempt in 0..retry_times {
let lease = runtime
.acquire_coordination_lease_inner(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<T, F, Fut>(
retry_times: i64,
retry_delay_ms: i64,
mut operation: F,
) -> RuntimeResult<T>
where
T: WorkspaceStatsSkippable,
F: FnMut() -> Fut,
Fut: std::future::Future<Output = RuntimeResult<T>>,
{
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>) -> RuntimeResult<bool> {
let row = sqlx::query(
r#"
SELECT pg_try_advisory_xact_lock(($1::bigint << 32) + $2::bigint) AS locked
"#,
)
.bind(WORKSPACE_STATS_LOCK_NAMESPACE)
.bind(WORKSPACE_STATS_REFRESH_LOCK_KEY)
.fetch_one(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats transaction lock failed", err))?;
Ok(row.get::<bool, _>("locked"))
}
async fn load_dirty(tx: &mut Transaction<'_, Postgres>, limit: i64) -> RuntimeResult<Vec<String>> {
let rows = sqlx::query(
r#"
SELECT workspace_id
FROM workspace_admin_stats_dirty
ORDER BY updated_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
"#,
)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats load dirty workspaces failed", err))?;
Ok(rows.into_iter().map(|row| row.get("workspace_id")).collect())
}
async fn count_dirty(tx: &mut Transaction<'_, Postgres>) -> RuntimeResult<i64> {
let row = sqlx::query("SELECT COUNT(*) AS total FROM workspace_admin_stats_dirty")
.fetch_one(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats count dirty workspaces failed", err))?;
Ok(row.get::<i64, _>("total"))
}
async fn clear_dirty(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> RuntimeResult<()> {
sqlx::query(
r#"
DELETE FROM workspace_admin_stats_dirty
WHERE workspace_id = ANY($1::varchar[])
"#,
)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats clear dirty workspaces failed", err))?;
Ok(())
}
async fn upsert_stats(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> RuntimeResult<()> {
if workspace_ids.is_empty() {
return Ok(());
}
sqlx::query(UPSERT_WORKSPACE_ADMIN_STATS_SQL)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats upsert stats failed", err))?;
Ok(())
}
async fn fetch_workspace_batch(
tx: &mut Transaction<'_, Postgres>,
last_sid: i64,
limit: i64,
) -> RuntimeResult<Vec<WorkspaceSid>> {
sqlx::query_as::<_, WorkspaceSid>(
r#"
SELECT id, sid
FROM workspaces
WHERE sid > $1
ORDER BY sid
LIMIT $2
"#,
)
.bind(last_sid)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats fetch workspace batch failed", err))
}
async fn write_daily_snapshot(tx: &mut Transaction<'_, Postgres>) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
INSERT INTO workspace_admin_stats_daily (
workspace_id,
date,
snapshot_size,
blob_size,
member_count,
updated_at
)
SELECT
workspace_id,
CURRENT_DATE,
snapshot_size,
blob_size,
member_count,
NOW()
FROM workspace_admin_stats
ON CONFLICT (workspace_id, date)
DO UPDATE SET
snapshot_size = EXCLUDED.snapshot_size,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
updated_at = EXCLUDED.updated_at
"#,
)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot failed", err))?;
Ok(result.rows_affected() as i64)
}