diff --git a/lib/components/fabro-store/src/error.rs b/lib/components/fabro-store/src/error.rs index 33126d474..43c26b44a 100644 --- a/lib/components/fabro-store/src/error.rs +++ b/lib/components/fabro-store/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { SessionAlreadyExists(String), #[error("run store is read-only")] ReadOnly, + #[error("event sequence limit of {max_seq} reached")] + EventSequenceExhausted { max_seq: u32 }, #[error("invalid key segment: {segment:?}")] InvalidKeySegment { segment: String }, #[error("failed to parse key: {0}")] diff --git a/lib/components/fabro-store/src/keys.rs b/lib/components/fabro-store/src/keys.rs index 2451df632..343cdc1d0 100644 --- a/lib/components/fabro-store/src/keys.rs +++ b/lib/components/fabro-store/src/keys.rs @@ -1,7 +1,10 @@ use std::fmt::{self, Write}; +use std::ops::Range; use fabro_types::{RunBlobId, RunId, SessionId}; +pub(crate) const MAX_EVENT_SEQ: u32 = 999_999; + #[derive(Debug, PartialEq, Eq)] pub(crate) struct SlateKey(String); @@ -23,6 +26,13 @@ impl SlateKey { self } + /// Exclusive end bound of this key's prefix keyspace: every key under + /// `self.into_prefix()` sorts below it and no other key sorts between. + fn into_prefix_end(mut self) -> Self { + self.0.push('\u{1}'); + self + } + #[cfg(test)] fn as_str(&self) -> &str { &self.0 @@ -52,6 +62,10 @@ pub(crate) fn run_events_prefix(run_id: &RunId) -> SlateKey { .into_prefix() } +// Sequence keys zero-pad `seq` to six digits so lexicographic key order +// matches numeric seq order through `MAX_EVENT_SEQ`. Seek-based event listing +// (`run_events_range`) depends on this invariant, so event allocation rejects +// larger sequences. pub(crate) fn run_event_key(run_id: &RunId, seq: u32, epoch_ms: i64) -> SlateKey { SlateKey::new("runs") .with(run_id) @@ -66,6 +80,17 @@ pub(crate) fn run_event_seq_prefix(run_id: &RunId, seq: u32) -> SlateKey { .with(format!("{seq:06}-")) } +/// Scan range covering the run's event keys from `start_seq` to the end of +/// the run's event namespace, so seek-based listing never touches keys of +/// other runs or namespaces. +pub(crate) fn run_events_range(run_id: &RunId, start_seq: u32) -> Range { + let end = SlateKey::new("runs") + .with(run_id) + .with("events") + .into_prefix_end(); + run_event_seq_prefix(run_id, start_seq)..end +} + pub(crate) fn blobs_prefix() -> SlateKey { SlateKey::new("blobs").with("sha256").into_prefix() } @@ -152,6 +177,24 @@ mod tests { assert_eq!(leaf, "000007-123"); } + #[test] + fn run_events_range_bounds_the_event_namespace() { + let run_id: RunId = "01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap(); + let range = run_events_range(&run_id, 2); + let contains = |key: &SlateKey| { + range.start.as_ref() <= key.as_ref() && key.as_ref() < range.end.as_ref() + }; + + assert!(!contains(&run_event_key(&run_id, 1, 123))); + assert!(contains(&run_event_key(&run_id, 2, 123))); + assert!(contains(&run_event_key(&run_id, MAX_EVENT_SEQ, 123))); + // Sibling namespaces of the same run sort outside the range. + assert!(!contains(&SlateKey::new("runs").with(run_id).with("state"))); + assert!(!contains( + &session_by_id_key(&fabro_types::SessionId::new()) + )); + } + #[test] fn parse_helpers_roundtrip() { let run_id: RunId = "01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap(); diff --git a/lib/components/fabro-store/src/slate/run_store.rs b/lib/components/fabro-store/src/slate/run_store.rs index a1d7f8441..7efa68aba 100644 --- a/lib/components/fabro-store/src/slate/run_store.rs +++ b/lib/components/fabro-store/src/slate/run_store.rs @@ -38,7 +38,9 @@ pub(crate) struct RunDatabaseInner { run_id: RunId, db: Db, blob_store: BlobStore, - event_seq: AtomicU32, + // `None` for reader-built inners: readers never append, so they carry no + // next-write sequence and any append through them fails as read-only. + event_seq: Option, close_lock: Mutex<()>, state_lock: Mutex<()>, projection_cache: Mutex, @@ -84,8 +86,13 @@ impl RunDatabase { shared_projection_cache: Arc, run_summary_store: Arc>>, ) -> Result { - let event_seq = - recover_next_seq(&db, keys::run_events_prefix(&run_id), keys::parse_event_seq).await?; + let event_seq = if read_only { + // Readers never append, so they do not need to scan the full event + // history to recover the next write sequence. + None + } else { + Some(AtomicU32::new(recover_next_seq(&db, &run_id).await?)) + }; let (event_tx, _) = broadcast::channel(DEFAULT_EVENT_TAIL_LIMIT.max(16)); let blob_store = BlobStore::new(Arc::new(db.clone())); Ok(Self { @@ -93,7 +100,7 @@ impl RunDatabase { run_id, db, blob_store, - event_seq: AtomicU32::new(event_seq), + event_seq, close_lock: Mutex::new(()), state_lock: Mutex::new(()), projection_cache: Mutex::new(EventProjectionCache::default()), @@ -240,15 +247,12 @@ impl RunDatabase { if start_seq < oldest_seq { return None; } - let mut events = recent_events + let events = recent_events .iter() .filter(|event| event.seq >= start_seq) .take(limit.saturating_add(1)) .cloned() .collect::>(); - if events.is_empty() && start_seq <= self.inner.event_seq.load(Ordering::SeqCst) { - events = Vec::new(); - } Some(events) } } @@ -287,7 +291,8 @@ impl RunDatabase { } async fn append_event_envelope_locked(&self, payload: &EventPayload) -> Result { - let seq = self.inner.event_seq.fetch_add(1, Ordering::SeqCst); + let event_seq = self.inner.event_seq.as_ref().ok_or(Error::ReadOnly)?; + let seq = allocate_event_seq(event_seq)?; let event = EventEnvelope { seq, event: RunEvent::try_from(payload)?, @@ -360,9 +365,11 @@ impl RunDatabase { } pub async fn list_events(&self) -> Result> { - self.list_events_from_with_limit(1, usize::MAX / 2).await + self.list_events_from_with_limit(1, usize::MAX).await } + /// Returns up to `limit + 1` events starting at `start_seq`. The extra + /// item lets callers compute `has_more` without a second read. pub async fn list_events_from_with_limit( &self, start_seq: u32, @@ -500,6 +507,16 @@ impl RunDatabase { } } +fn allocate_event_seq(event_seq: &AtomicU32) -> Result { + event_seq + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |seq| { + (seq <= keys::MAX_EVENT_SEQ).then_some(seq + 1) + }) + .map_err(|_| Error::EventSequenceExhausted { + max_seq: keys::MAX_EVENT_SEQ, + }) +} + fn apply_cached_projection_event( state: &mut Option, event: &EventEnvelope, @@ -512,19 +529,15 @@ fn apply_cached_projection_event( Ok(()) } -async fn recover_next_seq( - db: &R, - prefix: keys::SlateKey, - parse: fn(&str) -> Option, -) -> Result +async fn recover_next_seq(db: &R, run_id: &RunId) -> Result where R: DbRead + Sync, { - let mut iter = db.scan_prefix(prefix).await?; + let mut iter = db.scan_prefix(keys::run_events_prefix(run_id)).await?; let mut max_seq = 0; while let Some(entry) = iter.next().await? { - let key = key_to_string(&entry.key)?; - if let Some(seq) = parse(&key) { + let key = key_to_str(&entry.key)?; + if let Some(seq) = keys::parse_event_seq(key) { max_seq = max_seq.max(seq); } } @@ -535,11 +548,31 @@ async fn list_events_from(db: &R, run_id: &RunId, start_seq: u32) -> Result( + db: &R, + run_id: &RunId, + start_seq: u32, + limit: usize, +) -> Result> +where + R: DbRead + Sync, +{ + let max_events = limit.saturating_add(1); + // Seek to the page cursor and decode only the requested page. Zero-padded + // sequence keys scan in seq order, so no post-scan sort is needed. + let mut iter = db.scan(keys::run_events_range(run_id, start_seq)).await?; let mut events = Vec::new(); - while let Some(entry) = iter.next().await? { - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { + while events.len() < max_events { + let Some(entry) = iter.next().await? else { + break; + }; + let key = key_to_str(&entry.key)?; + let Some(seq) = keys::parse_event_seq(key) else { continue; }; if seq < start_seq { @@ -550,21 +583,6 @@ where event: serde_json::from_slice(&entry.value)?, }); } - events.sort_by_key(|event| event.seq); - Ok(events) -} - -async fn list_events_from_with_limit( - db: &R, - run_id: &RunId, - start_seq: u32, - limit: usize, -) -> Result> -where - R: DbRead + Sync, -{ - let mut events = list_events_from(db, run_id, start_seq).await?; - events.truncate(limit.saturating_add(1)); Ok(events) } @@ -615,8 +633,8 @@ where let mut iter = db.scan_prefix(keys::run_events_prefix(run_id)).await?; let mut events: Vec = Vec::new(); while let Some(entry) = iter.next().await? { - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { + let key = key_to_str(&entry.key)?; + let Some(seq) = keys::parse_event_seq(key) else { continue; }; if seq < start_seq { @@ -675,8 +693,8 @@ where let mut iter = db.scan_prefix(keys::run_events_prefix(run_id)).await?; let mut events = Vec::new(); while let Some(entry) = iter.next().await? { - let key = key_to_string(&entry.key)?; - let Some(seq) = keys::parse_event_seq(&key) else { + let key = key_to_str(&entry.key)?; + let Some(seq) = keys::parse_event_seq(key) else { continue; }; if seq < start_seq { @@ -710,8 +728,8 @@ where let mut iter = db.scan_prefix(keys::blobs_prefix()).await?; let mut blob_ids = Vec::new(); while let Some(entry) = iter.next().await? { - let key = key_to_string(&entry.key)?; - let Some(blob_id) = keys::parse_blob_id(&key) else { + let key = key_to_str(&entry.key)?; + let Some(blob_id) = keys::parse_blob_id(key) else { continue; }; blob_ids.push(blob_id); @@ -720,21 +738,22 @@ where Ok(blob_ids) } -fn key_to_string(key: &Bytes) -> Result { - String::from_utf8(key.to_vec()) +fn key_to_str(key: &Bytes) -> Result<&str> { + std::str::from_utf8(key) .map_err(|err| Error::Other(format!("stored key is not valid UTF-8: {err}"))) } #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::Ordering; use std::time::Duration; use fabro_types::{Graph, RunId, SessionId, StageId, WorkflowSettings, test_support}; use object_store::memory::InMemory; use serde_json::json; - use crate::{Database, EventPayload}; + use crate::{Database, Error, EventPayload, keys}; #[tokio::test] async fn list_blobs_reads_global_cas_namespace() { @@ -836,6 +855,89 @@ mod tests { run } + #[tokio::test] + async fn list_events_from_with_limit_does_not_read_past_limit_plus_one() { + let run = fresh_run().await; + let run_id = run.run_id(); + run.append_event(&stage_prompt_payload(&run_id, 1, Some("alpha"))) + .await + .unwrap(); + run.append_event(&stage_prompt_payload(&run_id, 2, Some("beta"))) + .await + .unwrap(); + run.inner + .db + .put(keys::run_event_key(&run_id, 4, 0), b"invalid json") + .await + .unwrap(); + + let events = super::list_events_from_with_limit(&run.inner.db, &run_id, 1, 2) + .await + .unwrap(); + + let seqs: Vec = events.iter().map(|event| event.seq).collect(); + assert_eq!(seqs, vec![1, 2, 3]); + } + + #[tokio::test] + async fn list_events_from_with_limit_seeks_to_start_sequence() { + let run = fresh_run().await; + let run_id = run.run_id(); + run.append_event(&stage_prompt_payload(&run_id, 1, Some("alpha"))) + .await + .unwrap(); + run.append_event(&stage_prompt_payload(&run_id, 2, Some("beta"))) + .await + .unwrap(); + let mut unreadable_earlier_key = keys::run_event_seq_prefix(&run_id, 2).as_ref().to_vec(); + unreadable_earlier_key.push(0xff); + run.inner + .db + .put(unreadable_earlier_key, b"invalid json") + .await + .unwrap(); + + let events = super::list_events_from_with_limit(&run.inner.db, &run_id, 3, 1) + .await + .unwrap(); + + let seqs: Vec = events.iter().map(|event| event.seq).collect(); + assert_eq!(seqs, vec![3]); + } + + #[tokio::test] + async fn append_event_rejects_sequences_beyond_key_order_limit() { + let run = fresh_run().await; + let run_id = run.run_id(); + run.inner + .event_seq + .as_ref() + .unwrap() + .store(keys::MAX_EVENT_SEQ, Ordering::SeqCst); + + let seq = run + .append_event(&stage_prompt_payload(&run_id, 1, Some("alpha"))) + .await + .unwrap(); + assert_eq!(seq, keys::MAX_EVENT_SEQ); + + let err = run + .append_event(&stage_prompt_payload(&run_id, 2, Some("beta"))) + .await + .unwrap_err(); + assert!(matches!( + err, + Error::EventSequenceExhausted { max_seq } + if max_seq == keys::MAX_EVENT_SEQ + )); + assert!( + run.get_event(keys::MAX_EVENT_SEQ + 1) + .await + .unwrap() + .is_none() + ); + } + #[tokio::test] async fn list_events_for_stage_returns_only_matching_events_in_seq_order() { let run = fresh_run().await;