From a6cd319b6191451cf66798f344957e23ec2fc588 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 20:39:39 +0700 Subject: [PATCH 01/29] docs: design spec for WALM-116 relayer write-stream redesign --- .gitignore | 3 + ...026-06-15-relayer-write-stream-redesign.md | 285 ++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md diff --git a/.gitignore b/.gitignore index d252afbc..27475874 100644 --- a/.gitignore +++ b/.gitignore @@ -137,3 +137,6 @@ claudedocs/ # Temp File temp/ + +# Superpowers brainstorming session files +.superpowers/ 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..887f9868 --- /dev/null +++ b/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md @@ -0,0 +1,285 @@ +# 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 will result 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, +} + +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 will mark 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. Each item's prep task releases its own permit after enqueueing its wallet job. + +#### `/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 prep (embed + encrypt) each fact, enqueue one `WalletOperation::UploadAndTransfer` per fact, and release each fact's permit after its wallet job is enqueued. +- **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 will be 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. | +| 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 & 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. From dfc87c6eb8964722ef60650a104630c50b10a4b0 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 20:50:34 +0700 Subject: [PATCH 02/29] docs: implementation plan for WALM-116 relayer write-stream redesign --- .gitignore | 3 + ...026-06-15-relayer-write-stream-redesign.md | 1168 +++++++++++++++++ 2 files changed, 1171 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md diff --git a/.gitignore b/.gitignore index 27475874..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 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..a8ab584e --- /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 will be dropped 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). From eb0fff002afe86f11634e60995e1cfd4f92c258b Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 20:54:45 +0700 Subject: [PATCH 03/29] feat(write-stream): add WriteStreamLimiter with unit tests --- services/server/src/services/mod.rs | 2 + services/server/src/services/write_stream.rs | 164 +++++++++++++++++++ 2 files changed, 166 insertions(+) create mode 100644 services/server/src/services/write_stream.rs diff --git a/services/server/src/services/mod.rs b/services/server/src/services/mod.rs index 9cdc683e..c5333463 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, WriteStreamPermit}; diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs new file mode 100644 index 00000000..f303b592 --- /dev/null +++ b/services/server/src/services/write_stream.rs @@ -0,0 +1,164 @@ +//! 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. +#[derive(Debug)] +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); + } +} From be47a756910bc20eb8053abfeb5e723e38b99456 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 20:57:22 +0700 Subject: [PATCH 04/29] fix(write-stream): reject acquire_many requests exceeding capacity --- services/server/src/services/write_stream.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index f303b592..13221092 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -18,6 +18,7 @@ const MAX_WRITE_STREAM_MAX_CONCURRENCY: usize = 100; pub enum AcquireError { Timeout, Closed, + WouldExceedCapacity { requested: usize, max: usize }, } impl std::fmt::Display for AcquireError { @@ -25,6 +26,10 @@ impl std::fmt::Display for AcquireError { 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}" + ), } } } @@ -92,7 +97,12 @@ impl WriteStreamLimiter { permits: 0, }); } - let n = n.min(self.max_permits); + if n > self.max_permits { + return Err(AcquireError::WouldExceedCapacity { + requested: n, + max: self.max_permits, + }); + } let permit = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) .await .map_err(|_: Elapsed| AcquireError::Timeout)? @@ -161,4 +171,11 @@ mod tests { 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 })); + } } From f996151703a6c08b962eeb019a73d6bcd3cd677c Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:01:40 +0700 Subject: [PATCH 05/29] refactor(write-stream): quality cleanups from review --- services/server/src/services/write_stream.rs | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 13221092..786d1284 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -8,13 +8,12 @@ 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)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AcquireError { Timeout, Closed, @@ -38,6 +37,7 @@ 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, permits: usize, @@ -67,10 +67,6 @@ impl WriteStreamLimiter { } } - pub fn default_limiter() -> Self { - Self::new(DEFAULT_WRITE_STREAM_MAX_CONCURRENCY) - } - pub fn max_permits(&self) -> usize { self.max_permits } @@ -105,7 +101,7 @@ impl WriteStreamLimiter { } let permit = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) .await - .map_err(|_: Elapsed| AcquireError::Timeout)? + .map_err(|_| AcquireError::Timeout)? .map_err(|_| AcquireError::Closed)?; permit.forget(); Ok(WriteStreamPermit { @@ -115,6 +111,12 @@ impl WriteStreamLimiter { } } +impl Default for WriteStreamLimiter { + fn default() -> Self { + Self::new(DEFAULT_WRITE_STREAM_MAX_CONCURRENCY) + } +} + #[cfg(test)] mod tests { use super::*; From da7b26e32df0fcd88d674ee65d7ddd25491ddd71 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:03:55 +0700 Subject: [PATCH 06/29] feat(write-stream): wire WriteStreamLimiter into AppState and Config --- services/server/src/main.rs | 14 +++++++++++++- services/server/src/types.rs | 27 ++++++++++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/services/server/src/main.rs b/services/server/src/main.rs index be16ba41..272f02c6 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,15 @@ 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 +508,7 @@ async fn main() { blob_cache_ttl, blob_cache_max_bytes, embedding_cache_ttl, + write_stream_limiter, }); tracing::info!( diff --git a/services/server/src/types.rs b/services/server/src/types.rs index 246b4005..4bc3b675 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,23 @@ fn parse_walrus_aggregator_urls(primary: &str, extra_csv: Option<&str>) -> Vec 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) +} + // ============================================================ // Sponsor Rate Limit Config // ============================================================ From d79f649e4a6dd7d7c702c675f67726ce8a7b56d8 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:07:14 +0700 Subject: [PATCH 07/29] feat(write-stream): add Prometheus metrics for limiter state --- services/server/src/observability.rs | 56 ++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/services/server/src/observability.rs b/services/server/src/observability.rs index 64b56001..f2f2a50d 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,20 @@ 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_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 From e1f89d0b15f5a6f715f9d01ba224cf2d1ef6038c Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:10:22 +0700 Subject: [PATCH 08/29] feat(write-stream): add saturated helper using RateLimited --- services/server/src/routes/mod.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index 9cfc330d..b68bad14 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -182,6 +182,16 @@ pub(super) fn zip_search_hit_fields_onto_hydrated( } } +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(), + ) +} + #[cfg(test)] mod tests { use super::{collect_bounded_results, truncate_str}; From bd0edf4cd5f8db7da98a03bd45cfbe20066fd103 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:11:40 +0700 Subject: [PATCH 09/29] refactor(write-stream): remove unused AcquireError import --- services/server/src/routes/mod.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index b68bad14..d78fa62a 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -182,8 +182,6 @@ pub(super) fn zip_search_hit_fields_onto_hydrated( } } -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); From a4d933f286811db09f42409219ced3dda268ab99 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:14:49 +0700 Subject: [PATCH 10/29] feat(write-stream): gate remember and remember/bulk with limiter --- services/server/src/routes/remember.rs | 30 ++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index d39e1b0b..09e3b117 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -101,6 +101,7 @@ fn spawn_prepare_remember_job( owner: String, namespace: String, agent_public_key: String, + _permit: crate::services::write_stream::WriteStreamPermit, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -202,6 +203,7 @@ fn spawn_prepare_bulk_remember_job( owner: String, agent_public_key: String, pending_items: Vec, + _permits: crate::services::write_stream::WriteStreamPermit, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -647,6 +649,18 @@ pub async fn remember( .await .map_err(|e| AppError::Internal(format!("Failed to create job row: {}", e)))?; + // 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) => permit, + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/remember")); + } + }; + spawn_prepare_remember_job( Arc::clone(&state), job_id.clone(), @@ -654,6 +668,7 @@ pub async fn remember( owner_owned, namespace_owned, auth.public_key.clone(), + permit, ); tracing::info!( @@ -784,11 +799,24 @@ pub async fn remember_bulk( let total = job_ids.len(); + 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) => permits, + Err(_) => { + return Err(crate::routes::write_stream_saturated("/api/remember/bulk")); + } + }; + spawn_prepare_bulk_remember_job( Arc::clone(&state), owner.clone(), auth.public_key.clone(), pending_items, + permits, ); tracing::info!("remember_bulk accepted: {} items owner={}", total, owner,); @@ -1077,6 +1105,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), } } From 4209e4226a7f34399877d62998a547c6e764979c Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:20:23 +0700 Subject: [PATCH 11/29] fix(write-stream): acquire permits before row insert; distinguish AcquireError variants --- services/server/src/routes/remember.rs | 82 ++++++++++++++++++-------- 1 file changed, 56 insertions(+), 26 deletions(-) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 09e3b117..45d7be09 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -101,7 +101,8 @@ fn spawn_prepare_remember_job( owner: String, namespace: String, agent_public_key: String, - _permit: crate::services::write_stream::WriteStreamPermit, + // 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 { @@ -195,6 +196,9 @@ fn spawn_prepare_remember_job( } else { work.await; } + + // Hold the permit until prep hands off to the durable wallet queue. + let _ = permit; }); } @@ -203,7 +207,8 @@ fn spawn_prepare_bulk_remember_job( owner: String, agent_public_key: String, pending_items: Vec, - _permits: crate::services::write_stream::WriteStreamPermit, + // Held until prep completes; releases all permits on drop. + permits: crate::services::write_stream::WriteStreamPermit, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -333,6 +338,9 @@ fn spawn_prepare_bulk_remember_job( } else { work.await; } + + // Hold the permits until prep hands off to the durable bulk queue. + let _ = permits; }); } @@ -637,6 +645,26 @@ pub async fn remember( let namespace_owned = namespace.clone(); let text = body.text; + // Acquire a write-stream permit before persisting the job row so a + // saturated write stream never leaves orphaned `running` rows. + let permit = match state + .write_stream_limiter + .acquire(state.config.write_stream_acquire_timeout) + .await + { + Ok(permit) => permit, + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/remember")); + } + Err(crate::services::write_stream::AcquireError::Closed) => { + return Err(AppError::Internal("write stream limiter closed".into())); + } + Err(crate::services::write_stream::AcquireError::WouldExceedCapacity { .. }) => { + // n=1 can never exceed capacity + return Err(AppError::Internal("write stream capacity exceeded".into())); + } + }; + let job_id = uuid::Uuid::new_v4().to_string(); sqlx::query( @@ -649,18 +677,6 @@ pub async fn remember( .await .map_err(|e| AppError::Internal(format!("Failed to create job row: {}", e)))?; - // 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) => permit, - Err(_) => { - return Err(crate::routes::write_stream_saturated("/api/remember")); - } - }; - spawn_prepare_remember_job( Arc::clone(&state), job_id.clone(), @@ -773,6 +789,32 @@ pub async fn remember_bulk( &owner[..10.min(owner.len())], ); + // Acquire write-stream permits before persisting rows so a saturated + // write stream never leaves orphaned `running` rows. + let item_count = body.items.len(); + let permits = match state + .write_stream_limiter + .acquire_many(item_count, state.config.write_stream_acquire_timeout) + .await + { + Ok(permits) => 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, + }) => { + return Err(AppError::BadRequest(format!( + "Bulk request of {} items exceeds write stream capacity of {}; reduce batch size", + requested, max + ))); + } + Err(crate::services::write_stream::AcquireError::Closed) => { + return Err(AppError::Internal("write stream limiter closed".into())); + } + }; + let mut job_ids: Vec = Vec::with_capacity(body.items.len()); let mut pending_items: Vec = Vec::with_capacity(body.items.len()); @@ -799,18 +841,6 @@ pub async fn remember_bulk( let total = job_ids.len(); - 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) => permits, - Err(_) => { - return Err(crate::routes::write_stream_saturated("/api/remember/bulk")); - } - }; - spawn_prepare_bulk_remember_job( Arc::clone(&state), owner.clone(), From b0b2ae28a8da5346b01007d1ede34c67fa266f52 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:24:12 +0700 Subject: [PATCH 12/29] feat(write-stream): gate remember/manual and analyze with limiter --- services/server/src/routes/analyze.rs | 38 ++++++++++++++++++++++++++ services/server/src/routes/remember.rs | 16 +++++++++++ 2 files changed, 54 insertions(+) diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 6cc7f921..a65bac8a 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -384,6 +384,21 @@ 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) => permit, + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/analyze")); + } + Err(_) => { + 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 @@ -489,6 +504,29 @@ pub async fn analyze( } rate_limit::check_storage_quota(&state, owner, total_encrypted_bytes).await?; + // Acquire write-stream permits before inserting rows. Analyze rows are + // created with status='pending' and the stale-job sweeper does not clean + // them up, so we must not persist them unless the work can actually run. + let _permits = match state + .write_stream_limiter + .acquire_many(facts.len(), state.config.write_stream_acquire_timeout) + .await + { + Ok(permits) => 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 }) => { + return Err(AppError::BadRequest(format!( + "Analyze extracted {} facts, exceeding write stream capacity of {}; reduce input", + requested, max + ))); + } + Err(_) => { + return Err(AppError::Internal("write stream limiter unavailable".into())); + } + }; + // Step 3: For each prepared fact — insert remember_jobs row + enqueue WalletJob. // Round-robin across wallet pool so facts upload in parallel. let mut job_ids: Vec = Vec::with_capacity(prepared.len()); diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 45d7be09..e71a7888 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -971,6 +971,22 @@ 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) => permit, + Err(crate::services::write_stream::AcquireError::Timeout) => { + return Err(crate::routes::write_stream_saturated("/api/remember/manual")); + } + Err(_) => { + 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. From 1ca4f16d35f800ad17395314b3d3ed34e2be6e1b Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:32:16 +0700 Subject: [PATCH 13/29] feat(write-stream): gate wallet worker upload with limiter --- services/server/src/jobs.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index 7bb2eab0..14e2f251 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -908,6 +908,22 @@ 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) => permit, + Err(_) => { + // Limiter closed or timeout — leave job in queue for retry. + return Err(WalletJobError::Transient( + "write stream permit unavailable; will retry".into(), + )); + } + }; + // ── Upload to Walrus via sidecar (using pinned wallet_index) ─ let upload_result = crate::storage::walrus::upload_blob( &state.http_client, From 6600f3002d421ec63c43b79cf9f751651bba2b4d Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:42:10 +0700 Subject: [PATCH 14/29] feat(write-stream): emit limiter state and acquisition metrics --- services/server/src/jobs.rs | 5 ++- services/server/src/main.rs | 17 ++++++++ services/server/src/routes/analyze.rs | 10 ++++- services/server/src/routes/mod.rs | 5 +++ services/server/src/routes/remember.rs | 15 +++++-- services/server/src/services/mod.rs | 2 +- services/server/src/services/write_stream.rs | 43 ++++++++++++++++---- 7 files changed, 82 insertions(+), 15 deletions(-) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index 14e2f251..068ed56d 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -915,7 +915,10 @@ async fn execute_upload_and_transfer( .acquire(std::time::Duration::from_secs(60)) .await { - Ok(permit) => permit, + 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( diff --git a/services/server/src/main.rs b/services/server/src/main.rs index 272f02c6..ab9239e4 100644 --- a/services/server/src/main.rs +++ b/services/server/src/main.rs @@ -520,6 +520,23 @@ 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)); + 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 diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index a65bac8a..02d756ba 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -389,7 +389,10 @@ pub async fn analyze( .acquire(std::time::Duration::from_secs(60)) .await { - Ok(permit) => permit, + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } Err(crate::services::write_stream::AcquireError::Timeout) => { return Err(crate::routes::write_stream_saturated("/api/analyze")); } @@ -512,7 +515,10 @@ pub async fn analyze( .acquire_many(facts.len(), state.config.write_stream_acquire_timeout) .await { - Ok(permits) => permits, + Ok(permits) => { + crate::routes::record_write_stream_acquired_success(); + permits + } Err(crate::services::write_stream::AcquireError::Timeout) => { return Err(crate::routes::write_stream_saturated("/api/analyze")); } diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index d78fa62a..c121a6da 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -185,11 +185,16 @@ 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(), ) } +pub(super) fn record_write_stream_acquired_success() { + crate::observability::record_write_stream_acquired("success"); +} + #[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 e71a7888..46d9080a 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -652,7 +652,10 @@ pub async fn remember( .acquire(state.config.write_stream_acquire_timeout) .await { - Ok(permit) => permit, + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } Err(crate::services::write_stream::AcquireError::Timeout) => { return Err(crate::routes::write_stream_saturated("/api/remember")); } @@ -797,7 +800,10 @@ pub async fn remember_bulk( .acquire_many(item_count, state.config.write_stream_acquire_timeout) .await { - Ok(permits) => permits, + Ok(permits) => { + crate::routes::record_write_stream_acquired_success(); + permits + } Err(crate::services::write_stream::AcquireError::Timeout) => { return Err(crate::routes::write_stream_saturated("/api/remember/bulk")); } @@ -978,7 +984,10 @@ pub async fn remember_manual( .acquire(state.config.write_stream_acquire_timeout) .await { - Ok(permit) => permit, + Ok(permit) => { + crate::routes::record_write_stream_acquired_success(); + permit + } Err(crate::services::write_stream::AcquireError::Timeout) => { return Err(crate::routes::write_stream_saturated("/api/remember/manual")); } diff --git a/services/server/src/services/mod.rs b/services/server/src/services/mod.rs index c5333463..118ec5ae 100644 --- a/services/server/src/services/mod.rs +++ b/services/server/src/services/mod.rs @@ -37,4 +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, WriteStreamPermit}; +pub use write_stream::{WriteStreamLimiter, WriteStreamPermit, WriteStreamSnapshot}; diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 786d1284..207b6462 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -5,6 +5,7 @@ //! /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; @@ -49,11 +50,20 @@ impl Drop for WriteStreamPermit { } } +/// Snapshot of the write-stream limiter state. +#[derive(Clone, Copy, Debug)] +pub struct WriteStreamSnapshot { + pub total: usize, + pub available: usize, + pub waiters: usize, +} + /// In-process concurrency limiter for the write stream. #[derive(Clone, Debug)] pub struct WriteStreamLimiter { semaphore: Arc, max_permits: usize, + waiters: Arc, } impl WriteStreamLimiter { @@ -64,6 +74,7 @@ impl WriteStreamLimiter { Self { semaphore: Arc::new(Semaphore::new(max_permits)), max_permits, + waiters: Arc::new(AtomicUsize::new(0)), } } @@ -75,6 +86,15 @@ impl WriteStreamLimiter { self.semaphore.available_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 @@ -99,15 +119,22 @@ impl WriteStreamLimiter { max: self.max_permits, }); } - let permit = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) + self.waiters.fetch_add(1, Ordering::Relaxed); + let result = tokio::time::timeout(timeout, self.semaphore.acquire_many(n as u32)) .await - .map_err(|_| AcquireError::Timeout)? - .map_err(|_| AcquireError::Closed)?; - permit.forget(); - Ok(WriteStreamPermit { - semaphore: Arc::clone(&self.semaphore), - permits: n, - }) + .map_err(|_| 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), + } } } From 91ebd61207268910be375599b434aab4aa63fad2 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:50:23 +0700 Subject: [PATCH 15/29] fix(write-stream): guard waiter count and clean lint warnings --- services/server/src/services/mod.rs | 2 +- services/server/src/services/write_stream.rs | 98 +++++++++++++++++--- 2 files changed, 85 insertions(+), 15 deletions(-) diff --git a/services/server/src/services/mod.rs b/services/server/src/services/mod.rs index 118ec5ae..d7782397 100644 --- a/services/server/src/services/mod.rs +++ b/services/server/src/services/mod.rs @@ -37,4 +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, WriteStreamPermit, WriteStreamSnapshot}; +pub use write_stream::WriteStreamLimiter; diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 207b6462..7712a3cf 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -58,6 +58,17 @@ pub struct WriteStreamSnapshot { pub waiters: usize, } +/// Decrements the waiter counter on drop, guarding against future cancellation. +struct WaiterGuard { + waiters: Arc, +} + +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 { @@ -68,9 +79,7 @@ pub struct WriteStreamLimiter { 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); + 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, @@ -82,10 +91,6 @@ impl WriteStreamLimiter { self.max_permits } - pub fn available_permits(&self) -> usize { - self.semaphore.available_permits() - } - /// Return a point-in-time snapshot of limiter state. pub fn snapshot(&self) -> WriteStreamSnapshot { WriteStreamSnapshot { @@ -120,11 +125,13 @@ impl WriteStreamLimiter { }); } self.waiters.fetch_add(1, Ordering::Relaxed); + let _waiter_guard = WaiterGuard { + waiters: 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)); - self.waiters.fetch_sub(1, Ordering::Relaxed); match result { Ok(permit) => { permit.forget(); @@ -152,9 +159,9 @@ mod tests { 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); + assert_eq!(limiter.snapshot().available, 0); drop(permit); - assert_eq!(limiter.available_permits(), 1); + assert_eq!(limiter.snapshot().available, 1); } #[tokio::test] @@ -178,11 +185,11 @@ mod tests { .await .unwrap_err(); assert!(matches!(err, AcquireError::Timeout)); - assert_eq!(limiter.available_permits(), 2); + assert_eq!(limiter.snapshot().available, 2); let p23 = limiter.acquire_many(2, Duration::from_secs(1)).await.unwrap(); - assert_eq!(limiter.available_permits(), 0); + assert_eq!(limiter.snapshot().available, 0); drop(p23); - assert_eq!(limiter.available_permits(), 2); + assert_eq!(limiter.snapshot().available, 2); } #[tokio::test] @@ -190,7 +197,70 @@ mod tests { 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); + 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 waiting = tokio::spawn({ + let limiter = limiter.clone(); + async move { + limiter.acquire(Duration::from_millis(200)).await + } + }); + + // Give the spawned task time to enter the wait. + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(limiter.snapshot().waiters, 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 waiting = tokio::spawn({ + let limiter = limiter.clone(); + async move { + limiter.acquire(Duration::from_millis(50)).await + } + }); + + // Wait long enough for the spawned task to be awaiting but not time out yet. + tokio::time::sleep(Duration::from_millis(25)).await; + assert_eq!(limiter.snapshot().waiters, 1); + + let result = waiting.await.unwrap(); + assert!(matches!(result, Err(AcquireError::Timeout))); + assert_eq!(limiter.snapshot().waiters, 0); } #[tokio::test] From e15ad6aa0a23f7f1e5309af6eeae2452bd88a1da Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 21:57:03 +0700 Subject: [PATCH 16/29] fix(write-stream): address Task 8 review feedback --- services/server/src/jobs.rs | 10 +++- services/server/src/main.rs | 1 + services/server/src/services/write_stream.rs | 51 ++++++++++++++++---- 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index 068ed56d..f0189235 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -919,12 +919,18 @@ async fn execute_upload_and_transfer( crate::routes::record_write_stream_acquired_success(); permit } - Err(_) => { - // Limiter closed or timeout — leave job in queue for retry. + 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(_) => { + // 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) ─ diff --git a/services/server/src/main.rs b/services/server/src/main.rs index ab9239e4..d81dbe3e 100644 --- a/services/server/src/main.rs +++ b/services/server/src/main.rs @@ -525,6 +525,7 @@ async fn main() { 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(); diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 7712a3cf..e1e2bdb3 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -63,6 +63,13 @@ 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); @@ -124,10 +131,7 @@ impl WriteStreamLimiter { max: self.max_permits, }); } - self.waiters.fetch_add(1, Ordering::Relaxed); - let _waiter_guard = WaiterGuard { - waiters: Arc::clone(&self.waiters), - }; + 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) @@ -226,15 +230,17 @@ mod tests { 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 } }); - - // Give the spawned task time to enter the wait. - tokio::time::sleep(Duration::from_millis(50)).await; + entered_rx.await.unwrap(); + tokio::task::yield_now().await; // let the spawned task reach the semaphore wait assert_eq!(limiter.snapshot().waiters, 1); let result = waiting.await.unwrap(); @@ -247,15 +253,16 @@ mod tests { 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 } }); - - // Wait long enough for the spawned task to be awaiting but not time out yet. - tokio::time::sleep(Duration::from_millis(25)).await; + entered_rx.await.unwrap(); + tokio::task::yield_now().await; assert_eq!(limiter.snapshot().waiters, 1); let result = waiting.await.unwrap(); @@ -263,6 +270,30 @@ mod tests { 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::task::yield_now().await; + assert_eq!(limiter.snapshot().waiters, 1); + + handle.abort(); + // Wait a bit for the abort to propagate. + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(limiter.snapshot().waiters, 0); + } + #[tokio::test] async fn clamps_out_of_range_values() { let low = WriteStreamLimiter::new(0); From 98d517bdd00bae652611eb1c6c150de36274ef12 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:02:32 +0700 Subject: [PATCH 17/29] refactor(write-stream): move metric helper to observability and harden tests --- services/server/src/jobs.rs | 2 +- services/server/src/observability.rs | 4 +++ services/server/src/routes/analyze.rs | 4 +-- services/server/src/routes/mod.rs | 4 --- services/server/src/routes/remember.rs | 6 ++-- services/server/src/services/write_stream.rs | 30 ++++++++++++++------ 6 files changed, 32 insertions(+), 18 deletions(-) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index f0189235..c3683456 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -916,7 +916,7 @@ async fn execute_upload_and_transfer( .await { Ok(permit) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permit } Err(crate::services::write_stream::AcquireError::Timeout) => { diff --git a/services/server/src/observability.rs b/services/server/src/observability.rs index f2f2a50d..98545c10 100644 --- a/services/server/src/observability.rs +++ b/services/server/src/observability.rs @@ -618,6 +618,10 @@ 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_rejected(route: &str) { WRITE_STREAM_REJECTED_TOTAL.with_label_values(&[route]).inc(); } diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 02d756ba..9095de2f 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -390,7 +390,7 @@ pub async fn analyze( .await { Ok(permit) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permit } Err(crate::services::write_stream::AcquireError::Timeout) => { @@ -516,7 +516,7 @@ pub async fn analyze( .await { Ok(permits) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permits } Err(crate::services::write_stream::AcquireError::Timeout) => { diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index c121a6da..1ad6b6a8 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -191,10 +191,6 @@ pub(super) fn write_stream_saturated(route: &str) -> AppError { ) } -pub(super) fn record_write_stream_acquired_success() { - crate::observability::record_write_stream_acquired("success"); -} - #[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 46d9080a..cc11cc7c 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -653,7 +653,7 @@ pub async fn remember( .await { Ok(permit) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permit } Err(crate::services::write_stream::AcquireError::Timeout) => { @@ -801,7 +801,7 @@ pub async fn remember_bulk( .await { Ok(permits) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permits } Err(crate::services::write_stream::AcquireError::Timeout) => { @@ -985,7 +985,7 @@ pub async fn remember_manual( .await { Ok(permit) => { - crate::routes::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success(); permit } Err(crate::services::write_stream::AcquireError::Timeout) => { diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index e1e2bdb3..0bb59bd2 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -240,8 +240,13 @@ mod tests { } }); entered_rx.await.unwrap(); - tokio::task::yield_now().await; // let the spawned task reach the semaphore wait - assert_eq!(limiter.snapshot().waiters, 1); + 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))); @@ -262,8 +267,13 @@ mod tests { } }); entered_rx.await.unwrap(); - tokio::task::yield_now().await; - assert_eq!(limiter.snapshot().waiters, 1); + 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))); @@ -285,12 +295,16 @@ mod tests { }); entered_rx.await.unwrap(); - tokio::task::yield_now().await; - assert_eq!(limiter.snapshot().waiters, 1); + 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(); - // Wait a bit for the abort to propagate. - tokio::time::sleep(Duration::from_millis(50)).await; + let _ = handle.await; assert_eq!(limiter.snapshot().waiters, 0); } From 4e1e279270bf1dbc43c53c3a17386b2e262042e0 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:08:13 +0700 Subject: [PATCH 18/29] feat(sidecar): raise WALRUS_UPLOAD_MAX_CONCURRENCY safety-net floor to 12 --- services/server/scripts/sidecar/config.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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, ); From 1f740c03af8dc24919a452b03f51bd3c211698f3 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:12:19 +0700 Subject: [PATCH 19/29] test(write-stream): add saturation unit test --- services/server/src/routes/remember.rs | 28 ++++++++++++++-- services/server/src/services/write_stream.rs | 35 +++++++++++++++++--- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index cc11cc7c..9e93e8c2 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -989,10 +989,14 @@ pub async fn remember_manual( permit } Err(crate::services::write_stream::AcquireError::Timeout) => { - return Err(crate::routes::write_stream_saturated("/api/remember/manual")); + return Err(crate::routes::write_stream_saturated( + "/api/remember/manual", + )); } Err(_) => { - return Err(AppError::Internal("write stream limiter unavailable".into())); + return Err(AppError::Internal( + "write stream limiter unavailable".into(), + )); } }; @@ -1211,4 +1215,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/write_stream.rs b/services/server/src/services/write_stream.rs index 0bb59bd2..aeafc5d9 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -86,7 +86,10 @@ pub struct WriteStreamLimiter { 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); + 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, @@ -155,6 +158,13 @@ impl Default for WriteStreamLimiter { } } +#[cfg(test)] +impl WriteStreamLimiter { + pub fn test_new(max_permits: usize) -> Self { + Self::new(max_permits) + } +} + #[cfg(test)] mod tests { use super::*; @@ -190,7 +200,10 @@ mod tests { .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(); + 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); @@ -199,7 +212,10 @@ mod tests { #[tokio::test] async fn zero_permits_noop() { let limiter = WriteStreamLimiter::new(1); - let guard = limiter.acquire_many(0, Duration::from_secs(1)).await.unwrap(); + let guard = limiter + .acquire_many(0, Duration::from_secs(1)) + .await + .unwrap(); assert_eq!(guard.permits, 0); assert_eq!(limiter.snapshot().available, 1); } @@ -319,7 +335,16 @@ mod tests { #[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 })); + let err = limiter + .acquire_many(5, Duration::from_secs(1)) + .await + .unwrap_err(); + assert!(matches!( + err, + AcquireError::WouldExceedCapacity { + requested: 5, + max: 3 + } + )); } } From f2c9da6e9c72d9083bb83c4e899426041fefcf7e Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:17:25 +0700 Subject: [PATCH 20/29] docs(relayer): document write-stream limiter config and metrics --- docs/relayer/observability.md | 5 +++++ docs/relayer/self-hosting.md | 9 +++++++++ 2 files changed, 14 insertions(+) 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` From 0285c8d48aa57541b7afbcf6b0a851d4ea5eb240 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:29:18 +0700 Subject: [PATCH 21/29] fix(write-stream): align permit acquisition with spec and split batch permits --- services/server/src/jobs.rs | 1 + services/server/src/routes/analyze.rs | 69 ++++++---- services/server/src/routes/remember.rs | 128 +++++++++++-------- services/server/src/services/write_stream.rs | 47 +++++++ 4 files changed, 166 insertions(+), 79 deletions(-) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index c3683456..571de42d 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -926,6 +926,7 @@ async fn execute_upload_and_transfer( )); } 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(), diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 9095de2f..dd1c32b1 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -397,6 +397,7 @@ pub async fn analyze( 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(), )); @@ -459,10 +460,41 @@ 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(); + 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_acquired("failure"); + return Err(AppError::BadRequest(format!( + "Analyze extracted {} facts, exceeding write stream capacity of {}; reduce input", + requested, max + ))); + } + 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() @@ -470,6 +502,9 @@ 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 { let embed_fut = state.embedder.embed(&fact.text); let encrypt_fut = crate::storage::seal::seal_encrypt( @@ -484,12 +519,16 @@ pub async fn analyze( // carry `importance` through the prep tuple so // the job payload below can persist it alongside the // ciphertext + vector. - Ok::<_, AppError>(( + let result = Ok::<_, AppError>(( fact.text, fact.importance, vector_result?, encrypted_result?, - )) + )); + // Hold the permit until this fact's prep hands off to the + // durable wallet queue (or fails). + let _ = permit; + result } }) .collect(); @@ -507,32 +546,6 @@ pub async fn analyze( } rate_limit::check_storage_quota(&state, owner, total_encrypted_bytes).await?; - // Acquire write-stream permits before inserting rows. Analyze rows are - // created with status='pending' and the stale-job sweeper does not clean - // them up, so we must not persist them unless the work can actually run. - let _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(); - 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 }) => { - return Err(AppError::BadRequest(format!( - "Analyze extracted {} facts, exceeding write stream capacity of {}; reduce input", - requested, max - ))); - } - Err(_) => { - return Err(AppError::Internal("write stream limiter unavailable".into())); - } - }; - // Step 3: For each prepared fact — insert remember_jobs row + enqueue WalletJob. // Round-robin across wallet pool so facts upload in parallel. let mut job_ids: Vec = Vec::with_capacity(prepared.len()); diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 9e93e8c2..ab4e9875 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -207,8 +207,9 @@ fn spawn_prepare_bulk_remember_job( owner: String, agent_public_key: String, pending_items: Vec, - // Held until prep completes; releases all permits on drop. - permits: crate::services::write_stream::WriteStreamPermit, + // One guard per item; each is held until that item's prep completes and + // releases its own slot. + mut permits: Vec, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -223,6 +224,9 @@ fn spawn_prepare_bulk_remember_job( .map(|item| { let state = Arc::clone(&state); let owner = owner.clone(); + let permit = permits + .pop() + .expect("permits length matches pending_items length"); async move { // bulk items can carry up to MAX_REMEMBER_TEXT_BYTES // each, so the same summarize-before-embed rule applies here. @@ -257,12 +261,16 @@ fn spawn_prepare_bulk_remember_job( ); let (vector_result, encrypted_result) = tokio::join!(embed_fut, encrypt_fut); - Ok::<_, AppError>(( + let result = Ok::<_, AppError>(( item.job_id, item.namespace, vector_result?, encrypted_result?, - )) + )); + // Hold the permit until this item's prep hands off + // to the durable bulk queue (or fails). + let _ = permit; + result } }) .collect(); @@ -338,9 +346,6 @@ fn spawn_prepare_bulk_remember_job( } else { work.await; } - - // Hold the permits until prep hands off to the durable bulk queue. - let _ = permits; }); } @@ -645,8 +650,23 @@ pub async fn remember( let namespace_owned = namespace.clone(); let text = body.text; - // Acquire a write-stream permit before persisting the job row so a - // saturated write stream never leaves orphaned `running` rows. + // 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( + "INSERT INTO remember_jobs (id, owner, namespace, status) VALUES ($1, $2, $3, 'running')", + ) + .bind(&job_id) + .bind(owner) + .bind(namespace) + .execute(state.db.pool()) + .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) @@ -660,26 +680,16 @@ pub async fn remember( 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_acquired("failure"); // n=1 can never exceed capacity return Err(AppError::Internal("write stream capacity exceeded".into())); } }; - let job_id = uuid::Uuid::new_v4().to_string(); - - sqlx::query( - "INSERT INTO remember_jobs (id, owner, namespace, status) VALUES ($1, $2, $3, 'running')", - ) - .bind(&job_id) - .bind(owner) - .bind(namespace) - .execute(state.db.pool()) - .await - .map_err(|e| AppError::Internal(format!("Failed to create job row: {}", e)))?; - spawn_prepare_remember_job( Arc::clone(&state), job_id.clone(), @@ -792,35 +802,8 @@ pub async fn remember_bulk( &owner[..10.min(owner.len())], ); - // Acquire write-stream permits before persisting rows so a saturated - // write stream never leaves orphaned `running` rows. - let item_count = body.items.len(); - let 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(); - 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, - }) => { - return Err(AppError::BadRequest(format!( - "Bulk request of {} items exceeds write stream capacity of {}; reduce batch size", - requested, max - ))); - } - Err(crate::services::write_stream::AcquireError::Closed) => { - return Err(AppError::Internal("write stream limiter closed".into())); - } - }; - + // 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()); @@ -847,12 +830,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 item_count = pending_items.len(); + 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(); + 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_acquired("failure"); + return Err(AppError::BadRequest(format!( + "Bulk request of {} items exceeds write stream capacity of {}; reduce batch size", + requested, max + ))); + } + 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, - permits, + item_permits, ); tracing::info!("remember_bulk accepted: {} items owner={}", total, owner,); @@ -994,6 +1019,7 @@ pub async fn remember_manual( )); } Err(_) => { + crate::observability::record_write_stream_acquired("failure"); return Err(AppError::Internal( "write stream limiter unavailable".into(), )); diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index aeafc5d9..6c8044f9 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -50,6 +50,23 @@ impl Drop for WriteStreamPermit { } } +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 { @@ -347,4 +364,34 @@ mod tests { } )); } + + #[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); + } } From 1212fbf935cd1dbe88e59b09fa974f35ac62ecb9 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:38:02 +0700 Subject: [PATCH 22/29] fix(write-stream): map capacity errors to 429 and correct permit lifetimes --- services/server/src/routes/analyze.rs | 17 ++++++++++------- services/server/src/routes/remember.rs | 26 +++++++++++++------------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index dd1c32b1..3c238331 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -476,11 +476,15 @@ pub async fn analyze( 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 }) => { + 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::BadRequest(format!( - "Analyze extracted {} facts, exceeding write stream capacity of {}; reduce input", - requested, max + return Err(AppError::RateLimited(format!( + "Analyze extracted {} facts, exceeding write stream capacity; reduce input", + requested ))); } Err(_) => { @@ -502,7 +506,7 @@ pub async fn analyze( let state = Arc::clone(&state); let owner = owner.clone(); let fact = fact.clone(); - let permit = permits + let _permit = permits .split_one() .expect("acquired permits equal facts.len()"); async move { @@ -525,9 +529,8 @@ pub async fn analyze( vector_result?, encrypted_result?, )); - // Hold the permit until this fact's prep hands off to the + // `_permit` is held until this fact's prep hands off to the // durable wallet queue (or fails). - let _ = permit; result } }) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index ab4e9875..6907465c 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -102,7 +102,7 @@ fn spawn_prepare_remember_job( namespace: String, agent_public_key: String, // Held until prep completes; releases the permit on drop. - permit: crate::services::write_stream::WriteStreamPermit, + _permit: crate::services::write_stream::WriteStreamPermit, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -196,9 +196,6 @@ fn spawn_prepare_remember_job( } else { work.await; } - - // Hold the permit until prep hands off to the durable wallet queue. - let _ = permit; }); } @@ -224,7 +221,7 @@ fn spawn_prepare_bulk_remember_job( .map(|item| { let state = Arc::clone(&state); let owner = owner.clone(); - let permit = permits + let _permit = permits .pop() .expect("permits length matches pending_items length"); async move { @@ -267,9 +264,8 @@ fn spawn_prepare_bulk_remember_job( vector_result?, encrypted_result?, )); - // Hold the permit until this item's prep hands off + // `_permit` is held until this item's prep hands off // to the durable bulk queue (or fails). - let _ = permit; result } }) @@ -684,9 +680,12 @@ pub async fn remember( 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 - return Err(AppError::Internal("write stream capacity exceeded".into())); + // 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(), + )); } }; @@ -848,12 +847,13 @@ pub async fn remember_bulk( } Err(crate::services::write_stream::AcquireError::WouldExceedCapacity { requested, - max, + max: _, }) => { + crate::observability::record_write_stream_rejected("/api/remember/bulk"); crate::observability::record_write_stream_acquired("failure"); - return Err(AppError::BadRequest(format!( - "Bulk request of {} items exceeds write stream capacity of {}; reduce batch size", - requested, max + return Err(AppError::RateLimited(format!( + "Bulk request of {} items exceeds write stream capacity; reduce batch size", + requested ))); } Err(crate::services::write_stream::AcquireError::Closed) => { From 13a29c019f3ef1302c9085ad283e732c54985dde Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:47:06 +0700 Subject: [PATCH 23/29] fix(write-stream): count permits per-item and document spec deviations --- ...026-06-15-relayer-write-stream-redesign.md | 12 ++++- services/server/src/observability.rs | 9 ++++ services/server/src/routes/analyze.rs | 2 +- services/server/src/routes/remember.rs | 2 +- services/server/src/types.rs | 54 +++++++++++++++++-- 5 files changed, 71 insertions(+), 8 deletions(-) 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 index 887f9868..401aaaa1 100644 --- a/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md +++ b/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md @@ -84,6 +84,12 @@ pub struct WriteStreamLimiter { 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; @@ -118,7 +124,8 @@ Initialize in `main.rs` after config load. 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. Each item's prep task releases its own permit after enqueueing its wallet job. +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` @@ -129,7 +136,7 @@ Initialize in `main.rs` after config load. #### `/api/analyze` -- **Production path**: after fact extraction, attempt to acquire `facts.len()` permits with a bounded timeout. If acquisition fails, return `429`. Then prep (embed + encrypt) each fact, enqueue one `WalletOperation::UploadAndTransfer` per fact, and release each fact's permit after its wallet job is enqueued. +- **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 @@ -197,6 +204,7 @@ Client POST /api/remember/bulk with N items | 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. | diff --git a/services/server/src/observability.rs b/services/server/src/observability.rs index 98545c10..31e8075a 100644 --- a/services/server/src/observability.rs +++ b/services/server/src/observability.rs @@ -622,6 +622,15 @@ 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(); } diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 3c238331..ac203db0 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -470,7 +470,7 @@ pub async fn analyze( .await { Ok(permits) => { - crate::observability::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success_n(facts.len()); permits } Err(crate::services::write_stream::AcquireError::Timeout) => { diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 6907465c..78fe7c30 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -839,7 +839,7 @@ pub async fn remember_bulk( .await { Ok(permits) => { - crate::observability::record_write_stream_acquired_success(); + crate::observability::record_write_stream_acquired_success_n(item_count); permits } Err(crate::services::write_stream::AcquireError::Timeout) => { diff --git a/services/server/src/types.rs b/services/server/src/types.rs index 4bc3b675..dfe8d210 100644 --- a/services/server/src/types.rs +++ b/services/server/src/types.rs @@ -344,16 +344,62 @@ fn parse_walrus_aggregator_urls(primary: &str, extra_csv: Option<&str>) -> Vec usize { std::env::var("WRITE_STREAM_MAX_CONCURRENCY") .ok() - .and_then(|v| v.trim().parse::().ok()) - .map(|v| v.clamp(1, 100)) + .and_then(|v| { + let trimmed = v.trim(); + if trimmed.is_empty() { + return None; + } + match trimmed.parse::() { + Ok(n) => { + if n < 1 || n > 100 { + tracing::warn!( + "WRITE_STREAM_MAX_CONCURRENCY={} is out of range; clamping to {}", + v, + n.clamp(1, 100) + ); + } + Some(n.clamp(1, 100)) + } + Err(_) => { + tracing::warn!( + "WRITE_STREAM_MAX_CONCURRENCY={} is invalid; using default 8", + v + ); + 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| v.trim().parse::().ok()) - .map(|v| v.clamp(100, 60_000)) + .and_then(|v| { + let trimmed = v.trim(); + if trimmed.is_empty() { + return None; + } + match trimmed.parse::() { + Ok(n) => { + if n < 100 || n > 60_000 { + tracing::warn!( + "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is out of range; clamping to {}", + v, + n.clamp(100, 60_000) + ); + } + Some(n.clamp(100, 60_000)) + } + Err(_) => { + tracing::warn!( + "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is invalid; using default 5000", + v + ); + None + } + } + }) .unwrap_or(5_000); std::time::Duration::from_millis(millis) } From cbfc15174cf93a0d4efb58c51b0978d91a2702bd Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 22:53:39 +0700 Subject: [PATCH 24/29] fix(write-stream): hold single remember permit through prep and silence new clippy lint --- services/server/src/routes/remember.rs | 4 +++- services/server/src/types.rs | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 78fe7c30..5439d8c0 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -102,10 +102,12 @@ fn spawn_prepare_remember_job( namespace: String, agent_public_key: String, // Held until prep completes; releases the permit on drop. - _permit: crate::services::write_stream::WriteStreamPermit, + 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 diff --git a/services/server/src/types.rs b/services/server/src/types.rs index dfe8d210..46f0aec9 100644 --- a/services/server/src/types.rs +++ b/services/server/src/types.rs @@ -351,7 +351,7 @@ pub(crate) fn parse_write_stream_max_concurrency() -> usize { } match trimmed.parse::() { Ok(n) => { - if n < 1 || n > 100 { + if !(1..=100).contains(&n) { tracing::warn!( "WRITE_STREAM_MAX_CONCURRENCY={} is out of range; clamping to {}", v, @@ -382,7 +382,7 @@ pub(crate) fn parse_write_stream_acquire_timeout() -> std::time::Duration { } match trimmed.parse::() { Ok(n) => { - if n < 100 || n > 60_000 { + if !(100..=60_000).contains(&n) { tracing::warn!( "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is out of range; clamping to {}", v, From 08190dcd949e7318a063f512b2bf6be8a4c6dc2d Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 23:02:14 +0700 Subject: [PATCH 25/29] feat(write-stream): gate legacy execute_remember upload path --- services/server/src/jobs.rs | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index 571de42d..917afc1f 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -1840,6 +1840,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, From 862c5cd65405004cdfd26c1cd08f532db21ba462 Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 23:07:23 +0700 Subject: [PATCH 26/29] chore(write-stream): minor cleanups from final review --- services/server/src/routes/remember.rs | 8 +++----- services/server/src/services/write_stream.rs | 2 +- services/server/src/types.rs | 8 ++++---- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 5439d8c0..e59c1670 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -208,7 +208,7 @@ fn spawn_prepare_bulk_remember_job( pending_items: Vec, // One guard per item; each is held until that item's prep completes and // releases its own slot. - mut permits: Vec, + item_permits: Vec, ) { let request_context = crate::observability::current_context(); tokio::spawn(async move { @@ -220,12 +220,10 @@ 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(); - let _permit = permits - .pop() - .expect("permits length matches pending_items length"); async move { // bulk items can carry up to MAX_REMEMBER_TEXT_BYTES // each, so the same summarize-before-embed rule applies here. diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 6c8044f9..84c0b1cd 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -41,7 +41,7 @@ impl std::error::Error for AcquireError {} #[must_use = "permit releases on drop"] pub struct WriteStreamPermit { semaphore: Arc, - permits: usize, + pub(crate) permits: usize, } impl Drop for WriteStreamPermit { diff --git a/services/server/src/types.rs b/services/server/src/types.rs index 46f0aec9..2ea54668 100644 --- a/services/server/src/types.rs +++ b/services/server/src/types.rs @@ -354,7 +354,7 @@ pub(crate) fn parse_write_stream_max_concurrency() -> usize { if !(1..=100).contains(&n) { tracing::warn!( "WRITE_STREAM_MAX_CONCURRENCY={} is out of range; clamping to {}", - v, + trimmed, n.clamp(1, 100) ); } @@ -363,7 +363,7 @@ pub(crate) fn parse_write_stream_max_concurrency() -> usize { Err(_) => { tracing::warn!( "WRITE_STREAM_MAX_CONCURRENCY={} is invalid; using default 8", - v + trimmed ); None } @@ -385,7 +385,7 @@ pub(crate) fn parse_write_stream_acquire_timeout() -> std::time::Duration { if !(100..=60_000).contains(&n) { tracing::warn!( "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is out of range; clamping to {}", - v, + trimmed, n.clamp(100, 60_000) ); } @@ -394,7 +394,7 @@ pub(crate) fn parse_write_stream_acquire_timeout() -> std::time::Duration { Err(_) => { tracing::warn!( "WRITE_STREAM_ACQUIRE_TIMEOUT_MS={} is invalid; using default 5000", - v + trimmed ); None } From 248733f2b83ab8e8f3e4ab06593a7add284d012c Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 23:18:37 +0700 Subject: [PATCH 27/29] fix(write-stream): fast-fail oversized bulk requests and run cargo fmt --- services/server/src/alerts.rs | 12 +++++++----- services/server/src/jobs.rs | 20 ++++++++------------ services/server/src/main.rs | 18 ++++++++---------- services/server/src/observability.rs | 8 ++++++-- services/server/src/routes/analyze.rs | 4 +++- services/server/src/routes/remember.rs | 13 +++++++++++-- services/server/src/services/extractor.rs | 13 ++++++++++--- services/server/src/services/write_stream.rs | 10 ++++++++-- 8 files changed, 61 insertions(+), 37 deletions(-) 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 917afc1f..3313f39e 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -1475,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 @@ -2205,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, }; @@ -2344,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", @@ -2490,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); } @@ -2498,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 d81dbe3e..88d41816 100644 --- a/services/server/src/main.rs +++ b/services/server/src/main.rs @@ -441,9 +441,8 @@ 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, - )); + 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(), @@ -546,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!( @@ -562,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; @@ -623,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 31e8075a..aab80f66 100644 --- a/services/server/src/observability.rs +++ b/services/server/src/observability.rs @@ -615,7 +615,9 @@ pub fn observe_write_stream_state(total: usize, available: usize, waiters: usize } pub fn record_write_stream_acquired(result: &str) { - WRITE_STREAM_ACQUIRED_TOTAL.with_label_values(&[result]).inc(); + WRITE_STREAM_ACQUIRED_TOTAL + .with_label_values(&[result]) + .inc(); } pub fn record_write_stream_acquired_success() { @@ -632,7 +634,9 @@ pub fn record_write_stream_acquired_success_n(n: usize) { } pub fn record_write_stream_rejected(route: &str) { - WRITE_STREAM_REJECTED_TOTAL.with_label_values(&[route]).inc(); + WRITE_STREAM_REJECTED_TOTAL + .with_label_values(&[route]) + .inc(); } fn record_http_request(method: &str, route: &str, status: StatusCode, elapsed: Duration) { diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index ac203db0..8683ea00 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -489,7 +489,9 @@ pub async fn analyze( } Err(_) => { crate::observability::record_write_stream_acquired("failure"); - return Err(AppError::Internal("write stream limiter unavailable".into())); + return Err(AppError::Internal( + "write stream limiter unavailable".into(), + )); } }; diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index e59c1670..c8887a87 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -794,10 +794,20 @@ pub async fn remember_bulk( } } + let item_count = body.items.len(); + if item_count > state.config.write_stream_max_concurrency { + 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", + item_count + ))); + } + let owner = &auth.owner; tracing::info!( "remember_bulk: {} items owner={}", - body.items.len(), + item_count, &owner[..10.min(owner.len())], ); @@ -832,7 +842,6 @@ pub async fn remember_bulk( // 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 item_count = pending_items.len(); let mut permits = match state .write_stream_limiter .acquire_many(item_count, state.config.write_stream_acquire_timeout) 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/write_stream.rs b/services/server/src/services/write_stream.rs index 84c0b1cd..6c40771a 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -368,7 +368,10 @@ mod tests { #[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(); + 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(); @@ -390,7 +393,10 @@ mod tests { #[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(); + 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); } From b140186083b34991b8c95d064796e646c7255b7a Mon Sep 17 00:00:00 2001 From: DalenMax Date: Mon, 15 Jun 2026 23:28:04 +0700 Subject: [PATCH 28/29] fix(write-stream): move per-item permits into async blocks and add guard test --- services/server/src/routes/analyze.rs | 10 +++--- services/server/src/routes/remember.rs | 18 +++-------- services/server/src/services/write_stream.rs | 33 ++++++++++++++++++++ 3 files changed, 43 insertions(+), 18 deletions(-) diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index 8683ea00..640131a1 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -512,6 +512,9 @@ pub async fn analyze( .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, @@ -525,15 +528,12 @@ pub async fn analyze( // carry `importance` through the prep tuple so // the job payload below can persist it alongside the // ciphertext + vector. - let result = Ok::<_, AppError>(( + Ok::<_, AppError>(( fact.text, fact.importance, vector_result?, encrypted_result?, - )); - // `_permit` is held until this fact's prep hands off to the - // durable wallet queue (or fails). - result + )) } }) .collect(); diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index c8887a87..df6ddd08 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -225,6 +225,9 @@ fn spawn_prepare_bulk_remember_job( 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 @@ -258,15 +261,12 @@ fn spawn_prepare_bulk_remember_job( ); let (vector_result, encrypted_result) = tokio::join!(embed_fut, encrypt_fut); - let result = Ok::<_, AppError>(( + Ok::<_, AppError>(( item.job_id, item.namespace, vector_result?, encrypted_result?, - )); - // `_permit` is held until this item's prep hands off - // to the durable bulk queue (or fails). - result + )) } }) .collect(); @@ -795,14 +795,6 @@ pub async fn remember_bulk( } let item_count = body.items.len(); - if item_count > state.config.write_stream_max_concurrency { - 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", - item_count - ))); - } let owner = &auth.owner; tracing::info!( diff --git a/services/server/src/services/write_stream.rs b/services/server/src/services/write_stream.rs index 6c40771a..be5bbe81 100644 --- a/services/server/src/services/write_stream.rs +++ b/services/server/src/services/write_stream.rs @@ -400,4 +400,37 @@ mod tests { 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); + } } From 5ba351a5b5f51431ba29e18884fb48292a7a2a7a Mon Sep 17 00:00:00 2001 From: DalenMax Date: Tue, 16 Jun 2026 09:45:32 +0700 Subject: [PATCH 29/29] docs(style): fix Sui style-guide violations in WALM-116 docs --- ...026-06-15-relayer-write-stream-redesign.md | 20 ++++++++--------- ...026-06-15-relayer-write-stream-redesign.md | 22 +++++++++---------- 2 files changed, 21 insertions(+), 21 deletions(-) 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 index a8ab584e..71d94e8d 100644 --- a/docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md +++ b/docs/superpowers/plans/2026-06-15-relayer-write-stream-redesign.md @@ -1,4 +1,4 @@ -# Relayer Write-Stream Redesign Implementation Plan +## 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. @@ -381,7 +381,7 @@ git commit -m "feat(write-stream): wire WriteStreamLimiter into AppState and Con --- -### Task 3: Add Prometheus Metrics for the Limiter +### Task 3: Add Prometheus metrics for the limiter **Files:** - Modify: `services/server/src/observability.rs` @@ -472,7 +472,7 @@ git commit -m "feat(write-stream): add Prometheus metrics for limiter state" --- -### Task 4: Add Shared Helper to Translate Permit Timeout to `AppError` +### Task 4: Add shared helper to translate permit timeout to `AppError` **Files:** - Modify: `services/server/src/routes/mod.rs` @@ -567,7 +567,7 @@ fn spawn_prepare_remember_job( ) { ``` -The permit will be dropped 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. +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** @@ -741,7 +741,7 @@ git commit -m "feat(write-stream): gate remember/manual and analyze with limiter --- -### Task 7: Gate Wallet Worker Upload Path +### Task 7: Gate wallet worker upload path **Files:** - Modify: `services/server/src/jobs.rs` @@ -795,7 +795,7 @@ git commit -m "feat(write-stream): gate wallet worker upload with limiter" --- -### Task 8: Add Snapshot Method and Wire Metric Emission +### Task 8: Add snapshot method and wire metric emission **Files:** - Modify: `services/server/src/services/write_stream.rs` @@ -961,7 +961,7 @@ git commit -m "feat(write-stream): emit limiter state and acquisition metrics" --- -### Task 9: Update Sidecar Safety-Net Configuration +### Task 9: Update sidecar safety-net configuration **Files:** - Modify: `services/server/scripts/sidecar/config.ts` @@ -1010,7 +1010,7 @@ git commit -m "feat(sidecar): raise default upload concurrency safety net" --- -### Task 10: Add Integration Test for Saturation Behavior +### Task 10: Add integration test for saturation behavior **Files:** - Modify: `services/server/src/routes/remember.rs` (add tests at the bottom) @@ -1066,7 +1066,7 @@ git commit -m "test(write-stream): add saturation unit test" --- -### Task 11: Full Test Suite and Lint +### Task 11: Full test suite and lint **Files:** - All modified files @@ -1107,7 +1107,7 @@ git commit -m "chore(write-stream): clippy fixes and test cleanup" --- -### Task 12: Documentation Updates +### Task 12: Documentation updates **Files:** - Modify: `docs/relayer/self-hosting.md` 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 index 401aaaa1..ed1a3018 100644 --- a/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md +++ b/docs/superpowers/specs/2026-06-15-relayer-write-stream-redesign.md @@ -1,4 +1,4 @@ -# Relayer Write-Stream Redesign +## Relayer write-stream redesign | | | |---|---| @@ -7,7 +7,7 @@ | **Owner** | Max Mai | | **Date** | 2026-06-15 | -## 1. Problem Statement +## 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. @@ -36,11 +36,11 @@ Move the effective upload-slot budget into the Rust relayer so that: - 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 +## 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 will result 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. +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 @@ -115,7 +115,7 @@ Initialize in `main.rs` after config load. 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 will mark it failed after `STALE_REMEMBER_JOB_AFTER`. +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). @@ -143,7 +143,7 @@ Initialize in `main.rs` after config load. `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 will be retried. Because permits are only held by active work, retries naturally back off until slots free up. +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 @@ -163,7 +163,7 @@ 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. Data flow ### 6.1 Single `/api/remember` @@ -199,7 +199,7 @@ Client POST /api/remember/bulk with N items ← return 202 + job_ids ``` -## 7. Error Handling +## 7. Error handling | Scenario | Behavior | |---|---| @@ -273,7 +273,7 @@ Logs: Run a targeted benchmark that previously produced `queuedWalrusUploads > 20`. Verify the metric stays near zero and p95 `/api/remember` latency is stable. -## 11. Risks & Mitigations +## 11. Risks and mitigations | Risk | Mitigation | |---|---| @@ -282,12 +282,12 @@ Run a targeted benchmark that previously produced `queuedWalrusUploads > 20`. Ve | 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 +## 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 +## 13. Open questions None. Design approved for implementation planning.