feat(server): impl storage runtime

2026-07-12 23:56:36 +08:00 · 2026-07-01 03:04:14 +08:00
parent a821f67fc9
commit 0e20823311
90 changed files with 5439 additions and 4292 deletions
@@ -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);
+}