diff --git a/.gitignore b/.gitignore index d252afbc..2dced961 100644 --- a/.gitignore +++ b/.gitignore @@ -125,6 +125,9 @@ railway-*.md # Personal/internal planning notes — never include in repo plans/ .local-plans/ +# But allow Superpowers plans/specs that are part of the design record. +!docs/superpowers/plans/ +!docs/superpowers/specs/ # Benchmark archives + working analysis notes # Kept locally under @@ -137,3 +140,6 @@ claudedocs/ # Temp File temp/ + +# Superpowers brainstorming session files +.superpowers/ diff --git a/docs/relayer/observability.md b/docs/relayer/observability.md index 6bc7bcb1..8849ed40 100644 --- a/docs/relayer/observability.md +++ b/docs/relayer/observability.md @@ -66,6 +66,11 @@ Core relayer metrics: | `memwal_sidecar_failures_total` | `operation`, `reason` | Sidecar transport and HTTP failures | | `memwal_db_query_duration_seconds` | `operation`, `status` | PostgreSQL and pgvector query latency | | `memwal_db_pool_connections` | `state` | PostgreSQL pool `open` and `idle` gauges | +| `memwal_write_stream_permits_total` | none | Configured permit ceiling | +| `memwal_write_stream_permits_available` | none | Permits currently free | +| `memwal_write_stream_waiters_total` | none | Tasks waiting for a permit | +| `memwal_write_stream_acquired_total` | `result` | Permit acquisition outcomes | +| `memwal_write_stream_rejected_total` | `route` | Requests rejected with `429` | Example Prometheus scrape config: diff --git a/docs/relayer/self-hosting.md b/docs/relayer/self-hosting.md index d47c2781..58fc02e2 100644 --- a/docs/relayer/self-hosting.md +++ b/docs/relayer/self-hosting.md @@ -98,6 +98,15 @@ By default, the relayer enforces rate limits and storage quotas via Redis to pre - `RATE_LIMIT_STORAGE_BYTES` — max storage per user in bytes (default: 1 GB, `1073741824`) - `REDIS_URL` — required to track sliding windows for rate limits (default: `redis://localhost:6379`) +### Write-stream concurrency + +| Variable | Default | Description | +|---|---|---| +| `WRITE_STREAM_MAX_CONCURRENCY` | `8` | Maximum concurrent active write operations (prep + upload) | +| `WRITE_STREAM_ACQUIRE_TIMEOUT_MS` | `5000` | Handler wait for a write slot before returning `429` | + +Tune `WRITE_STREAM_MAX_CONCURRENCY` down to reduce sidecar pressure, or up to increase throughput when the sidecar has headroom. Keep it below `WALRUS_UPLOAD_MAX_CONCURRENCY` so the sidecar safety net remains meaningful. + ### Defaults - `PORT` defaults to `8000` diff --git a/docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md b/docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md new file mode 100644 index 00000000..71d94e8d --- /dev/null +++ b/docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md @@ -0,0 +1,1168 @@ +## Relayer write-stream redesign implementation plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Move the Walrus upload concurrency budget from the TypeScript sidecar into the Rust relayer so write-stream overload cannot build unbounded sidecar queues. + +**Architecture:** Add a `tokio::sync::Semaphore`-based `WriteStreamLimiter` to `AppState`. Handlers acquire a permit before starting prep work; wallet workers acquire a permit before calling `/walrus/upload`. The same permit pool caps total active writes. The sidecar keeps its existing upload limiter as a higher safety net. + +**Tech Stack:** Rust, Axum, Tokio, Apalis, Prometheus, PostgreSQL, TypeScript/Express sidecar. + +--- + +## File Structure + +| File | Responsibility | +|---|---| +| `services/server/src/services/write_stream.rs` (create) | `WriteStreamLimiter`, permit guard, acquisition error, unit tests. | +| `services/server/src/services/mod.rs` (modify) | Export the new module and its public types. | +| `services/server/src/types.rs` (modify) | Add `write_stream_limiter` to `AppState`; add config fields for `WRITE_STREAM_MAX_CONCURRENCY` and `WRITE_STREAM_ACQUIRE_TIMEOUT_MS`. | +| `services/server/src/main.rs` (modify) | Construct and inject `WriteStreamLimiter` into `AppState`. | +| `services/server/src/observability.rs` (modify) | Register new Prometheus metrics for permit state and acquisition outcomes. | +| `services/server/src/routes/mod.rs` (modify) | Add small shared helper to translate a permit acquisition timeout into a `429` `AppError`. | +| `services/server/src/routes/remember.rs` (modify) | Gate `/api/remember`, `/api/remember/bulk`, `/api/remember/manual` with the limiter. | +| `services/server/src/routes/analyze.rs` (modify) | Gate `/api/analyze` production and benchmark paths with the limiter. | +| `services/server/src/jobs.rs` (modify) | Gate `execute_wallet_job` upload path with the limiter. | + +--- + +### Task 1: Create `WriteStreamLimiter` + +**Files:** +- Create: `services/server/src/services/write_stream.rs` +- Modify: `services/server/src/services/mod.rs` + +- [ ] **Step 1: Write the failing unit-test file** + +Create `services/server/src/services/write_stream.rs` with tests first: + +```rust +//! Write-stream concurrency limiter. +//! +//! Owns the single in-process budget for active write operations +//! (prep + upload). Every memory item that will result in a sidecar +//! /walrus/upload call must acquire a permit before starting work and +//! release it when the active phase ends. + +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Semaphore; +use tokio::time::error::Elapsed; + +const DEFAULT_WRITE_STREAM_MAX_CONCURRENCY: usize = 8; +const MIN_WRITE_STREAM_MAX_CONCURRENCY: usize = 1; +const MAX_WRITE_STREAM_MAX_CONCURRENCY: usize = 100; + +#[derive(Debug)] +pub enum AcquireError { + Timeout, + Closed, +} + +impl std::fmt::Display for AcquireError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AcquireError::Timeout => write!(f, "write stream concurrency limit reached"), + AcquireError::Closed => write!(f, "write stream limiter closed"), + } + } +} + +impl std::error::Error for AcquireError {} + +/// Guard that releases one or more permits when dropped. +pub struct WriteStreamPermit { + semaphore: Arc, + permits: usize, +} + +impl Drop for WriteStreamPermit { + fn drop(&mut self) { + self.semaphore.add_permits(self.permits); + } +} + +/// In-process concurrency limiter for the write stream. +#[derive(Clone, Debug)] +pub struct WriteStreamLimiter { + semaphore: Arc, + max_permits: usize, +} + +impl WriteStreamLimiter { + pub fn new(max_permits: usize) -> Self { + let max_permits = max_permits + .max(MIN_WRITE_STREAM_MAX_CONCURRENCY) + .min(MAX_WRITE_STREAM_MAX_CONCURRENCY); + Self { + semaphore: Arc::new(Semaphore::new(max_permits)), + max_permits, + } + } + + pub fn default_limiter() -> Self { + Self::new(DEFAULT_WRITE_STREAM_MAX_CONCURRENCY) + } + + pub fn max_permits(&self) -> usize { + self.max_permits + } + + pub fn available_permits(&self) -> usize { + self.semaphore.available_permits() + } + + /// Acquire a single permit, waiting up to `timeout`. + pub async fn acquire(&self, timeout: Duration) -> Result { + self.acquire_many(1, timeout).await + } + + /// Acquire `n` permits atomically with respect to this call, waiting up to `timeout`. + /// The returned guard releases all `n` permits on drop. + pub async fn acquire_many( + &self, + n: usize, + timeout: Duration, + ) -> Result { + if n == 0 { + return Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: 0, + }); + } + let n = n.min(self.max_permits); + let permit = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) + .await + .map_err(|_: Elapsed| AcquireError::Timeout)? + .map_err(|_| AcquireError::Closed)?; + permit.forget(); + Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: n, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn acquire_returns_permit_when_available() { + let limiter = WriteStreamLimiter::new(1); + let permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + assert_eq!(limiter.available_permits(), 0); + drop(permit); + assert_eq!(limiter.available_permits(), 1); + } + + #[tokio::test] + async fn acquire_times_out_when_exhausted() { + let limiter = WriteStreamLimiter::new(1); + let _permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + let err = limiter + .acquire(Duration::from_millis(50)) + .await + .unwrap_err(); + assert!(matches!(err, AcquireError::Timeout)); + } + + #[tokio::test] + async fn acquire_many_returns_all_or_none() { + let limiter = WriteStreamLimiter::new(3); + let _p1 = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + // asking for 3 when only 2 are free should time out + let err = limiter + .acquire_many(3, Duration::from_millis(50)) + .await + .unwrap_err(); + assert!(matches!(err, AcquireError::Timeout)); + assert_eq!(limiter.available_permits(), 2); + let p23 = limiter.acquire_many(2, Duration::from_secs(1)).await.unwrap(); + assert_eq!(limiter.available_permits(), 0); + drop(p23); + assert_eq!(limiter.available_permits(), 2); + } + + #[tokio::test] + async fn zero_permits_noop() { + let limiter = WriteStreamLimiter::new(1); + let guard = limiter.acquire_many(0, Duration::from_secs(1)).await.unwrap(); + assert_eq!(guard.permits, 0); + assert_eq!(limiter.available_permits(), 1); + } + + #[tokio::test] + async fn clamps_out_of_range_values() { + let low = WriteStreamLimiter::new(0); + assert_eq!(low.max_permits(), MIN_WRITE_STREAM_MAX_CONCURRENCY); + let high = WriteStreamLimiter::new(10_000); + assert_eq!(high.max_permits(), MAX_WRITE_STREAM_MAX_CONCURRENCY); + } +} +``` + +- [ ] **Step 2: Run the tests to verify they compile and fail** + +Run: + +```bash +cd services/server +cargo test --lib services::write_stream -- --nocapture +``` + +Expected: compile succeeds, tests pass (this module is self-contained, so the first test run already passes). + +- [ ] **Step 3: Export the module from `services/mod.rs`** + +Modify `services/server/src/services/mod.rs`: + +```rust +pub mod embedder; +pub mod extractor; +pub mod llm_chat; +pub mod ranker; +pub mod write_stream; + +// Placeholder module — reserved namespace for the consolidator +pub mod consolidator; + +pub use embedder::{Embedder, OpenAiEmbedder}; +pub use extractor::{Extractor, LlmExtractor}; +pub use ranker::{CompositeRanker, Ranker}; +pub use write_stream::{WriteStreamLimiter, WriteStreamPermit, WriteStreamSnapshot}; +``` + +- [ ] **Step 4: Run the tests again** + +Run: + +```bash +cd services/server +cargo test --lib services::write_stream -- --nocapture +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add services/server/src/services/write_stream.rs services/server/src/services/mod.rs +git commit -m "feat(write-stream): add WriteStreamLimiter with unit tests" +``` + +--- + +### Task 2: Wire `WriteStreamLimiter` into `AppState` and `Config` + +**Files:** +- Modify: `services/server/src/types.rs` +- Modify: `services/server/src/main.rs` + +- [ ] **Step 1: Add config fields and parsing** + +In `services/server/src/types.rs`, add to the `Config` struct: + +```rust +/// Maximum concurrent active write-stream operations (prep + upload). +pub write_stream_max_concurrency: usize, +/// How long a handler waits for a write-stream permit before returning 429. +pub write_stream_acquire_timeout: std::time::Duration, +``` + +Add helper functions near the other env-parsing helpers: + +```rust +pub(crate) fn parse_write_stream_max_concurrency() -> usize { + std::env::var("WRITE_STREAM_MAX_CONCURRENCY") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .map(|v| v.clamp(1, 100)) + .unwrap_or(8) +} + +pub(crate) fn parse_write_stream_acquire_timeout() -> std::time::Duration { + let millis = std::env::var("WRITE_STREAM_ACQUIRE_TIMEOUT_MS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .map(|v| v.clamp(100, 60_000)) + .unwrap_or(5_000); + std::time::Duration::from_millis(millis) +} +``` + +In `Config::from_env()` (around line 300), add: + +```rust +write_stream_max_concurrency: parse_write_stream_max_concurrency(), +write_stream_acquire_timeout: parse_write_stream_acquire_timeout(), +``` + +- [ ] **Step 2: Add `write_stream_limiter` to `AppState`** + +In `services/server/src/types.rs`, add inside `AppState`: + +```rust +/// In-process concurrency limiter for write operations. +pub write_stream_limiter: Arc, +``` + +- [ ] **Step 3: Import `WriteStreamLimiter` in `types.rs`** + +At the top of `services/server/src/types.rs`, change: + +```rust +use crate::services::{Embedder, Extractor, Ranker}; +``` + +to: + +```rust +use crate::services::{Embedder, Extractor, Ranker, WriteStreamLimiter}; +``` + +- [ ] **Step 4: Construct the limiter in `main.rs`** + +In `services/server/src/main.rs`, after the `config` is wrapped in `Arc` (around line 440), add: + +```rust +let write_stream_limiter = Arc::new(WriteStreamLimiter::new( + config.write_stream_max_concurrency, +)); +tracing::info!( + " write stream limiter: max_concurrency={} acquire_timeout_ms={}", + write_stream_limiter.max_permits(), + config.write_stream_acquire_timeout.as_millis(), +); +``` + +Then add `write_stream_limiter` to the `AppState` initialization: + +```rust +let state = Arc::new(AppState { + db, + config: Arc::clone(&config), + http_client, + key_pool, + alerts, + engine, + embedder, + extractor, + ranker, + redis, + fallback_rate_limit: tokio::sync::Mutex::new(crate::rate_limit::InMemoryFallback::default()), + remember_job_storage: remember_job_storage.clone(), + wallet_storage: wallet_storage.clone(), + bulk_job_storage: bulk_job_storage.clone(), + blob_cache_ttl, + blob_cache_max_bytes, + embedding_cache_ttl, + write_stream_limiter, +}); +``` + +- [ ] **Step 5: Compile to verify wiring** + +Run: + +```bash +cd services/server +cargo check +``` + +Expected: no errors. + +- [ ] **Step 6: Commit** + +```bash +git add services/server/src/types.rs services/server/src/main.rs +git commit -m "feat(write-stream): wire WriteStreamLimiter into AppState and Config" +``` + +--- + +### Task 3: Add Prometheus metrics for the limiter + +**Files:** +- Modify: `services/server/src/observability.rs` + +- [ ] **Step 1: Register static metrics** + +Add near the other `LazyLock` metric definitions in `services/server/src/observability.rs`: + +```rust +static WRITE_STREAM_PERMITS_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_permits_total", + "Total configured write-stream permits." + ) + .expect("register memwal_write_stream_permits_total") +}); + +static WRITE_STREAM_PERMITS_AVAILABLE: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_permits_available", + "Currently available write-stream permits." + ) + .expect("register memwal_write_stream_permits_available") +}); + +static WRITE_STREAM_WAITERS_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_waiters_total", + "Tasks currently waiting for a write-stream permit." + ) + .expect("register memwal_write_stream_waiters_total") +}); + +static WRITE_STREAM_ACQUIRED_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "memwal_write_stream_acquired_total", + "Write-stream permit acquisition outcomes.", + &["result"] + ) + .expect("register memwal_write_stream_acquired_total") +}); + +static WRITE_STREAM_REJECTED_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "memwal_write_stream_rejected_total", + "Requests rejected because the write stream is saturated.", + &["route"] + ) + .expect("register memwal_write_stream_rejected_total") +}); +``` + +- [ ] **Step 2: Add observation helper** + +Add a public helper function: + +```rust +pub fn observe_write_stream_state(total: usize, available: usize, waiters: usize) { + WRITE_STREAM_PERMITS_TOTAL.set(total as i64); + WRITE_STREAM_PERMITS_AVAILABLE.set(available as i64); + WRITE_STREAM_WAITERS_TOTAL.set(waiters as i64); +} + +pub fn record_write_stream_acquired(result: &str) { + WRITE_STREAM_ACQUIRED_TOTAL.with_label_values(&[result]).inc(); +} + +pub fn record_write_stream_rejected(route: &str) { + WRITE_STREAM_REJECTED_TOTAL.with_label_values(&[route]).inc(); +} +``` + +- [ ] **Step 3: Compile** + +```bash +cd services/server +cargo check +``` + +Expected: no errors. + +- [ ] **Step 4: Commit** + +```bash +git add services/server/src/observability.rs +git commit -m "feat(write-stream): add Prometheus metrics for limiter state" +``` + +--- + +### Task 4: Add shared helper to translate permit timeout to `AppError` + +**Files:** +- Modify: `services/server/src/routes/mod.rs` + +- [ ] **Step 1: Add helper function** + +At the end of `services/server/src/routes/mod.rs`, before the `#[cfg(test)]` block, add: + +```rust +use crate::services::write_stream::AcquireError; + +/// Convert a write-stream permit acquisition timeout into a 429 response. +pub(super) fn write_stream_saturated(route: &str) -> AppError { + crate::observability::record_write_stream_rejected(route); + AppError::RateLimited( + "Write stream concurrency limit reached; retry after a short delay".into(), + ) +} +``` + +- [ ] **Step 2: Confirm `AppError::RateLimited` exists** + +`services/server/src/types.rs` already has `AppError::RateLimited(String)`, which maps to HTTP 429. No new variant is needed. + +- [ ] **Step 3: Compile** + +```bash +cd services/server +cargo check +``` + +Expected: no errors. + +- [ ] **Step 4: Commit** + +```bash +git add services/server/src/routes/mod.rs services/server/src/types.rs +git commit -m "feat(write-stream): add saturated helper using RateLimited" +``` + +--- + +### Task 5: Gate `/api/remember` and `/api/remember/bulk` + +**Files:** +- Modify: `services/server/src/routes/remember.rs` + +- [ ] **Step 1: Modify `remember` handler** + +In `services/server/src/routes/remember.rs`, in the `remember` handler (around line 616), after inserting the `remember_jobs` row and before calling `spawn_prepare_remember_job`, add permit acquisition: + +```rust +// Acquire a write-stream permit before starting prep work. +let permit = match state + .write_stream_limiter + .acquire(state.config.write_stream_acquire_timeout) + .await +{ + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/remember")); + } +}; + +spawn_prepare_remember_job( + Arc::clone(&state), + job_id.clone(), + text, + owner_owned, + namespace_owned, + auth.public_key.clone(), + permit, +); +``` + +- [ ] **Step 2: Update `spawn_prepare_remember_job` signature and release** + +Change the function signature: + +```rust +fn spawn_prepare_remember_job( + state: Arc, + job_id: String, + text: String, + owner: String, + namespace: String, + agent_public_key: String, + _permit: crate::services::write_stream::WriteStreamPermit, +) { +``` + +The permit drops when the spawned task finishes, which is after the wallet job is enqueued. This is intentional: prep work is gated, and the permit is released once prep hands off to durable queue. + +- [ ] **Step 3: Modify `remember_bulk` handler** + +In `remember_bulk`, after inserting all rows and before spawning prep, add: + +```rust +let item_count = pending_items.len(); +let permits = match state + .write_stream_limiter + .acquire_many(item_count, state.config.write_stream_acquire_timeout) + .await +{ + Ok(permits) => { + crate::routes::record_write_stream_acquired_success(); + permits + } + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/remember/bulk")); + } +}; +``` + +Then pass the permits to `spawn_prepare_bulk_remember_job`: + +```rust +spawn_prepare_bulk_remember_job( + Arc::clone(&state), + owner.clone(), + auth.public_key.clone(), + pending_items, + permits, +); +``` + +- [ ] **Step 4: Update `spawn_prepare_bulk_remember_job` signature** + +Change: + +```rust +fn spawn_prepare_bulk_remember_job( + state: Arc, + owner: String, + agent_public_key: String, + pending_items: Vec, + _permits: crate::services::write_stream::WriteStreamPermit, +) { +``` + +The permit guard is dropped when the prep task completes. + +- [ ] **Step 5: Compile and run remember route tests** + +```bash +cd services/server +cargo check +cargo test --lib routes::remember -- --nocapture +``` + +Expected: compile succeeds, existing tests pass. + +- [ ] **Step 6: Commit** + +```bash +git add services/server/src/routes/remember.rs +git commit -m "feat(write-stream): gate remember and remember/bulk with limiter" +``` + +--- + +### Task 6: Gate `/api/remember/manual` and `/api/analyze` + +**Files:** +- Modify: `services/server/src/routes/remember.rs` +- Modify: `services/server/src/routes/analyze.rs` + +- [ ] **Step 1: Gate `remember_manual`** + +In `remember_manual`, after validation and before calling `engine.store_blob`, add: + +```rust +let _permit = match state + .write_stream_limiter + .acquire(state.config.write_stream_acquire_timeout) + .await +{ + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/remember/manual")); + } +}; +``` + +The `_permit` variable is dropped at the end of the handler after the synchronous `store_blob` call completes. + +- [ ] **Step 2: Gate `analyze` production path** + +In `services/server/src/routes/analyze.rs`, after fact extraction and **before** inserting `remember_jobs` rows, add: + +```rust +let fact_count = facts.len(); +let _permits = match state + .write_stream_limiter + .acquire_many(fact_count, state.config.write_stream_acquire_timeout) + .await +{ + Ok(permits) => { + crate::routes::record_write_stream_acquired_success(); + permits + } + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/analyze")); + } +}; +``` + +This is done before row insertion because analyze rows are created with `status='pending'` and the stale-job sweeper does not clean up pending rows. The permit guard is dropped after the prep/enqueue loop finishes. + +- [ ] **Step 3: Gate `analyze` benchmark path** + +In the benchmark-mode branch (around line 371), wrap each fact's synchronous `store_blob` with a permit. Replace the `store_tasks` map closure body: + +```rust +async move { + let vector = state.embedder.embed(&fact.text).await?; + let _permit = state + .write_stream_limiter + .acquire(std::time::Duration::from_secs(60)) + .await + .map_err(|_| { + AppError::Internal("write stream limiter unavailable".into()) + })?; + crate::routes::record_write_stream_acquired_success(); + let mref = state + .engine + .store_blob( + &owner, + &namespace, + fact.text.as_bytes(), + &vector, + fact.importance, + Some(&agent_pk), + ) + .await?; + Ok::<_, AppError>(AnalyzeAcceptedFact { + text: fact.text, + id: mref.id.clone(), + job_id: mref.id, + }) +} +``` + +- [ ] **Step 4: Compile and run tests** + +```bash +cd services/server +cargo check +cargo test --lib routes::analyze -- --nocapture +``` + +Expected: compile succeeds, existing tests pass. + +- [ ] **Step 5: Commit** + +```bash +git add services/server/src/routes/remember.rs services/server/src/routes/analyze.rs +git commit -m "feat(write-stream): gate remember/manual and analyze with limiter" +``` + +--- + +### Task 7: Gate wallet worker upload path + +**Files:** +- Modify: `services/server/src/jobs.rs` + +- [ ] **Step 1: Acquire permit in `execute_upload_and_transfer`** + +In `services/server/src/jobs.rs`, in `execute_upload_and_transfer`, after decoding the encrypted bytes and before the `upload_blob` call, add: + +```rust +// Acquire a write-stream permit before hitting the sidecar. This ensures +// the sidecar upload queue cannot grow beyond the Rust-managed budget. +let _permit = match state + .write_stream_limiter + .acquire(std::time::Duration::from_secs(60)) + .await +{ + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } + Err(_) => { + // Limiter closed or timeout — leave job in queue for retry. + return Err(WalletJobError::Transient( + "write stream permit unavailable; will retry".into(), + ) + .into_apalis_error()); + } +}; +``` + +- [ ] **Step 2: Ensure permit is released on all paths** + +The `_permit` guard is dropped when `execute_upload_and_transfer` returns. Verify that all early returns in the function use `?` or explicit `return` that goes through the function scope exit. No extra code needed if the guard is declared in the function body. + +- [ ] **Step 3: Compile and run tests** + +```bash +cd services/server +cargo check +cargo test --lib jobs -- --nocapture +``` + +Expected: compile succeeds, tests pass. + +- [ ] **Step 4: Commit** + +```bash +git add services/server/src/jobs.rs +git commit -m "feat(write-stream): gate wallet worker upload with limiter" +``` + +--- + +### Task 8: Add snapshot method and wire metric emission + +**Files:** +- Modify: `services/server/src/services/write_stream.rs` +- Modify: `services/server/src/main.rs` +- Modify: `services/server/src/routes/mod.rs` +- Modify: `services/server/src/jobs.rs` + +To avoid a circular module dependency (`types` → `services/write_stream` → `observability` → `types`), the limiter itself does not call observability. Callers record acquisition outcomes, and a background task polls a snapshot for gauges. + +- [ ] **Step 1: Track waiter count and add snapshot** + +Add to `services/server/src/services/write_stream.rs`: + +```rust +use std::sync::atomic::{AtomicUsize, Ordering}; + +#[derive(Clone, Debug)] +pub struct WriteStreamLimiter { + semaphore: Arc, + max_permits: usize, + waiters: Arc, +} + +#[derive(Clone, Copy, Debug)] +pub struct WriteStreamSnapshot { + pub total: usize, + pub available: usize, + pub waiters: usize, +} +``` + +Update constructors: + +```rust +pub fn new(max_permits: usize) -> Self { + let max_permits = max_permits + .max(MIN_WRITE_STREAM_MAX_CONCURRENCY) + .min(MAX_WRITE_STREAM_MAX_CONCURRENCY); + Self { + semaphore: Arc::new(Semaphore::new(max_permits)), + max_permits, + waiters: Arc::new(AtomicUsize::new(0)), + } +} +``` + +Add method: + +```rust +pub fn snapshot(&self) -> WriteStreamSnapshot { + WriteStreamSnapshot { + total: self.max_permits, + available: self.semaphore.available_permits(), + waiters: self.waiters.load(Ordering::Relaxed), + } +} +``` + +- [ ] **Step 2: Update `acquire_many` to maintain waiter count** + +```rust +pub async fn acquire_many( + &self, + n: usize, + timeout: Duration, +) -> Result { + if n == 0 { + return Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: 0, + }); + } + let n = n.min(self.max_permits); + self.waiters.fetch_add(1, Ordering::Relaxed); + let result = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) + .await + .map_err(|_: Elapsed| AcquireError::Timeout) + .and_then(|res| res.map_err(|_| AcquireError::Closed)); + self.waiters.fetch_sub(1, Ordering::Relaxed); + match result { + Ok(permit) => { + permit.forget(); + Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: n, + }) + } + Err(e) => Err(e), + } +} +``` + +- [ ] **Step 3: Record acquisition outcomes in route helper** + +Update `write_stream_saturated` in `services/server/src/routes/mod.rs`: + +```rust +pub(super) fn write_stream_saturated(route: &str) -> AppError { + crate::observability::record_write_stream_rejected(route); + crate::observability::record_write_stream_acquired("timeout"); + AppError::RateLimited( + "Write stream concurrency limit reached; retry after a short delay".into(), + ) +} +``` + +Add a success-recording helper: + +```rust +pub(super) fn record_write_stream_acquired_success() { + crate::observability::record_write_stream_acquired("success"); +} +``` + +- [ ] **Step 4: Record success in callers** + +In `services/server/src/routes/remember.rs`, after each successful acquisition: + +```rust +crate::routes::record_write_stream_acquired_success(); +``` + +Same in `services/server/src/routes/analyze.rs` and `services/server/src/jobs.rs`. + +- [ ] **Step 5: Add background snapshot task in `main.rs`** + +After `AppState` is constructed, spawn: + +```rust +{ + let state = Arc::clone(&state); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); + loop { + interval.tick().await; + let snap = state.write_stream_limiter.snapshot(); + crate::observability::observe_write_stream_state( + snap.total, + snap.available, + snap.waiters, + ); + } + }); +} +``` + +- [ ] **Step 6: Compile and test** + +```bash +cd services/server +cargo check +cargo test --lib services::write_stream -- --nocapture +``` + +Expected: PASS. + +- [ ] **Step 7: Commit** + +```bash +git add services/server/src/services/write_stream.rs services/server/src/main.rs services/server/src/routes/mod.rs services/server/src/routes/remember.rs services/server/src/routes/analyze.rs services/server/src/jobs.rs +git commit -m "feat(write-stream): emit limiter state and acquisition metrics" +``` + +--- + +### Task 9: Update sidecar safety-net configuration + +**Files:** +- Modify: `services/server/scripts/sidecar/config.ts` + +- [ ] **Step 1: Change sidecar default so it is a higher safety net** + +In `services/server/scripts/sidecar/config.ts`, change: + +```typescript +export const WALRUS_UPLOAD_MAX_CONCURRENCY = parsePositiveIntEnv( + "WALRUS_UPLOAD_MAX_CONCURRENCY", + Math.max(1, SERVER_SUI_PRIVATE_KEYS.length || 1), + 1, + 100, +); +``` + +to: + +```typescript +export const WALRUS_UPLOAD_MAX_CONCURRENCY = parsePositiveIntEnv( + "WALRUS_UPLOAD_MAX_CONCURRENCY", + Math.max(12, SERVER_SUI_PRIVATE_KEYS.length || 1), + 1, + 100, +); +``` + +This keeps the sidecar ceiling above the Rust default of 8 so it only fires as defense-in-depth. + +- [ ] **Step 2: Compile sidecar TypeScript** + +```bash +cd services/server/scripts +npx tsc --noEmit sidecar/config.ts +``` + +Expected: no type errors. + +- [ ] **Step 3: Commit** + +```bash +git add services/server/scripts/sidecar/config.ts +git commit -m "feat(sidecar): raise default upload concurrency safety net" +``` + +--- + +### Task 10: Add integration test for saturation behavior + +**Files:** +- Modify: `services/server/src/routes/remember.rs` (add tests at the bottom) + +- [ ] **Step 1: Add a helper to create a test limiter** + +If not already possible, add a `#[cfg(test)]` helper in `services/server/src/services/write_stream.rs`: + +```rust +#[cfg(test)] +impl WriteStreamLimiter { + pub fn test_new(max_permits: usize) -> Self { + Self::new(max_permits) + } +} +``` + +- [ ] **Step 2: Add saturation unit test in `remember.rs` tests** + +At the bottom of `services/server/src/routes/remember.rs` `mod tests`, add a test that exercises the limiter directly (since full handler integration tests need running DB/sidecar): + +```rust +#[tokio::test] +async fn write_stream_limiter_blocks_beyond_capacity() { + use std::time::Duration; + use crate::services::write_stream::WriteStreamLimiter; + + let limiter = WriteStreamLimiter::test_new(2); + let p1 = limiter.acquire(Duration::from_millis(10)).await.unwrap(); + let p2 = limiter.acquire(Duration::from_millis(10)).await.unwrap(); + let timeout = limiter.acquire(Duration::from_millis(10)).await.unwrap_err(); + assert!(matches!(timeout, crate::services::write_stream::AcquireError::Timeout)); + drop(p1); + drop(p2); +} +``` + +- [ ] **Step 3: Run tests** + +```bash +cd services/server +cargo test --lib routes::remember::tests::write_stream_limiter_blocks_beyond_capacity -- --nocapture +``` + +Expected: PASS. + +- [ ] **Step 4: Commit** + +```bash +git add services/server/src/routes/remember.rs services/server/src/services/write_stream.rs +git commit -m "test(write-stream): add saturation unit test" +``` + +--- + +### Task 11: Full test suite and lint + +**Files:** +- All modified files + +- [ ] **Step 1: Run Rust tests** + +```bash +cd services/server +cargo test --lib +``` + +Expected: all tests pass. + +- [ ] **Step 2: Run clippy** + +```bash +cd services/server +cargo clippy --all-targets -- -D warnings +``` + +Expected: no warnings. Fix any clippy lints that appear (likely around unused imports or the `_permit` naming). + +- [ ] **Step 3: Run sidecar TypeScript tests** + +```bash +cd services/server/scripts +npm test +``` + +Expected: existing tests pass (the sidecar behavior did not change functionally). + +- [ ] **Step 4: Commit any fixes** + +```bash +git add -A +git commit -m "chore(write-stream): clippy fixes and test cleanup" +``` + +--- + +### Task 12: Documentation updates + +**Files:** +- Modify: `docs/relayer/self-hosting.md` +- Modify: `docs/relayer/observability.md` + +- [ ] **Step 1: Document new environment variables** + +In `docs/relayer/self-hosting.md`, add a section: + +```markdown +### Write-stream concurrency + +| Variable | Default | Description | +|---|---|---| +| `WRITE_STREAM_MAX_CONCURRENCY` | `8` | Maximum concurrent active write operations (prep + upload) | +| `WRITE_STREAM_ACQUIRE_TIMEOUT_MS` | `5000` | Handler wait for a write slot before returning `429` | + +Tune `WRITE_STREAM_MAX_CONCURRENCY` down to reduce sidecar pressure, or up to increase throughput when the sidecar has headroom. Keep it below `WALRUS_UPLOAD_MAX_CONCURRENCY` so the sidecar safety net remains meaningful. +``` + +- [ ] **Step 2: Document new metrics** + +In `docs/relayer/observability.md`, add to the metrics table: + +```markdown +| `memwal_write_stream_permits_total` | none | Configured permit ceiling | +| `memwal_write_stream_permits_available` | none | Permits currently free | +| `memwal_write_stream_waiters_total` | none | Tasks waiting for a permit | +| `memwal_write_stream_acquired_total` | `result` | Permit acquisition outcomes | +| `memwal_write_stream_rejected_total` | `route` | Requests rejected with `429` | +``` + +- [ ] **Step 3: Commit** + +```bash +git add docs/relayer/self-hosting.md docs/relayer/observability.md +git commit -m "docs(relayer): document write-stream limiter config and metrics" +``` + +--- + +## Self-Review + +- [ ] **Spec coverage:** + - `WriteStreamLimiter` created → Task 1. + - `AppState` and `Config` wiring → Task 2. + - Metrics → Task 3 and Task 8. + - `/api/remember` gated → Task 5. + - `/api/remember/bulk` gated → Task 5. + - `/api/remember/manual` gated → Task 6. + - `/api/analyze` gated → Task 6. + - Wallet worker upload gated → Task 7. + - Sidecar safety net retained → Task 9. + - Tests → Task 10 and Task 11. + - Docs → Task 12. +- [ ] **Placeholder scan:** No TBD/TODO/"implement later". +- [ ] **Type consistency:** `WriteStreamLimiter::acquire` and `acquire_many` are used consistently across all call sites. +- [ ] **Known follow-ups not in this plan:** distributed limiter for multi-instance deployments (future work in spec). diff --git a/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md b/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md new file mode 100644 index 00000000..ed1a3018 --- /dev/null +++ b/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md @@ -0,0 +1,293 @@ +## Relayer write-stream redesign + +| | | +|---|---| +| **Linear** | [WALM-116](https://linear.app/mysten-labs/issue/WALM-116/relayer-redesign-the-write-stream) | +| **Status** | Design approved, pending implementation plan | +| **Owner** | Max Mai | +| **Date** | 2026-06-15 | + +## 1. Problem statement + +The relayer write stream (`/api/remember`, `/api/remember/bulk`, `/api/remember/manual`, and the store path inside `/api/analyze`) can overload the TypeScript sidecar's Walrus upload path. + +Observed symptoms: + +- The sidecar `/health` endpoint reports `queuedWalrusUploads` climbing during bursts (past incidents reached ~120 queued uploads). +- Sidecar upload-slot acquisition times out after 120 s, returning `503` to the Rust worker. +- Apalis wallet workers retry quickly, burning their retry budget while the backlog drains. +- Congestion-requeue logic exists, but it is reactive: it only kicks in after the queue is already deep. + +The root cause is that **concurrency control lives in the sidecar**, downstream of an unbounded enqueue. The Rust workers can produce upload jobs faster than the sidecar can consume them, so the queue grows. + +## 2. Goal + +Move the effective upload-slot budget into the Rust relayer so that: + +1. The number of concurrently active write operations is bounded by a single, explicit configuration. +2. No embedding/encryption prep work starts unless an upload slot is available. +3. The sidecar `/walrus/upload` queue stays near zero under normal load. +4. Existing durability and retry semantics are preserved. + +## 3. Non-Goals + +- This design does **not** add distributed coordination across multiple relayer instances. It targets the current single-instance deployment. +- It does **not** change the authentication, encryption, embedding, or vector-search paths. +- It does **not** redesign the read stream (`/api/recall`, `/api/ask`). +- It does **not** remove the Apalis job queue; it only gates when work enters and leaves the queue. + +## 4. Design overview + +Add a `WriteStreamLimiter` to `AppState`. It is a thin wrapper around `tokio::sync::Semaphore` with `WRITE_STREAM_MAX_CONCURRENCY` permits. + +Every memory item that results in a sidecar `/walrus/upload` call must acquire one permit before starting prep work, and must release it when the active work phase ends. The same permit pool is shared between prep tasks and upload workers, so the total number of active writes is always bounded. + +```mermaid +flowchart LR + Client["Client"] + Handler["Axum handler
/api/remember"] + Limiter["WriteStreamLimiter
tokio::sync::Semaphore"] + Prep["Prep task
summarize → embed + encrypt"] + Queue["Apalis wallet_jobs"] + Worker["Wallet worker"] + Sidecar["Sidecar /walrus/upload"] + + Client --> Handler + Handler -->|insert remember_jobs| DB[(Postgres)] + Handler -->|acquire permit| Limiter + Limiter -->|permit held| Prep + Prep -->|enqueue + release permit| Queue + Queue -->|dequeue when permit free| Worker + Worker -->|acquire permit| Limiter + Worker -->|upload| Sidecar + Sidecar -->|release permit| Worker +``` + +The sidecar keeps its existing global and per-wallet upload semaphores as a safety net, configured with a slightly higher ceiling than the Rust limit. + +## 5. Components + +### 5.1 `WriteStreamLimiter` + +New file: `services/server/src/services/write_stream.rs` + +Responsibilities: + +- Own the write-stream concurrency budget. +- Provide `acquire(timeout)` and `acquire_many(n, timeout)`. +- Return a permit guard that releases the permit on drop. +- Expose metrics for active/queued permit waiters. + +```rust +pub struct WriteStreamLimiter { + semaphore: Arc, + max_permits: usize, +} + +pub enum AcquireError { + Timeout, + Closed, + WouldExceedCapacity { requested: usize, max: usize }, +} + +impl WriteStreamLimiter { + pub fn new(max_permits: usize) -> Self; + pub async fn acquire(&self, timeout: Duration) -> Result; + pub async fn acquire_many(&self, n: usize, timeout: Duration) -> Result; +} +``` + +The permit guard implements `Drop` to release the permit even if the task panics or the request is cancelled. + +### 5.2 `AppState` + +Add a field: + +```rust +pub write_stream_limiter: Arc, +``` + +Initialize in `main.rs` after config load. + +### 5.3 Route handlers + +#### `/api/remember` + +1. Validate request and insert `remember_jobs` row (`status='running'`). +2. Acquire one permit from `write_stream_limiter` with a bounded timeout. +3. If acquisition times out, return `429 Too Many Requests` immediately. The row remains `running`; the stale-job sweeper marks it failed after `STALE_REMEMBER_JOB_AFTER`. +4. Pass the permit to `spawn_prepare_remember_job`. +5. The prep task releases the permit after enqueueing the wallet job (or on failure). + +#### `/api/remember/bulk` + +1. Validate items and insert one `remember_jobs` row per item. +2. Attempt to acquire `items.len()` permits with a bounded timeout. +3. If the full set cannot be acquired, return `429`. No prep work starts. +4. Split the acquired batch permit into one guard per item. Each item holds its own permit through embedding/SEAL-encryption prep. +5. The item's permit is released when its prep completes (successfully, by handing off to the durable bulk queue) or fails. The wallet worker later reacquires a permit before the actual sidecar upload. + +#### `/api/remember/manual` + +1. Validate request. +2. Acquire one permit. +3. Call `engine.store_blob(...)` synchronously (this calls the sidecar `/walrus/upload` directly). +4. Release permit on completion. + +#### `/api/analyze` + +- **Production path**: after fact extraction, attempt to acquire `facts.len()` permits with a bounded timeout. If acquisition fails, return `429`. Then split the batch permit into one guard per fact. Each fact holds its permit through embedding/SEAL-encryption prep and releases it when that fact's prep completes (or fails). The wallet worker later reacquires a permit before the actual sidecar upload. +- **Benchmark mode (`BENCHMARK_MODE=true`)**: each fact calls `engine.store_blob(...)` synchronously. Acquire one permit per fact, hold it for the duration of the synchronous store, and release on completion. + +### 5.4 Wallet worker + +`execute_wallet_job` acquires one permit before calling `upload_blob`, and releases it after the sidecar call completes (success or terminal failure). + +If permit acquisition times out, the job is left in the Apalis queue and is retried. Because permits are only held by active work, retries naturally back off until slots free up. + +### 5.5 Sidecar safety net + +The sidecar keeps: + +- `WALRUS_UPLOAD_MAX_CONCURRENCY` (default = number of keys, same as today) +- `WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` (default `1`) +- `WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS` (default `120_000`) + +Recommended production tuning: + +```text +WRITE_STREAM_MAX_CONCURRENCY = 8 +WALRUS_UPLOAD_MAX_CONCURRENCY = 12 +WALRUS_UPLOAD_PER_WALLET_CONCURRENCY = 1 +``` + +The sidecar limit should be higher than the Rust limit so it only fires if Rust leaks a slot or if another caller (MCP, direct) bypasses Rust. + +## 6. Data flow + +### 6.1 Single `/api/remember` + +``` +Client POST /api/remember + → Axum handler validates body + → INSERT remember_jobs (status='running') + → write_stream_limiter.acquire(timeout) + └─ timeout → 429 + → spawn_prepare_remember_job(permit) + ├─ summarize-for-embedding (if needed) + ├─ embed(text_summary) || SEAL encrypt(original) + ├─ check storage quota + ├─ enqueue WalletOperation::UploadAndTransfer + └─ release permit + → Apalis worker dequeues job + ├─ write_stream_limiter.acquire(timeout) + ├─ POST /walrus/upload + ├─ update remember_jobs (done / failed) + └─ release permit + ← Client polls /api/remember/:job_id +``` + +### 6.2 `/api/remember/bulk` + +``` +Client POST /api/remember/bulk with N items + → validate all items + → INSERT N remember_jobs rows + → write_stream_limiter.acquire_many(N, timeout) + └─ timeout → 429 (no prep starts) + → spawn N prep tasks, each with one permit + ← return 202 + job_ids +``` + +## 7. Error handling + +| Scenario | Behavior | +|---|---| +| Slot acquisition timeout in handler | Return `429 Too Many Requests`. The `remember_jobs` row remains `running`; stale-job sweeper handles abandoned rows. | +| Request needs more permits than configured ceiling (`WouldExceedCapacity`) | Return `429 Too Many Requests` immediately. The request is never queued and no prep work starts. | +| Prep failure (embed/encrypt/quota) | Release slot immediately, mark `remember_jobs` as `failed`. | +| Wallet job transient failure | Apalis retries. Each retry attempt reacquires a slot before calling the sidecar. | +| Wallet job permanent failure / object locked | Mark `remember_jobs` as `failed`; release slot. | +| Sidecar safety-net timeout | Treated as transient; worker releases slot and Apalis retries. | +| Server crash while slot held | In-memory permit is lost. Surviving queued Apalis jobs reacquire fresh slots on restart. Recovery is bounded by the number of in-flight jobs. | +| Panic in prep or worker task | Permit guard drops on panic, releasing the slot. | + +## 8. Configuration + +New environment variables: + +| Variable | Default | Min | Max | Description | +|---|---|---|---|---| +| `WRITE_STREAM_MAX_CONCURRENCY` | `8` | `1` | `100` | Maximum concurrent active writes (prep + upload). | +| `WRITE_STREAM_ACQUIRE_TIMEOUT_MS` | `5_000` | `100` | `60_000` | How long a handler waits for a permit before returning `429`. | + +Existing variables that remain relevant: + +| Variable | Notes | +|---|---| +| `WALLET_JOB_CONCURRENCY` | Number of Apalis workers. Should be ≥ `WRITE_STREAM_MAX_CONCURRENCY`. | +| `WALRUS_UPLOAD_MAX_CONCURRENCY` | Sidecar safety net; should be > `WRITE_STREAM_MAX_CONCURRENCY`. | +| `WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` | Keep at `1` to avoid Sui object-lock collisions. | +| `WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS` | Sidecar safety-net timeout. Can be reduced once Rust owns the budget. | + +## 9. Observability + +Add Prometheus metrics: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `memwal_write_stream_permits_total` | Gauge | none | Total permits configured. | +| `memwal_write_stream_permits_available` | Gauge | none | Currently available permits. | +| `memwal_write_stream_waiters_total` | Gauge | none | Tasks waiting for a permit. | +| `memwal_write_stream_acquired_total` | Counter | `result="success\|timeout\|failure"` | Permit acquisition outcomes. | +| `memwal_write_stream_rejected_total` | Counter | `route` | Requests rejected with `429` due to slot exhaustion. | + +Logs: + +- `write_stream: acquired permit` / `write_stream: permit timeout` at `info` level. +- Include `route`, `owner_prefix`, `job_id`, and `wait_ms` where applicable. + +## 10. Testing + +### 10.1 Unit tests + +- `WriteStreamLimiter::new(0)` falls back to minimum. +- `acquire` returns a permit when available. +- `acquire` times out correctly. +- `acquire_many` returns all permits or none. +- Permit guard releases on drop and on panic. + +### 10.2 Integration tests + +- **Saturation test**: `WRITE_STREAM_MAX_CONCURRENCY=2`, slow mock sidecar, 10 concurrent `/api/remember` requests. Assert at most 2 concurrent `/walrus/upload` calls and the rest receive `429`. +- **Bulk atomicity test**: `WRITE_STREAM_MAX_CONCURRENCY=3`, `/api/remember/bulk` with 5 items. Assert `429` and zero embeddings/encrypts. +- **Recovery test**: kill a prep task mid-flight, assert permit is released and other requests can proceed. + +### 10.3 Regression tests + +- `cargo test` passes. +- Sidecar TS test suite passes. +- Existing end-to-end tests for remember/recall pass. + +### 10.4 Load validation + +Run a targeted benchmark that previously produced `queuedWalrusUploads > 20`. Verify the metric stays near zero and p95 `/api/remember` latency is stable. + +## 11. Risks and mitigations + +| Risk | Mitigation | +|---|---| +| Lower throughput because prep and upload share the same pool. | Tune `WRITE_STREAM_MAX_CONCURRENCY` and `WALLET_JOB_CONCURRENCY`; if needed, a future iteration can split prep and upload slots. | +| Bulk requests starve single-item requests. | Bulk must acquire all permits upfront; clients should expect `429` on large bursts. Consider lowering `MAX_BULK_ITEMS` if needed. | +| Slot leak due to a bug. | Permit guard uses `Drop`; add metrics for available permits to detect leaks; sidecar safety net catches overflow. | +| Retry storm after a crash. | Bounded by in-flight jobs at crash time; Apalis retry budget still applies. | + +## 12. Future work + +- **Distributed limiter**: If the relayer scales horizontally, replace the in-process semaphore with a Redis-backed semaphore so instances share the budget. +- **Two-tier slots**: If throughput measurements show prep is the bottleneck, split into prep-slot and upload-slot pools. +- **Admission by account**: Consider per-account write slots to prevent one noisy neighbor from consuming the global budget. + +## 13. Open questions + +None. Design approved for implementation planning. diff --git a/services/server/scripts/sidecar/config.ts b/services/server/scripts/sidecar/config.ts index 26e3d796..38f9e0b5 100644 --- a/services/server/scripts/sidecar/config.ts +++ b/services/server/scripts/sidecar/config.ts @@ -142,9 +142,11 @@ export const UPLOAD_RELAY_TIP_CACHE_TTL_MS = (() => { // Walrus upload concurrency limits // ============================================================ +// Keep the sidecar ceiling above the Rust WriteStreamLimiter default (8) +// so the sidecar only acts as defense-in-depth. export const WALRUS_UPLOAD_MAX_CONCURRENCY = parsePositiveIntEnv( "WALRUS_UPLOAD_MAX_CONCURRENCY", - Math.max(1, SERVER_SUI_PRIVATE_KEYS.length || 1), + Math.max(12, SERVER_SUI_PRIVATE_KEYS.length || 1), 1, 100, ); diff --git a/services/server/src/alerts.rs b/services/server/src/alerts.rs index 49224a01..249a86b8 100644 --- a/services/server/src/alerts.rs +++ b/services/server/src/alerts.rs @@ -458,8 +458,12 @@ impl SlackPayload { (Some(before), Some(after)) => { format!("*On-chain system version:* `{}` → `{}`\n", before, after) } - (None, Some(after)) => format!("*On-chain system version (after refresh):* `{}`\n", after), - (Some(before), None) => format!("*On-chain system version (before refresh):* `{}`\n", before), + (None, Some(after)) => { + format!("*On-chain system version (after refresh):* `{}`\n", after) + } + (Some(before), None) => { + format!("*On-chain system version (before refresh):* `{}`\n", before) + } (None, None) => String::new(), }; let details = format!( @@ -672,9 +676,7 @@ Congestion-requeued jobs ride it out with minutes-scale backoff, but sustained s let threshold_wal = format_wal_amount(alert.threshold); let required_line = alert .required - .map(|required| { - format!("*Required:* {} WAL\n", format_wal_amount(required)) - }) + .map(|required| format!("*Required:* {} WAL\n", format_wal_amount(required))) .unwrap_or_default(); let action = "*Action (ops):* top up WAL for this relayer wallet before retrying. If the wallet is being topped up, rotate or temporarily remove that key from pool to keep uploads flowing." diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index 7bb2eab0..3313f39e 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -908,6 +908,32 @@ async fn execute_upload_and_transfer( encrypted.len(), ); + // Acquire a write-stream permit before hitting the sidecar. This ensures + // the sidecar upload queue cannot grow beyond the Rust-managed budget. + let _permit = match state + .write_stream_limiter + .acquire(std::time::Duration::from_secs(60)) + .await + { + Ok(permit) => { + crate::observability::record_write_stream_acquired_success(); + permit + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + crate::observability::record_write_stream_acquired("timeout"); + return Err(WalletJobError::Transient( + "write stream permit unavailable; will retry".into(), + )); + } + Err(_) => { + crate::observability::record_write_stream_acquired("failure"); + // Limiter closed — leave job in queue for retry. + return Err(WalletJobError::Transient( + "write stream limiter closed; will retry".into(), + )); + } + }; + // ── Upload to Walrus via sidecar (using pinned wallet_index) ─ let upload_result = crate::storage::walrus::upload_blob( &state.http_client, @@ -1449,11 +1475,7 @@ async fn maybe_alert_walrus_low_wal_balance( error: msg.to_string(), }; - if let Err(err) = state - .alerts - .notify_walrus_low_wal_balance(alert) - .await - { + if let Err(err) = state.alerts.notify_walrus_low_wal_balance(alert).await { tracing::warn!( "[wallet-job:upload] failed to send Slack alert for low WAL balance: {}", err @@ -1814,6 +1836,33 @@ pub async fn execute_remember( wallet_index_for_upload_attempt(starting_key_index, attempt_info.current, key_pool_len) .expect("non-empty key pool must yield wallet index"); + // Acquire a write-stream permit before hitting the sidecar. Legacy + // RememberJob rows share the same upload budget as UploadAndTransfer jobs. + let _permit = match state + .write_stream_limiter + .acquire(std::time::Duration::from_secs(60)) + .await + { + Ok(permit) => { + crate::observability::record_write_stream_acquired_success(); + permit + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + crate::observability::record_write_stream_acquired("timeout"); + return Err(WalletJobError::Transient( + "write stream permit unavailable; will retry".into(), + ) + .into_apalis_error()); + } + Err(_) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(WalletJobError::Transient( + "write stream limiter closed; will retry".into(), + ) + .into_apalis_error()); + } + }; + // ── Step 3: walrus upload (the slow part ~2-3s) ─────────────── let upload_result = crate::storage::walrus::upload_blob( &state.http_client, @@ -2152,7 +2201,7 @@ mod tests { classify_wallet_remember_handoff_failure, congestion_backoff_secs, escalate_if_gas_pool_exhausted, gas_pool_exhaustion_threshold, is_walrus_package_version_mismatch, mark_remember_job_failed, parse_locked_object_info, - wallet_index_for_upload_attempt, wallet_job_request, parse_wal_balance_alert_info, + parse_wal_balance_alert_info, wallet_index_for_upload_attempt, wallet_job_request, WalletJob, WalletJobAttemptInfo, WalletJobError, WalletOperation, MAX_ATTEMPTS, MAX_CONGESTION_REQUEUES, }; @@ -2291,7 +2340,9 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt // The whole requeue budget must outlive a realistic backlog drain // (the 2026-06-10 queue needed ~8 minutes at ~16 uploads/min). - let total: u64 = (0..MAX_CONGESTION_REQUEUES).map(congestion_backoff_secs).sum(); + let total: u64 = (0..MAX_CONGESTION_REQUEUES) + .map(congestion_backoff_secs) + .sum(); assert!( total >= 20 * 60, "congestion requeue budget should span >= 20 minutes, got {}s", @@ -2437,7 +2488,8 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt #[test] fn parse_wal_balance_alert_info_extracts_required_and_available() { - let parsed = parse_wal_balance_alert_info(LOW_WAL_BALANCE_ERR).expect("expected low WAL signal"); + let parsed = + parse_wal_balance_alert_info(LOW_WAL_BALANCE_ERR).expect("expected low WAL signal"); assert_eq!(parsed.required, Some(64367730)); assert_eq!(parsed.available, 10708877); } @@ -2445,10 +2497,7 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt #[test] fn classify_low_wal_balance_is_dedicated_transient() { let classified = WalletJobError::classify_sidecar_error(LOW_WAL_BALANCE_ERR); - assert!(matches!( - classified, - WalletJobError::WalrusBalanceLow(_) - )); + assert!(matches!(classified, WalletJobError::WalrusBalanceLow(_))); assert!(!classified.aborts_retries()); assert!(!classified.is_permanent()); assert_eq!(classified.kind(), "walrus_balance_low"); diff --git a/services/server/src/main.rs b/services/server/src/main.rs index be16ba41..88d41816 100644 --- a/services/server/src/main.rs +++ b/services/server/src/main.rs @@ -32,7 +32,9 @@ use jobs::{ execute_bulk_remember, execute_wallet_job, BulkRememberJob, MetaTransferJob, RememberJob, WalletJobStorage, }; -use services::{CompositeRanker, Embedder, Extractor, LlmExtractor, OpenAiEmbedder, Ranker}; +use services::{ + CompositeRanker, Embedder, Extractor, LlmExtractor, OpenAiEmbedder, Ranker, WriteStreamLimiter, +}; use storage::db::VectorDb; use types::{ AppState, Config, KeyPool, DEFAULT_BLOB_CACHE_MAX_BYTES, DEFAULT_BLOB_CACHE_TTL_SECS, @@ -439,6 +441,14 @@ async fn main() { // Wrap the immutable config so the MemoryEngine + handlers share it. let config = Arc::new(config); + let write_stream_limiter = + Arc::new(WriteStreamLimiter::new(config.write_stream_max_concurrency)); + tracing::info!( + " write stream limiter: max_concurrency={} acquire_timeout_ms={}", + write_stream_limiter.max_permits(), + config.write_stream_acquire_timeout.as_millis(), + ); + // Select the persistence engine. Production = WalrusSealEngine (SEAL // encrypt happens in the handler/client; the engine uploads the // ciphertext to Walrus and indexes the row, with the Redis blob @@ -497,6 +507,7 @@ async fn main() { blob_cache_ttl, blob_cache_max_bytes, embedding_cache_ttl, + write_stream_limiter, }); tracing::info!( @@ -508,6 +519,24 @@ async fn main() { } ); + // Background task: emit write-stream limiter state every 5s. + { + let state = Arc::clone(&state); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + let snap = state.write_stream_limiter.snapshot(); + crate::observability::observe_write_stream_state( + snap.total, + snap.available, + snap.waiters, + ); + } + }); + } + // Sidecar upload-queue saturation monitor. The watchdog above only // checks that /health answers; during the 2026-06-10 congestion incident // it stayed green while 120 uploads queued and jobs burned their retry @@ -516,8 +545,7 @@ async fn main() { // throttle the burst) — before queued requests outlive the sidecar's // 120s acquire timeout and start failing. let saturation_threshold = parse_env_u64("SIDECAR_QUEUE_SATURATION_THRESHOLD", 20, 1, 10_000); - let saturation_consecutive = - parse_env_u32("SIDECAR_QUEUE_SATURATION_CONSECUTIVE", 4, 1, 100); + let saturation_consecutive = parse_env_u32("SIDECAR_QUEUE_SATURATION_CONSECUTIVE", 4, 1, 100); let saturation_interval_secs = parse_env_u64("SIDECAR_QUEUE_SATURATION_INTERVAL_SECS", 30, 5, 300); tracing::info!( @@ -532,9 +560,8 @@ async fn main() { let monitor_alerts = Arc::clone(&state.alerts); let monitor_network = config.sui_network.clone(); tokio::spawn(async move { - let mut interval = tokio::time::interval(std::time::Duration::from_secs( - saturation_interval_secs, - )); + let mut interval = + tokio::time::interval(std::time::Duration::from_secs(saturation_interval_secs)); let mut consecutive_saturated = 0u32; loop { interval.tick().await; @@ -593,8 +620,9 @@ async fn main() { threshold: saturation_threshold, consecutive_checks: consecutive_saturated, }; - if let Err(err) = - monitor_alerts.notify_walrus_upload_queue_saturated(alert).await + if let Err(err) = monitor_alerts + .notify_walrus_upload_queue_saturated(alert) + .await { tracing::warn!(" sidecar: saturation alert delivery failed: {}", err); } diff --git a/services/server/src/observability.rs b/services/server/src/observability.rs index 64b56001..aab80f66 100644 --- a/services/server/src/observability.rs +++ b/services/server/src/observability.rs @@ -176,6 +176,48 @@ static DB_POOL: LazyLock = LazyLock::new(|| { .expect("register memwal_db_pool_connections") }); +static WRITE_STREAM_PERMITS_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_permits_total", + "Total configured write-stream permits." + ) + .expect("register memwal_write_stream_permits_total") +}); + +static WRITE_STREAM_PERMITS_AVAILABLE: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_permits_available", + "Currently available write-stream permits." + ) + .expect("register memwal_write_stream_permits_available") +}); + +static WRITE_STREAM_WAITERS_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge!( + "memwal_write_stream_waiters_total", + "Tasks currently waiting for a write-stream permit." + ) + .expect("register memwal_write_stream_waiters_total") +}); + +static WRITE_STREAM_ACQUIRED_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "memwal_write_stream_acquired_total", + "Write-stream permit acquisition outcomes.", + &["result"] + ) + .expect("register memwal_write_stream_acquired_total") +}); + +static WRITE_STREAM_REJECTED_TOTAL: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "memwal_write_stream_rejected_total", + "Requests rejected because the write stream is saturated.", + &["route"] + ) + .expect("register memwal_write_stream_rejected_total") +}); + pub fn init_tracing() -> TelemetryGuard { let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| "memwal_server=info,tower_http=info".into()); @@ -566,6 +608,37 @@ pub fn update_db_pool_metrics(pool: &sqlx::PgPool) { .set(pool.num_idle() as i64); } +pub fn observe_write_stream_state(total: usize, available: usize, waiters: usize) { + WRITE_STREAM_PERMITS_TOTAL.set(total as i64); + WRITE_STREAM_PERMITS_AVAILABLE.set(available as i64); + WRITE_STREAM_WAITERS_TOTAL.set(waiters as i64); +} + +pub fn record_write_stream_acquired(result: &str) { + WRITE_STREAM_ACQUIRED_TOTAL + .with_label_values(&[result]) + .inc(); +} + +pub fn record_write_stream_acquired_success() { + record_write_stream_acquired("success"); +} + +pub fn record_write_stream_acquired_success_n(n: usize) { + if n == 0 { + return; + } + WRITE_STREAM_ACQUIRED_TOTAL + .with_label_values(&["success"]) + .inc_by(n as u64); +} + +pub fn record_write_stream_rejected(route: &str) { + WRITE_STREAM_REJECTED_TOTAL + .with_label_values(&[route]) + .inc(); +} + fn record_http_request(method: &str, route: &str, status: StatusCode, elapsed: Duration) { let status = status.as_u16().to_string(); HTTP_REQUESTS_TOTAL diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 6cc7f921..640131a1 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -384,6 +384,25 @@ pub async fn analyze( let fact = fact.clone(); async move { let vector = state.embedder.embed(&fact.text).await?; + let _permit = match state + .write_stream_limiter + .acquire(std::time::Duration::from_secs(60)) + .await + { + Ok(permit) => { + crate::observability::record_write_stream_acquired_success(); + permit + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/analyze")); + } + Err(_) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::Internal( + "write stream limiter unavailable".into(), + )); + } + }; // importance is threaded through the engine // (see store_blob signature in engine::MemoryEngine). // The PlaintextEngine persists it on the new @@ -441,10 +460,47 @@ pub async fn analyze( )); } + // Acquire write-stream permits before starting prep work. Each fact's + // embed + SEAL task will hold one permit until its prep completes, so the + // total active prep work is bounded by the write-stream budget alongside + // the wallet upload workers. + let mut permits = match state + .write_stream_limiter + .acquire_many(facts.len(), state.config.write_stream_acquire_timeout) + .await + { + Ok(permits) => { + crate::observability::record_write_stream_acquired_success_n(facts.len()); + permits + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/analyze")); + } + Err(crate::services::write_stream::AcquireError::WouldExceedCapacity { + requested, + max: _, + }) => { + crate::observability::record_write_stream_rejected("/api/analyze"); + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::RateLimited(format!( + "Analyze extracted {} facts, exceeding write stream capacity; reduce input", + requested + ))); + } + Err(_) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::Internal( + "write stream limiter unavailable".into(), + )); + } + }; + // Step 2: embed + SEAL encrypt all facts concurrently (no wallet needed yet). // This is the fast part (~300-500ms), done in the request handler so: // - No plaintext stored in job payload // - Exact ciphertext size known for quota check + // Each fact holds its own permit; the permit drops when that fact's prep + // completes, releasing the slot for the next fact or an upload worker. let auth_pubkey_base = auth.public_key.clone(); let prep_tasks: Vec<_> = facts .iter() @@ -452,7 +508,13 @@ pub async fn analyze( let state = Arc::clone(&state); let owner = owner.clone(); let fact = fact.clone(); + let _permit = permits + .split_one() + .expect("acquired permits equal facts.len()"); async move { + // Hold the permit until this fact's prep hands off to the durable wallet queue. + let _permit = _permit; + let embed_fut = state.embedder.embed(&fact.text); let encrypt_fut = crate::storage::seal::seal_encrypt( &state.http_client, diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index 9cfc330d..1ad6b6a8 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -182,6 +182,15 @@ pub(super) fn zip_search_hit_fields_onto_hydrated( } } +/// Convert a write-stream permit acquisition timeout into a 429 response. +pub(super) fn write_stream_saturated(route: &str) -> AppError { + crate::observability::record_write_stream_rejected(route); + crate::observability::record_write_stream_acquired("timeout"); + AppError::RateLimited( + "Write stream concurrency limit reached; retry after a short delay".into(), + ) +} + #[cfg(test)] mod tests { use super::{collect_bounded_results, truncate_str}; diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index d39e1b0b..df6ddd08 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -101,9 +101,13 @@ fn spawn_prepare_remember_job( owner: String, namespace: String, agent_public_key: String, + // Held until prep completes; releases the permit on drop. + permit: crate::services::write_stream::WriteStreamPermit, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { + // Hold the permit until prep hands off to the durable wallet queue. + let _permit = permit; let work = async move { let result: Result<(), AppError> = async { // texts beyond the embedder's context window must be @@ -202,6 +206,9 @@ fn spawn_prepare_bulk_remember_job( owner: String, agent_public_key: String, pending_items: Vec, + // One guard per item; each is held until that item's prep completes and + // releases its own slot. + item_permits: Vec, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -213,10 +220,14 @@ fn spawn_prepare_bulk_remember_job( let result: Result<(), AppError> = async { let prep_tasks: Vec<_> = pending_items .into_iter() - .map(|item| { + .zip(item_permits) + .map(|(item, _permit)| { let state = Arc::clone(&state); let owner = owner.clone(); async move { + // Hold the permit until this item's prep hands off to the durable bulk queue. + let _permit = _permit; + // bulk items can carry up to MAX_REMEMBER_TEXT_BYTES // each, so the same summarize-before-embed rule applies here. let needs_summary = item.text.len() > SUMMARIZE_THRESHOLD_BYTES @@ -635,6 +646,8 @@ pub async fn remember( let namespace_owned = namespace.clone(); let text = body.text; + // Insert the remember_jobs row first. If permit acquisition times out, the + // row remains `running` and the stale-job sweeper will mark it failed. let job_id = uuid::Uuid::new_v4().to_string(); sqlx::query( @@ -647,6 +660,35 @@ pub async fn remember( .await .map_err(|e| AppError::Internal(format!("Failed to create job row: {}", e)))?; + // Acquire a write-stream permit after persisting the job row. On timeout + // the caller gets 429 immediately while the row stays running for the + // stale-job sweeper. + let permit = match state + .write_stream_limiter + .acquire(state.config.write_stream_acquire_timeout) + .await + { + Ok(permit) => { + crate::observability::record_write_stream_acquired_success(); + permit + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/remember")); + } + Err(crate::services::write_stream::AcquireError::Closed) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::Internal("write stream limiter closed".into())); + } + Err(crate::services::write_stream::AcquireError::WouldExceedCapacity { .. }) => { + crate::observability::record_write_stream_rejected("/api/remember"); + crate::observability::record_write_stream_acquired("failure"); + // n=1 can never exceed capacity, but keep the arm consistent. + return Err(AppError::RateLimited( + "Write stream capacity exceeded; retry after a short delay".into(), + )); + } + }; + spawn_prepare_remember_job( Arc::clone(&state), job_id.clone(), @@ -654,6 +696,7 @@ pub async fn remember( owner_owned, namespace_owned, auth.public_key.clone(), + permit, ); tracing::info!( @@ -751,13 +794,17 @@ pub async fn remember_bulk( } } + let item_count = body.items.len(); + let owner = &auth.owner; tracing::info!( "remember_bulk: {} items owner={}", - body.items.len(), + item_count, &owner[..10.min(owner.len())], ); + // Insert remember_jobs rows first. If permit acquisition times out, the + // rows remain `running` and the stale-job sweeper will mark them failed. let mut job_ids: Vec = Vec::with_capacity(body.items.len()); let mut pending_items: Vec = Vec::with_capacity(body.items.len()); @@ -784,11 +831,54 @@ pub async fn remember_bulk( let total = job_ids.len(); + // Acquire write-stream permits after persisting rows. On timeout the + // caller gets 429 immediately while the rows stay running for the + // stale-job sweeper. + let mut permits = match state + .write_stream_limiter + .acquire_many(item_count, state.config.write_stream_acquire_timeout) + .await + { + Ok(permits) => { + crate::observability::record_write_stream_acquired_success_n(item_count); + permits + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/remember/bulk")); + } + Err(crate::services::write_stream::AcquireError::WouldExceedCapacity { + requested, + max: _, + }) => { + crate::observability::record_write_stream_rejected("/api/remember/bulk"); + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::RateLimited(format!( + "Bulk request of {} items exceeds write stream capacity; reduce batch size", + requested + ))); + } + Err(crate::services::write_stream::AcquireError::Closed) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::Internal("write stream limiter closed".into())); + } + }; + + // Split the acquired batch permit into one guard per item so each item's + // prep task releases its own slot when its embedding/encryption finishes. + let item_permits: Vec = (0..item_count) + .map(|_| { + permits + .split_one() + .expect("acquired permits equal pending_items.len()") + }) + .collect(); + spawn_prepare_bulk_remember_job( Arc::clone(&state), owner.clone(), auth.public_key.clone(), pending_items, + item_permits, ); tracing::info!("remember_bulk accepted: {} items owner={}", total, owner,); @@ -913,6 +1003,30 @@ pub async fn remember_manual( // the engine owns persistence, not policy). rate_limit::check_storage_quota(&state, owner, encrypted_bytes.len() as i64).await?; + // Gate on write-stream capacity so a saturated stream can't be driven + // through the synchronous manual path. + let _permit = match state + .write_stream_limiter + .acquire(state.config.write_stream_acquire_timeout) + .await + { + Ok(permit) => { + crate::observability::record_write_stream_acquired_success(); + permit + } + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated( + "/api/remember/manual", + )); + } + Err(_) => { + crate::observability::record_write_stream_acquired("failure"); + return Err(AppError::Internal( + "write stream limiter unavailable".into(), + )); + } + }; + // Persist via the storage engine: Walrus upload (pool key pays gas, // configured storage epochs, immediate transfer to owner) -> Postgres index row. // Same logic as before, now in engine/walrus_seal.rs::store_blob. @@ -1077,6 +1191,8 @@ mod tests { sponsor_rate_limit: crate::types::SponsorRateLimitConfig::default(), allowed_origins: String::new(), benchmark_mode: false, + write_stream_max_concurrency: 8, + write_stream_acquire_timeout: std::time::Duration::from_millis(5_000), } } @@ -1126,4 +1242,24 @@ mod tests { assert!(seen.iter().all(|len| *len <= SUMMARIZE_CHUNK_BYTES + 1024)); assert!(seen.iter().all(|len| *len < MAX_REMEMBER_TEXT_BYTES / 4)); } + + #[tokio::test] + async fn write_stream_limiter_blocks_beyond_capacity() { + use crate::services::write_stream::WriteStreamLimiter; + use std::time::Duration; + + let limiter = WriteStreamLimiter::test_new(2); + let p1 = limiter.acquire(Duration::from_millis(10)).await.unwrap(); + let p2 = limiter.acquire(Duration::from_millis(10)).await.unwrap(); + let timeout = limiter + .acquire(Duration::from_millis(10)) + .await + .unwrap_err(); + assert!(matches!( + timeout, + crate::services::write_stream::AcquireError::Timeout + )); + drop(p1); + drop(p2); + } } diff --git a/services/server/src/services/extractor.rs b/services/server/src/services/extractor.rs index ad900533..a0783fbe 100644 --- a/services/server/src/services/extractor.rs +++ b/services/server/src/services/extractor.rs @@ -345,7 +345,10 @@ impl LlmExtractor { // client, and the SDK/harness will not retry a 202 with // facts=[]. The explicit string `"NONE"` from the LLM remains // the valid no-facts path (handled by `parse_extracted_facts`). - let raw_content = api_resp.choices.first().and_then(|c| c.message.content.as_deref()); + let raw_content = api_resp + .choices + .first() + .and_then(|c| c.message.content.as_deref()); let content = match raw_content { Some(s) => s.trim().to_string(), None => { @@ -1433,9 +1436,13 @@ mod tests { // Transient assert!(is_upstream_status_transient(StatusCode::TOO_MANY_REQUESTS)); - assert!(is_upstream_status_transient(StatusCode::INTERNAL_SERVER_ERROR)); + assert!(is_upstream_status_transient( + StatusCode::INTERNAL_SERVER_ERROR + )); assert!(is_upstream_status_transient(StatusCode::BAD_GATEWAY)); - assert!(is_upstream_status_transient(StatusCode::SERVICE_UNAVAILABLE)); + assert!(is_upstream_status_transient( + StatusCode::SERVICE_UNAVAILABLE + )); assert!(is_upstream_status_transient(StatusCode::GATEWAY_TIMEOUT)); // Not transient — real bugs / config errors that retrying won't fix diff --git a/services/server/src/services/mod.rs b/services/server/src/services/mod.rs index 9cdc683e..d7782397 100644 --- a/services/server/src/services/mod.rs +++ b/services/server/src/services/mod.rs @@ -27,6 +27,7 @@ pub mod embedder; pub mod extractor; pub mod llm_chat; pub mod ranker; +pub mod write_stream; // Placeholder module — reserved namespace for the consolidator // (linked-memory-ids + supersede logic). Doc-only until a real caller @@ -36,3 +37,4 @@ pub mod consolidator; pub use embedder::{Embedder, OpenAiEmbedder}; pub use extractor::{Extractor, LlmExtractor}; pub use ranker::{CompositeRanker, Ranker}; +pub use write_stream::WriteStreamLimiter; diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs new file mode 100644 index 00000000..be5bbe81 --- /dev/null +++ b/services/server/src/services/write_stream.rs @@ -0,0 +1,436 @@ +//! Write-stream concurrency limiter. +//! +//! Owns the single in-process budget for active write operations +//! (prep + upload). Every memory item that will result in a sidecar +//! /walrus/upload call must acquire a permit before starting work and +//! release it when the active phase ends. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Semaphore; + +const DEFAULT_WRITE_STREAM_MAX_CONCURRENCY: usize = 8; +const MIN_WRITE_STREAM_MAX_CONCURRENCY: usize = 1; +const MAX_WRITE_STREAM_MAX_CONCURRENCY: usize = 100; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AcquireError { + Timeout, + Closed, + WouldExceedCapacity { requested: usize, max: usize }, +} + +impl std::fmt::Display for AcquireError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AcquireError::Timeout => write!(f, "write stream concurrency limit reached"), + AcquireError::Closed => write!(f, "write stream limiter closed"), + AcquireError::WouldExceedCapacity { requested, max } => write!( + f, + "write stream request for {requested} permits exceeds capacity of {max}" + ), + } + } +} + +impl std::error::Error for AcquireError {} + +/// Guard that releases one or more permits when dropped. +#[derive(Debug)] +#[must_use = "permit releases on drop"] +pub struct WriteStreamPermit { + semaphore: Arc, + pub(crate) permits: usize, +} + +impl Drop for WriteStreamPermit { + fn drop(&mut self) { + self.semaphore.add_permits(self.permits); + } +} + +impl WriteStreamPermit { + /// Split off a single permit into a new guard. + /// + /// Returns `None` if this guard currently holds zero permits. + /// The total number of permits held by both guards remains unchanged. + pub fn split_one(&mut self) -> Option { + if self.permits == 0 { + return None; + } + self.permits -= 1; + Some(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: 1, + }) + } +} + +/// Snapshot of the write-stream limiter state. +#[derive(Clone, Copy, Debug)] +pub struct WriteStreamSnapshot { + pub total: usize, + pub available: usize, + pub waiters: usize, +} + +/// Decrements the waiter counter on drop, guarding against future cancellation. +struct WaiterGuard { + waiters: Arc, +} + +impl WaiterGuard { + fn new(waiters: Arc) -> Self { + waiters.fetch_add(1, Ordering::Relaxed); + Self { waiters } + } +} + +impl Drop for WaiterGuard { + fn drop(&mut self) { + self.waiters.fetch_sub(1, Ordering::Relaxed); + } +} + +/// In-process concurrency limiter for the write stream. +#[derive(Clone, Debug)] +pub struct WriteStreamLimiter { + semaphore: Arc, + max_permits: usize, + waiters: Arc, +} + +impl WriteStreamLimiter { + pub fn new(max_permits: usize) -> Self { + let max_permits = max_permits.clamp( + MIN_WRITE_STREAM_MAX_CONCURRENCY, + MAX_WRITE_STREAM_MAX_CONCURRENCY, + ); + Self { + semaphore: Arc::new(Semaphore::new(max_permits)), + max_permits, + waiters: Arc::new(AtomicUsize::new(0)), + } + } + + pub fn max_permits(&self) -> usize { + self.max_permits + } + + /// Return a point-in-time snapshot of limiter state. + pub fn snapshot(&self) -> WriteStreamSnapshot { + WriteStreamSnapshot { + total: self.max_permits, + available: self.semaphore.available_permits(), + waiters: self.waiters.load(Ordering::Relaxed), + } + } + + /// Acquire a single permit, waiting up to `timeout`. + pub async fn acquire(&self, timeout: Duration) -> Result { + self.acquire_many(1, timeout).await + } + + /// Acquire `n` permits atomically with respect to this call, waiting up to `timeout`. + /// The returned guard releases all `n` permits on drop. + pub async fn acquire_many( + &self, + n: usize, + timeout: Duration, + ) -> Result { + if n == 0 { + return Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: 0, + }); + } + if n > self.max_permits { + return Err(AcquireError::WouldExceedCapacity { + requested: n, + max: self.max_permits, + }); + } + let _waiter_guard = WaiterGuard::new(Arc::clone(&self.waiters)); + let result = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) + .await + .map_err(|_| AcquireError::Timeout) + .and_then(|res| res.map_err(|_| AcquireError::Closed)); + match result { + Ok(permit) => { + permit.forget(); + Ok(WriteStreamPermit { + semaphore: Arc::clone(&self.semaphore), + permits: n, + }) + } + Err(e) => Err(e), + } + } +} + +impl Default for WriteStreamLimiter { + fn default() -> Self { + Self::new(DEFAULT_WRITE_STREAM_MAX_CONCURRENCY) + } +} + +#[cfg(test)] +impl WriteStreamLimiter { + pub fn test_new(max_permits: usize) -> Self { + Self::new(max_permits) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn acquire_returns_permit_when_available() { + let limiter = WriteStreamLimiter::new(1); + let permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + assert_eq!(limiter.snapshot().available, 0); + drop(permit); + assert_eq!(limiter.snapshot().available, 1); + } + + #[tokio::test] + async fn acquire_times_out_when_exhausted() { + let limiter = WriteStreamLimiter::new(1); + let _permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + let err = limiter + .acquire(Duration::from_millis(50)) + .await + .unwrap_err(); + assert!(matches!(err, AcquireError::Timeout)); + } + + #[tokio::test] + async fn acquire_many_returns_all_or_none() { + let limiter = WriteStreamLimiter::new(3); + let _p1 = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + // asking for 3 when only 2 are free should time out + let err = limiter + .acquire_many(3, Duration::from_millis(50)) + .await + .unwrap_err(); + assert!(matches!(err, AcquireError::Timeout)); + assert_eq!(limiter.snapshot().available, 2); + let p23 = limiter + .acquire_many(2, Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!(limiter.snapshot().available, 0); + drop(p23); + assert_eq!(limiter.snapshot().available, 2); + } + + #[tokio::test] + async fn zero_permits_noop() { + let limiter = WriteStreamLimiter::new(1); + let guard = limiter + .acquire_many(0, Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!(guard.permits, 0); + assert_eq!(limiter.snapshot().available, 1); + } + + #[tokio::test] + async fn snapshot_reports_state() { + let limiter = WriteStreamLimiter::new(3); + let before = limiter.snapshot(); + assert_eq!(before.total, 3); + assert_eq!(before.available, 3); + assert_eq!(before.waiters, 0); + + let permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + let during = limiter.snapshot(); + assert_eq!(during.total, 3); + assert_eq!(during.available, 2); + assert_eq!(during.waiters, 0); + drop(permit); + + let after = limiter.snapshot(); + assert_eq!(after.total, 3); + assert_eq!(after.available, 3); + assert_eq!(after.waiters, 0); + } + + #[tokio::test] + async fn waiters_increments_while_waiting() { + let limiter = WriteStreamLimiter::new(1); + let _permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let waiting = tokio::spawn({ + let limiter = limiter.clone(); + async move { + // signal that we are about to acquire + let _ = entered_tx.send(()); + limiter.acquire(Duration::from_millis(200)).await + } + }); + entered_rx.await.unwrap(); + tokio::time::timeout(Duration::from_millis(100), async { + while limiter.snapshot().waiters != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("waiter count should reach 1"); + + let result = waiting.await.unwrap(); + assert!(matches!(result, Err(AcquireError::Timeout))); + assert_eq!(limiter.snapshot().waiters, 0); + } + + #[tokio::test] + async fn waiters_decrements_on_timeout() { + let limiter = WriteStreamLimiter::new(1); + let _permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let waiting = tokio::spawn({ + let limiter = limiter.clone(); + async move { + let _ = entered_tx.send(()); + limiter.acquire(Duration::from_millis(50)).await + } + }); + entered_rx.await.unwrap(); + tokio::time::timeout(Duration::from_millis(100), async { + while limiter.snapshot().waiters != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("waiter count should reach 1"); + + let result = waiting.await.unwrap(); + assert!(matches!(result, Err(AcquireError::Timeout))); + assert_eq!(limiter.snapshot().waiters, 0); + } + + #[tokio::test] + async fn waiters_decrements_on_cancellation() { + let limiter = WriteStreamLimiter::new(1); + let _permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let handle = tokio::spawn({ + let limiter = limiter.clone(); + async move { + let _ = entered_tx.send(()); + limiter.acquire(Duration::from_secs(60)).await + } + }); + + entered_rx.await.unwrap(); + tokio::time::timeout(Duration::from_millis(100), async { + while limiter.snapshot().waiters != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("waiter count should reach 1"); + + handle.abort(); + let _ = handle.await; + assert_eq!(limiter.snapshot().waiters, 0); + } + + #[tokio::test] + async fn clamps_out_of_range_values() { + let low = WriteStreamLimiter::new(0); + assert_eq!(low.max_permits(), MIN_WRITE_STREAM_MAX_CONCURRENCY); + let high = WriteStreamLimiter::new(10_000); + assert_eq!(high.max_permits(), MAX_WRITE_STREAM_MAX_CONCURRENCY); + } + + #[tokio::test] + async fn acquire_many_errors_when_requested_exceeds_capacity() { + let limiter = WriteStreamLimiter::new(3); + let err = limiter + .acquire_many(5, Duration::from_secs(1)) + .await + .unwrap_err(); + assert!(matches!( + err, + AcquireError::WouldExceedCapacity { + requested: 5, + max: 3 + } + )); + } + + #[tokio::test] + async fn split_one_divides_guard_into_single_permits() { + let limiter = WriteStreamLimiter::new(3); + let mut guard = limiter + .acquire_many(3, Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!(limiter.snapshot().available, 0); + + let p1 = guard.split_one().unwrap(); + let p2 = guard.split_one().unwrap(); + let p3 = guard.split_one().unwrap(); + assert!(guard.split_one().is_none()); + + // All permits still held; none released yet. + assert_eq!(limiter.snapshot().available, 0); + + drop(p1); + assert_eq!(limiter.snapshot().available, 1); + drop(p2); + assert_eq!(limiter.snapshot().available, 2); + drop(p3); + assert_eq!(limiter.snapshot().available, 3); + } + + #[tokio::test] + async fn split_one_on_zero_permits_returns_none() { + let limiter = WriteStreamLimiter::new(1); + let mut guard = limiter + .acquire_many(0, Duration::from_secs(1)) + .await + .unwrap(); + assert!(guard.split_one().is_none()); + assert_eq!(limiter.snapshot().available, 1); + } + + #[tokio::test] + async fn permit_moved_into_spawned_async_block_blocks_acquire() { + let limiter = WriteStreamLimiter::new(1); + let permit = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let handle = tokio::spawn(async move { + // This binding must be inside the async move block so the guard is + // held until the task completes. If it were dropped at the closure + // boundary, the second acquire below would succeed. + let _permit = permit; + let _ = started_tx.send(()); + tokio::time::sleep(Duration::from_millis(200)).await; + }); + + // Wait until the spawned task has definitely taken ownership of the permit. + started_rx.await.unwrap(); + + // The permit is still held, so a concurrent acquire must time out. + let err = limiter + .acquire(Duration::from_millis(50)) + .await + .unwrap_err(); + assert!(matches!(err, AcquireError::Timeout)); + + handle.await.unwrap(); + + // After the task drops the guard, the permit is released. + assert_eq!(limiter.snapshot().available, 1); + let permit2 = limiter.acquire(Duration::from_secs(1)).await.unwrap(); + drop(permit2); + } +} diff --git a/services/server/src/types.rs b/services/server/src/types.rs index 246b4005..2ea54668 100644 --- a/services/server/src/types.rs +++ b/services/server/src/types.rs @@ -5,7 +5,7 @@ use crate::alerts::AlertManager; use crate::engine::MemoryEngine; use crate::jobs::{BulkRememberJobStorage, RememberJobStorage, WalletJobStorage}; use crate::rate_limit::RateLimitConfig; -use crate::services::{Embedder, Extractor, Ranker}; +use crate::services::{Embedder, Extractor, Ranker, WriteStreamLimiter}; use crate::storage::db::VectorDb; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -125,6 +125,8 @@ pub struct AppState { pub blob_cache_max_bytes: usize, /// Redis TTL for recall query embedding cache entries. pub embedding_cache_ttl: std::time::Duration, + /// In-process concurrency limiter for write operations. + pub write_stream_limiter: Arc, } // ============================================================ @@ -229,6 +231,10 @@ pub struct Config { /// bypassing SEAL + Walrus. **Not for production.** Off by default; /// set `BENCHMARK_MODE=true` to enable. Surfaced via `GET /health`. pub benchmark_mode: bool, + /// Maximum concurrent active write-stream operations (prep + upload). + pub write_stream_max_concurrency: usize, + /// How long a handler waits for a write-stream permit before returning 429. + pub write_stream_acquire_timeout: std::time::Duration, } impl Config { @@ -299,6 +305,8 @@ impl Config { benchmark_mode: std::env::var("BENCHMARK_MODE") .map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes")) .unwrap_or(false), + write_stream_max_concurrency: parse_write_stream_max_concurrency(), + write_stream_acquire_timeout: parse_write_stream_acquire_timeout(), } } } @@ -333,6 +341,69 @@ fn parse_walrus_aggregator_urls(primary: &str, extra_csv: Option<&str>) -> Vec usize { + std::env::var("WRITE_STREAM_MAX_CONCURRENCY") + .ok() + .and_then(|v| { + let trimmed = v.trim(); + if trimmed.is_empty() { + return None; + } + match trimmed.parse::() { + Ok(n) => { + if !(1..=100).contains(&n) { + tracing::warn!( + "WRITE_STREAM_MAX_CONCURRENCY={} is out of range; clamping to {}", + trimmed, + n.clamp(1, 100) + ); + } + Some(n.clamp(1, 100)) + } + Err(_) => { + tracing::warn!( + "WRITE_STREAM_MAX_CONCURRENCY={} is invalid; using default 8", + trimmed + ); + None + } + } + }) + .unwrap_or(8) +} + +pub(crate) fn parse_write_stream_acquire_timeout() -> std::time::Duration { + let millis = std::env::var("WRITE_STREAM_ACQUIRE_TIMEOUT_MS") + .ok() + .and_then(|v| { + let trimmed = v.trim(); + if trimmed.is_empty() { + return None; + } + match trimmed.parse::() { + Ok(n) => { + if !(100..=60_000).contains(&n) { + tracing::warn!( + "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is out of range; clamping to {}", + trimmed, + n.clamp(100, 60_000) + ); + } + Some(n.clamp(100, 60_000)) + } + Err(_) => { + tracing::warn!( + "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is invalid; using default 5000", + trimmed + ); + None + } + } + }) + .unwrap_or(5_000); + std::time::Duration::from_millis(millis) +} + // ============================================================ // Sponsor Rate Limit Config // ============================================================