diff --git a/lib/apps/fabro-server/src/run_files.rs b/lib/apps/fabro-server/src/run_files.rs index 1f07d09af..126408e43 100644 --- a/lib/apps/fabro-server/src/run_files.rs +++ b/lib/apps/fabro-server/src/run_files.rs @@ -1194,14 +1194,8 @@ fn to_sha_wrapper(sha: &str) -> RunFilesMetaToSha { async fn load_projection( state: &Arc, run_id: &RunId, -) -> std::result::Result { - let cached = state - .store_ref() - .get_cached_run(run_id) - .await - .map_err(|err| ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))? - .ok_or_else(|| ApiError::not_found("Run not found."))?; - Ok((*cached.projection).clone()) +) -> std::result::Result, ApiError> { + Ok(state.cached_run(run_id).await?.projection) } async fn reconnect_run_sandbox( diff --git a/lib/apps/fabro-server/src/server.rs b/lib/apps/fabro-server/src/server.rs index ed67c1898..407169fbc 100644 --- a/lib/apps/fabro-server/src/server.rs +++ b/lib/apps/fabro-server/src/server.rs @@ -85,8 +85,8 @@ use fabro_slack::threads::ThreadRegistry; use fabro_slack::{blocks as slack_blocks, connection as slack_connection}; use fabro_static::EnvVars; use fabro_store::{ - ArtifactKey, ArtifactStore, Database, EventEnvelope, EventPayload, NodeArtifact, - PendingInterviewRecord, RunSummaryStore, StageArtifactEntry, StageId, + ArtifactKey, ArtifactStore, CachedRunProjection, Database, EventEnvelope, EventPayload, + NodeArtifact, PendingInterviewRecord, RunSummaryStore, StageArtifactEntry, StageId, }; #[cfg(test)] use fabro_types::BlockedReason; @@ -1520,6 +1520,18 @@ impl AppState { &self.stores.runs } + /// Current cached projection for `run_id`, with the standard HTTP error + /// mapping: storage failures become 500s and a missing run becomes the + /// canonical 404. + pub(crate) async fn cached_run(&self, run_id: &RunId) -> Result { + self.stores + .runs + .get_cached_run(run_id) + .await + .map_err(|err| ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))? + .ok_or_else(|| ApiError::not_found("Run not found.")) + } + pub(crate) fn session_runtimes(&self) -> &SessionRuntimeManager { &self.session_runtimes } @@ -2686,7 +2698,7 @@ async fn delete_run_internal( } async fn load_durable_run_status(state: &AppState, id: &RunId) -> Option { - let cached = state.stores.runs.get_cached_run(id).await.ok()??; + let cached = state.cached_run(id).await.ok()?; Some(cached.projection.status) } @@ -3736,15 +3748,10 @@ async fn load_pending_interview( run_id: RunId, qid: &str, ) -> Result { - let cached = match state.stores.runs.get_cached_run(&run_id).await { - Ok(Some(cached)) => cached, - Ok(None) => return Err(ApiError::not_found("Run not found.").into_response()), - Err(err) => { - return Err( - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response(), - ); - } - }; + let cached = state + .cached_run(&run_id) + .await + .map_err(IntoResponse::into_response)?; let Some(record) = cached.projection.pending_interviews.get(qid) else { return Err(ApiError::new( StatusCode::CONFLICT, @@ -4520,7 +4527,7 @@ async fn append_control_request( /// run is currently archived. Returns `None` otherwise (including when the run /// doesn't exist — the caller's own not-found handling will surface that). async fn reject_if_archived(state: &AppState, run_id: &RunId) -> Option { - let cached = state.stores.runs.get_cached_run(run_id).await.ok()??; + let cached = state.cached_run(run_id).await.ok()?; cached.projection.archived_at.is_some().then(|| { ApiError::new( StatusCode::CONFLICT, diff --git a/lib/apps/fabro-server/src/server/handler/artifacts.rs b/lib/apps/fabro-server/src/server/handler/artifacts.rs index 0d7403e0e..9d7b049fb 100644 --- a/lib/apps/fabro-server/src/server/handler/artifacts.rs +++ b/lib/apps/fabro-server/src/server/handler/artifacts.rs @@ -68,15 +68,12 @@ async fn get_checkpoint( Ok(id) => id, Err(response) => return response, }; - match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => match cached.projection.current_checkpoint() { + match state.cached_run(&id).await { + Ok(cached) => match cached.projection.current_checkpoint() { Some(cp) => (StatusCode::OK, Json(cp.clone())).into_response(), None => (StatusCode::OK, Json(serde_json::json!(null))).into_response(), }, - Ok(None) => ApiError::not_found("Run not found.").into_response(), - Err(err) => { - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() - } + Err(err) => err.into_response(), } } @@ -118,17 +115,12 @@ async fn read_run_blob( } } -async fn load_run_spec(state: &AppState, run_id: &RunId) -> Result { - let cached = state - .stores - .runs - .get_cached_run(run_id) +async fn ensure_run_exists(state: &AppState, run_id: &RunId) -> Result<(), Response> { + state + .cached_run(run_id) .await - .map_err(|err| { - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() - })? - .ok_or_else(|| ApiError::not_found("Run not found.").into_response())?; - Ok(cached.projection.spec.clone()) + .map(|_| ()) + .map_err(IntoResponse::into_response) } async fn list_run_artifacts( @@ -140,7 +132,7 @@ async fn list_run_artifacts( Ok(id) => id, Err(response) => return response, }; - if let Err(response) = load_run_spec(state.as_ref(), &id).await { + if let Err(response) = ensure_run_exists(state.as_ref(), &id).await { return response; } @@ -186,7 +178,7 @@ async fn list_stage_artifacts( Ok(stage_id) => stage_id, Err(response) => return response, }; - if let Err(response) = load_run_spec(state.as_ref(), &id).await { + if let Err(response) = ensure_run_exists(state.as_ref(), &id).await { return response; } @@ -568,7 +560,7 @@ async fn put_stage_artifact( if let Some(response) = reject_if_archived(state.as_ref(), &id).await { return response; } - if let Err(response) = load_run_spec(state.as_ref(), &id).await.map(|_| ()) { + if let Err(response) = ensure_run_exists(state.as_ref(), &id).await { return response; } let retry = match required_query_param(params.retry.as_ref(), "retry") { @@ -636,7 +628,7 @@ async fn get_stage_artifact( Ok(path) => path, Err(response) => return response, }; - if let Err(response) = load_run_spec(state.as_ref(), &id).await { + if let Err(response) = ensure_run_exists(state.as_ref(), &id).await { return response; } diff --git a/lib/apps/fabro-server/src/server/handler/billing.rs b/lib/apps/fabro-server/src/server/handler/billing.rs index 557e1180a..ba7b71625 100644 --- a/lib/apps/fabro-server/src/server/handler/billing.rs +++ b/lib/apps/fabro-server/src/server/handler/billing.rs @@ -5,9 +5,9 @@ use chrono::{DateTime, Utc}; use fabro_types::{RunProjection, StageHandler, StageProjection, StageState, StageTiming}; use super::super::{ - ApiError, AppState, BillingByModel, BillingStageRef, IntoResponse, Json, ListResponse, - PaginationParams, Path, Query, RequiredUser, Response, Router, RunBilling, RunBillingStage, - RunBillingTotals, RunId, State, StatusCode, get, parse_run_id_path, run_stage_from_stage_id, + AppState, BillingByModel, BillingStageRef, IntoResponse, Json, ListResponse, PaginationParams, + Path, Query, RequiredUser, Response, Router, RunBilling, RunBillingStage, RunBillingTotals, + RunId, State, StatusCode, get, parse_run_id_path, run_stage_from_stage_id, }; pub(super) fn routes() -> Router> { @@ -27,13 +27,9 @@ async fn list_run_stages( Err(response) => return response, }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let projection = cached.projection; @@ -70,13 +66,9 @@ async fn get_run_billing( State(state): State>, Path(id): Path, ) -> Response { - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let projection = cached.projection; diff --git a/lib/apps/fabro-server/src/server/handler/graph.rs b/lib/apps/fabro-server/src/server/handler/graph.rs index e2982abed..364b19cad 100644 --- a/lib/apps/fabro-server/src/server/handler/graph.rs +++ b/lib/apps/fabro-server/src/server/handler/graph.rs @@ -243,12 +243,9 @@ async fn load_run_dot_source(state: &AppState, id: &RunId) -> Result, id: &RunId, ) -> Result { - let cached = state - .stores - .runs - .get_cached_run(id) - .await - .map_err(|err| ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))? - .ok_or_else(|| ApiError::not_found("Run not found."))?; + let cached = state.cached_run(id).await?; cached.projection.pull_request.clone().ok_or_else(|| { ApiError::with_code( StatusCode::NOT_FOUND, @@ -294,19 +288,9 @@ async fn create_run_pull_request( let Ok(run_store) = state.stores.runs.open_run(&id).await else { return ApiError::not_found("Run not found.").into_response(); }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => { - return ApiError::new( - StatusCode::INTERNAL_SERVER_ERROR, - "Run projection unavailable.", - ) - .into_response(); - } - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let run_state = cached.projection.as_ref(); let inputs = match RunPrInputs::extract(run_state, body.force) { @@ -406,19 +390,9 @@ async fn unlink_run_pull_request( let Ok(run_store) = state.stores.runs.open_run(&id).await else { return ApiError::not_found("Run not found.").into_response(); }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => { - return ApiError::new( - StatusCode::INTERNAL_SERVER_ERROR, - "Run projection unavailable.", - ) - .into_response(); - } - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let Some(pull_request) = cached.projection.pull_request.clone() else { return ApiError::with_code( diff --git a/lib/apps/fabro-server/src/server/handler/runs.rs b/lib/apps/fabro-server/src/server/handler/runs.rs index 66e1f10be..82e3660cd 100644 --- a/lib/apps/fabro-server/src/server/handler/runs.rs +++ b/lib/apps/fabro-server/src/server/handler/runs.rs @@ -942,13 +942,9 @@ async fn get_run_settings( Ok(id) => id, Err(response) => return response, }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; ( StatusCode::OK, @@ -961,8 +957,8 @@ async fn get_questions( RequireRunManagementTarget(id, _actor): RequireRunManagementTarget, State(state): State>, ) -> Response { - match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => { + match state.cached_run(&id).await { + Ok(cached) => { let questions = cached .projection .pending_interviews @@ -971,10 +967,7 @@ async fn get_questions( .collect::>(); (StatusCode::OK, Json(ListResponse::new(questions))).into_response() } - Ok(None) => ApiError::not_found("Run not found.").into_response(), - Err(err) => { - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() - } + Err(err) => err.into_response(), } } @@ -1006,12 +999,9 @@ async fn get_run_state( RequireRunManagementTarget(id, _actor): RequireRunManagementTarget, State(state): State>, ) -> Response { - match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => Json((*cached.projection).clone()).into_response(), - Ok(None) => ApiError::not_found("Run not found.").into_response(), - Err(err) => { - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() - } + match state.cached_run(&id).await { + Ok(cached) => Json(&*cached.projection).into_response(), + Err(err) => err.into_response(), } } @@ -1046,13 +1036,9 @@ async fn get_run_stage_context_window( Ok(stage_id) => stage_id, Err(response) => return response, }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let Some(stage) = cached.projection.stage(&stage_id) else { return ApiError::not_found("Stage not found.").into_response(); @@ -1106,16 +1092,9 @@ async fn get_run_stage_command_log( return ApiError::bad_request("limit must be greater than 0").into_response(); } let limit = query.limit.min(MAX_COMMAND_LOG_LIMIT); - let Ok(run_store) = state.stores.runs.open_run_reader(&id).await else { - return ApiError::not_found("Run not found.").into_response(); - }; - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; let Some(node) = cached.projection.stage(&stage_id) else { return ApiError::not_found("Stage not found.").into_response(); @@ -1153,7 +1132,14 @@ async fn get_run_stage_command_log( } if let Some(cas_ref) = cas_ref { - let text = match read_json_string_blob(&run_store.clone().into(), &cas_ref).await { + let run_store = match state.stores.runs.open_run_reader(&id).await { + Ok(run_store) => run_store, + Err(err) => { + return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) + .into_response(); + } + }; + let text = match read_json_string_blob(&run_store.into(), &cas_ref).await { Ok(Some(text)) => text, Ok(None) => String::new(), Err(err) => { diff --git a/lib/apps/fabro-server/src/server/handler/sandbox.rs b/lib/apps/fabro-server/src/server/handler/sandbox.rs index 96fab6b16..c2027571e 100644 --- a/lib/apps/fabro-server/src/server/handler/sandbox.rs +++ b/lib/apps/fabro-server/src/server/handler/sandbox.rs @@ -952,14 +952,9 @@ async fn load_run_sandbox_instance( run_id: &RunId, ) -> Result { let cached = state - .stores - .runs - .get_cached_run(run_id) + .cached_run(run_id) .await - .map_err(|err| { - ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response() - })? - .ok_or_else(|| ApiError::not_found("Run not found.").into_response())?; + .map_err(IntoResponse::into_response)?; cached .projection .sandbox diff --git a/lib/apps/fabro-server/src/server/handler/worker_control.rs b/lib/apps/fabro-server/src/server/handler/worker_control.rs index 982f05053..7f70e9297 100644 --- a/lib/apps/fabro-server/src/server/handler/worker_control.rs +++ b/lib/apps/fabro-server/src/server/handler/worker_control.rs @@ -35,13 +35,9 @@ async fn worker_control_stream( Query(query): Query, ws: WebSocketUpgrade, ) -> Response { - let cached = match state.stores.runs.get_cached_run(&id).await { - Ok(Some(cached)) => cached, - Ok(None) => return ApiError::not_found("Run not found.").into_response(), - Err(err) => { - return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) - .into_response(); - } + let cached = match state.cached_run(&id).await { + Ok(cached) => cached, + Err(err) => return err.into_response(), }; if cached.projection.archived_at.is_some() { return ApiError::new(StatusCode::CONFLICT, "Run is archived.").into_response(); diff --git a/lib/components/fabro-store/src/run_state.rs b/lib/components/fabro-store/src/run_state.rs index b37903299..58f11bbf0 100644 --- a/lib/components/fabro-store/src/run_state.rs +++ b/lib/components/fabro-store/src/run_state.rs @@ -1,5 +1,6 @@ use std::collections::{BTreeMap, HashMap}; use std::str::FromStr; +use std::sync::Arc; use chrono::{DateTime, Utc}; use fabro_types::run_event::{ @@ -26,7 +27,9 @@ use crate::{Error, EventEnvelope, Result}; #[derive(Debug, Clone, Default)] pub(crate) struct EventProjectionCache { pub last_seq: u32, - pub state: Option, + // Arc-shared with the shared projection cache so opening a run does not + // deep-copy the projection; mutated copy-on-write via `Arc::make_mut`. + pub state: Option>, } pub trait RunProjectionReducer { diff --git a/lib/components/fabro-store/src/slate/projection_cache.rs b/lib/components/fabro-store/src/slate/projection_cache.rs index c5bf4b2e3..cd9efd482 100644 --- a/lib/components/fabro-store/src/slate/projection_cache.rs +++ b/lib/components/fabro-store/src/slate/projection_cache.rs @@ -171,13 +171,18 @@ impl RunProjectionCache { .map(|entry| state.with_children_count(entry)) } - pub(crate) async fn last_seq(&self, run_id: &RunId) -> Option { + /// Projection and last sequence for `run_id`, without the summary clone + /// and children count that `get` computes under the cache mutex. + pub(crate) async fn projection_snapshot( + &self, + run_id: &RunId, + ) -> Option<(Arc, u32)> { self.state .lock() .await .entries .get(run_id) - .map(|entry| entry.last_seq) + .map(|entry| (Arc::clone(&entry.projection), entry.last_seq)) } pub(crate) async fn get_summary(&self, run_id: &RunId, now: DateTime) -> Option { diff --git a/lib/components/fabro-store/src/slate/run_store.rs b/lib/components/fabro-store/src/slate/run_store.rs index 2b982f0df..73d7bd2c6 100644 --- a/lib/components/fabro-store/src/slate/run_store.rs +++ b/lib/components/fabro-store/src/slate/run_store.rs @@ -6,7 +6,7 @@ use bytes::Bytes; use chrono::Utc; use fabro_types::{RunBlobId, RunEvent, RunId, SessionId}; use futures::Stream; -use slatedb::{Db, DbRead}; +use slatedb::{Db, DbIterator, DbRead}; use tokio::sync::{Mutex, broadcast, mpsc}; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{error, warn}; @@ -84,18 +84,16 @@ impl RunDatabase { shared_projection_cache: Arc, run_summary_store: Arc>>, ) -> Result { - let cached_projection = shared_projection_cache.get(&run_id).await; - let projection_cache = - cached_projection - .as_ref() - .map_or_else(EventProjectionCache::default, |cached| { - EventProjectionCache { - last_seq: cached.last_seq, - state: Some((*cached.projection).clone()), - } - }); + let cached_projection = shared_projection_cache.projection_snapshot(&run_id).await; + let projection_cache = cached_projection.as_ref().map_or_else( + EventProjectionCache::default, + |(projection, last_seq)| EventProjectionCache { + last_seq: *last_seq, + state: Some(Arc::clone(projection)), + }, + ); let event_seq = match (&cached_projection, read_only) { - (Some(cached), _) => cached.last_seq.saturating_add(1), + (Some((_, last_seq)), _) => last_seq.saturating_add(1), (None, true) => { // Readers never append, so they do not need to scan the full event // history to recover the next write sequence. @@ -187,12 +185,12 @@ impl RunDatabase { ))) } - async fn projected_state(&self) -> Result { + async fn projected_state(&self) -> Result> { let _state_guard = self.inner.state_lock.lock().await; self.projected_state_locked().await } - async fn projected_state_locked(&self) -> Result { + async fn projected_state_locked(&self) -> Result> { let next_seq = { let cache = self.inner.projection_cache.lock().await; cache.last_seq.saturating_add(1) @@ -249,7 +247,7 @@ impl RunDatabase { let state = RunProjection::apply_events(&events)?; let mut projection_cache = self.inner.projection_cache.lock().await; - projection_cache.state = Some(state); + projection_cache.state = Some(Arc::new(state)); projection_cache.last_seq = last_seq; Ok(()) } @@ -384,20 +382,12 @@ impl RunDatabase { } /// Returns the newest stored event sequence without reading event bodies - /// when a current local or shared projection is available. + /// when a current projection is available. pub async fn last_event_seq(&self) -> Result> { let local_last_seq = self.inner.projection_cache.lock().await.last_seq; if local_last_seq > 0 { return Ok(Some(local_last_seq)); } - if let Some(last_seq) = self - .inner - .shared_projection_cache - .last_seq(&self.inner.run_id) - .await - { - return Ok(Some(last_seq)); - } let next_seq = recover_next_seq( &self.inner.db, @@ -426,11 +416,10 @@ impl RunDatabase { /// Returns up to `limit + 1` events for the given stage visit, /// starting at `start_seq`. The `+1` lets callers compute `has_more`. /// - /// Implementation note: scans the unbounded run-event prefix and - /// filters by stage identity *before* applying `limit`, so a stage with - /// matches sparsely scattered late in the event log still returns its - /// full slice (no premature truncation from a generic `limit`-bounded - /// scan). + /// Implementation note: filters by stage identity *before* applying + /// `limit`, so a stage with matches sparsely scattered late in the event + /// log still returns its full slice (no premature truncation from a + /// generic `limit`-bounded scan). pub async fn list_events_for_stage_from_with_limit( &self, stage_id: &StageId, @@ -541,18 +530,20 @@ impl RunDatabase { } pub async fn state(&self) -> Result { - self.projected_state().await + Ok(Arc::unwrap_or_clone(self.projected_state().await?)) } } fn apply_cached_projection_event( - state: &mut Option, + state: &mut Option>, event: &EventEnvelope, ) -> Result<()> { if let Some(projection) = state { - projection.apply_event(event)?; + Arc::make_mut(projection).apply_event(event)?; } else { - *state = Some(RunProjection::apply_events(std::slice::from_ref(event))?); + *state = Some(Arc::new(RunProjection::apply_events( + std::slice::from_ref(event), + )?)); } Ok(()) } @@ -576,31 +567,55 @@ where Ok(max_seq.saturating_add(1).max(1)) } +/// Cursor over a run's stored events starting at `start_seq`, yielding raw +/// `(seq, payload)` entries in ascending sequence order (event keys embed a +/// zero-padded sequence, so key order matches sequence order). +struct EventScan { + iter: DbIterator, + event_prefix: keys::SlateKey, + start_seq: u32, +} + +impl EventScan { + async fn seek(db: &R, run_id: &RunId, start_seq: u32) -> Result + where + R: DbRead + Sync, + { + let iter = db + .scan(keys::run_event_seq_prefix(run_id, start_seq)..) + .await?; + Ok(Self { + iter, + event_prefix: keys::run_events_prefix(run_id), + start_seq, + }) + } + + async fn next(&mut self) -> Result> { + while let Some(entry) = self.iter.next().await? { + if !entry.key.starts_with(self.event_prefix.as_ref()) { + return Ok(None); + } + let key = key_to_string(&entry.key)?; + let Some(seq) = keys::parse_event_seq(&key) else { + continue; + }; + if seq < self.start_seq { + continue; + } + return Ok(Some((seq, entry.value))); + } + Ok(None) + } +} + async fn list_events_from(db: &R, run_id: &RunId, start_seq: u32) -> Result> where R: DbRead + Sync, { - let event_prefix = keys::run_events_prefix(run_id); - let mut iter = db - .scan(keys::run_event_seq_prefix(run_id, start_seq)..) - .await?; - let mut events = Vec::new(); - while let Some(entry) = iter.next().await? { - if !entry.key.starts_with(event_prefix.as_ref()) { - break; - } - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { - continue; - }; - if seq < start_seq { - continue; - } - events.push(EventEnvelope { - seq, - event: serde_json::from_slice(&entry.value)?, - }); - } + let mut events = list_events_from_with_limit(db, run_id, start_seq, usize::MAX / 2).await?; + // Key order matches sequence order only through the 6-digit zero padding + // in event keys; this keeps full-history replays correct past it. events.sort_by_key(|event| event.seq); Ok(events) } @@ -614,31 +629,18 @@ async fn list_events_from_with_limit( where R: DbRead + Sync, { - let event_prefix = keys::run_events_prefix(run_id); let max_events = limit.saturating_add(1); - // Seek to the page cursor and decode only the requested page plus the - // sentinel used to compute `has_more`. - let mut iter = db - .scan(keys::run_event_seq_prefix(run_id, start_seq)..) - .await?; + // Decode only the requested page plus the sentinel used to compute + // `has_more`. + let mut scan = EventScan::seek(db, run_id, start_seq).await?; let mut events = Vec::new(); while events.len() < max_events { - let Some(entry) = iter.next().await? else { + let Some((seq, value)) = scan.next().await? else { break; }; - if !entry.key.starts_with(event_prefix.as_ref()) { - break; - } - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { - continue; - }; - if seq < start_seq { - continue; - } events.push(EventEnvelope { seq, - event: serde_json::from_slice(&entry.value)?, + event: serde_json::from_slice(&value)?, }); } Ok(events) @@ -670,10 +672,9 @@ async fn list_events_for_stage_from_with_limit( where R: DbRead + Sync, { - // Scan without a storage-level item limit from the requested cursor: - // filtering by stage identity with a generic limit-bounded scan would - // silently drop matches whenever the stage's events are sparse late in - // the event log. + // Filter by stage identity *before* applying `limit`: a generic + // limit-bounded scan would silently drop matches whenever the stage's + // events are sparse late in the event log. // // We probe just the stage identity fields with a small partial deserialize and // only run the full `RunEvent` parse on matches. Most events in a run @@ -689,23 +690,13 @@ where let stage_id_string = stage_id.to_string(); let max_events = limit.saturating_add(1); - let event_prefix = keys::run_events_prefix(run_id); - let mut iter = db - .scan(keys::run_event_seq_prefix(run_id, start_seq)..) - .await?; - let mut events: Vec = Vec::new(); - while let Some(entry) = iter.next().await? { - if !entry.key.starts_with(event_prefix.as_ref()) { + let mut scan = EventScan::seek(db, run_id, start_seq).await?; + let mut events = Vec::new(); + while events.len() < max_events { + let Some((seq, value)) = scan.next().await? else { break; - } - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { - continue; }; - if seq < start_seq { - continue; - } - let probe: StageIdProbe = serde_json::from_slice(&entry.value)?; + let probe: StageIdProbe = serde_json::from_slice(&value)?; let matches_stage_id = probe.stage_id == Some(stage_id_string.as_str()); let matches_legacy_node_id = probe.stage_id.is_none() && stage_id.visit() == 1 @@ -713,25 +704,9 @@ where if !matches_stage_id && !matches_legacy_node_id { continue; } - let event: RunEvent = serde_json::from_slice(&entry.value)?; - let envelope = EventEnvelope { seq, event }; - if events.len() < max_events { - events.push(envelope); - continue; - } - - if let Some((max_index, max_seq)) = events - .iter() - .enumerate() - .max_by_key(|(_, existing)| existing.seq) - .map(|(index, existing)| (index, existing.seq)) - { - if seq < max_seq { - events[max_index] = envelope; - } - } + let event: RunEvent = serde_json::from_slice(&value)?; + events.push(EventEnvelope { seq, event }); } - events.sort_by_key(|event| event.seq); Ok(events) } @@ -755,24 +730,13 @@ where let session_id_string = session_id.to_string(); let max_events = limit.saturating_add(1); - let event_prefix = keys::run_events_prefix(run_id); - let mut iter = db - .scan(keys::run_event_seq_prefix(run_id, start_seq)..) - .await?; + let mut scan = EventScan::seek(db, run_id, start_seq).await?; let mut events = Vec::new(); - while let Some(entry) = iter.next().await? { - if !entry.key.starts_with(event_prefix.as_ref()) { + while events.len() < max_events { + let Some((seq, value)) = scan.next().await? else { break; - } - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { - continue; }; - if seq < start_seq { - continue; - } - - let probe: SessionEventProbe = serde_json::from_slice(&entry.value)?; + let probe: SessionEventProbe = serde_json::from_slice(&value)?; if probe.session_id != Some(session_id_string.as_str()) || !probe .event_name @@ -781,12 +745,9 @@ where continue; } - let event: RunEvent = serde_json::from_slice(&entry.value)?; + let event: RunEvent = serde_json::from_slice(&value)?; if event.body.is_run_session_event() { events.push(EventEnvelope { seq, event }); - if events.len() >= max_events { - break; - } } } Ok(events)