scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-09-21 00:21:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor(store): route test helpers through server-owned runs

Browse source
This commit is contained in:
Bryan Helmkamp 2026-04-05 02:53:17 -04:00
parent 9597a115b7
commit 50aff2787b
No known key found for this signature in database
6 changed files with 625 additions and 2300 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

179
lib/crates/fabro-cli/tests/it/cmd/support.rs
View file

@ -1,18 +1,96 @@
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::process::Output;
use std::sync::Arc;
use std::time::{Duration, Instant};
use fabro_store::{EventEnvelope, RunProjection, SlateRunStore, SlateStore};
use fabro_store::EventEnvelope;
use fabro_test::TestContext;
use fabro_types::RunId;
use object_store::local::LocalFileSystem;
use fabro_types::{
Checkpoint, Conclusion, NodeStatusRecord, PullRequestRecord, Retro, RunRecord,
RunStatusRecord, SandboxRecord, StageId, StartRecord,
};
use serde_json::Value;
use shlex::try_quote;
const COMMAND_TIMEOUT: Duration = Duration::from_secs(30);
#[allow(dead_code)]
#[derive(Debug, Clone, Default, serde::Deserialize)]
pub(crate) struct RunProjection {
#[serde(default)]
pub run: Option<RunRecord>,
#[serde(default)]
pub graph_source: Option<String>,
#[serde(default)]
pub start: Option<StartRecord>,
#[serde(default)]
pub status: Option<RunStatusRecord>,
#[serde(default)]
pub checkpoint: Option<Checkpoint>,
#[serde(default)]
pub checkpoints: Vec<(u32, Checkpoint)>,
#[serde(default)]
pub conclusion: Option<Conclusion>,
#[serde(default)]
pub retro: Option<Retro>,
#[serde(default)]
pub retro_prompt: Option<String>,
#[serde(default)]
pub retro_response: Option<String>,
#[serde(default)]
pub sandbox: Option<SandboxRecord>,
#[serde(default)]
pub final_patch: Option<String>,
#[serde(default)]
pub pull_request: Option<PullRequestRecord>,
#[serde(default)]
pub nodes: std::collections::HashMap<String, NodeState>,
}
#[allow(dead_code)]
#[derive(Debug, Clone, Default, serde::Deserialize)]
pub(crate) struct NodeState {
#[serde(default)]
pub prompt: Option<String>,
#[serde(default)]
pub response: Option<String>,
#[serde(default)]
pub status: Option<NodeStatusRecord>,
#[serde(default)]
pub provider_used: Option<serde_json::Value>,
#[serde(default)]
pub diff: Option<String>,
#[serde(default)]
pub script_invocation: Option<serde_json::Value>,
#[serde(default)]
pub script_timing: Option<serde_json::Value>,
#[serde(default)]
pub parallel_results: Option<serde_json::Value>,
#[serde(default)]
pub stdout: Option<String>,
#[serde(default)]
pub stderr: Option<String>,
}
#[derive(Debug, Clone, Default, serde::Deserialize)]
struct RunSummaryRecord {
run_id: String,
#[serde(default)]
labels: std::collections::HashMap<String, String>,
}
impl RunProjection {
pub(crate) fn iter_nodes(&self) -> impl Iterator<Item = (StageId, &NodeState)> {
self.nodes
.iter()
.filter_map(|(stage_id, state)| stage_id.parse::<StageId>().ok().map(|id| (id, state)))
}
pub(crate) fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
}
pub(crate) struct RunSetup {
pub(crate) run_id: String,
pub(crate) run_dir: PathBuf,
@ -209,12 +287,7 @@ pub(crate) fn setup_detached_dry_run(context: &TestContext) -> RunSetup {
.to_string();
let run = resolve_run(context, &run_id);
let deadline = Instant::now() + COMMAND_TIMEOUT;
while {
let store = run_store(&run.run_dir);
block_on(store.list_events())
.ok()
.is_none_or(|events| events.is_empty())
} {
while run_events(&run.run_dir).is_empty() {
assert!(
Instant::now() < deadline,
"timed out waiting for store events for {run_id}"
@ -332,13 +405,7 @@ worktree_mode = "never"
);
let run = run_local_workflow(context, &workspace_dir, "run.toml");
let store = run_store(&run.run_dir);
assert!(
block_on(store.state())
.ok()
.and_then(|state| state.sandbox)
.is_some()
);
assert!(run_state(&run.run_dir).sandbox.is_some());
WorkspaceRunSetup { run, workspace_dir }
}
@ -423,9 +490,9 @@ pub(crate) fn write_gated_workflow(path: &Path, name: &str, goal: &str) -> Workf
pub(crate) fn wait_for_status(run_dir: &Path, expected: &[&str]) -> String {
let deadline = Instant::now() + COMMAND_TIMEOUT;
loop {
if let Some(status) = block_on(run_store(run_dir).state())
.ok()
.and_then(|state| state.status.map(|record| record.status.to_string()))
if let Some(status) = run_state(run_dir)
.status
.map(|record| record.status.to_string())
{
if expected.iter().any(|candidate| *candidate == status) {
return status;
@ -461,23 +528,15 @@ pub(crate) fn run_count_for_test_case(context: &TestContext) -> usize {
}
fn run_dirs_for_test_case(context: &TestContext) -> Vec<PathBuf> {
let runs_dir = context.storage_dir.join("runs");
let entries = match std::fs::read_dir(&runs_dir) {
Ok(entries) => entries,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Vec::new(),
Err(err) => panic!("failed to read {}: {err}", runs_dir.display()),
};
entries
.filter_map(Result::ok)
.map(|entry| entry.path())
.filter(|path| path.is_dir())
.filter(|path| {
std::panic::catch_unwind(|| run_state(path)).ok().and_then(|state| state.run).is_some_and(|run| {
run.labels
.get("fabro_test_case")
.is_some_and(|value| value == context.test_case_id())
})
let runs: Vec<RunSummaryRecord> =
block_on(get_server_json_for_storage(&context.storage_dir, "/api/v1/runs"));
runs.into_iter()
.filter(|run| {
run.labels
.get("fabro_test_case")
.is_some_and(|value| value == context.test_case_id())
})
.filter_map(|run| find_run_dir(&context.storage_dir, &run.run_id))
.collect()
}
@ -550,28 +609,50 @@ fn block_on<T>(future: impl std::future::Future<Output = T>) -> T {
.block_on(future)
}
fn run_store(run_dir: &Path) -> SlateRunStore {
fn server_http_client(storage_dir: &Path) -> reqwest::Client {
reqwest::ClientBuilder::new()
.unix_socket(storage_dir.join("fabro.sock"))
.no_proxy()
.build()
.expect("test HTTP client should build")
}
async fn get_server_json<T: serde::de::DeserializeOwned>(run_dir: &Path, path: &str) -> T {
let runs_dir = run_dir.parent().expect("run dir should have parent");
let storage_dir = runs_dir.parent().expect("runs dir should have parent");
let run_id: RunId = infer_run_id(run_dir).parse().expect("run id should parse");
let object_store = Arc::new(
LocalFileSystem::new_with_prefix(storage_dir.join("store"))
.expect("test store path should be accessible"),
get_server_json_for_storage(storage_dir, path).await
}
async fn get_server_json_for_storage<T: serde::de::DeserializeOwned>(
storage_dir: &Path,
path: &str,
) -> T {
let response = server_http_client(storage_dir)
.get(format!("http://fabro{path}"))
.send()
.await
.expect("server request should succeed");
assert!(
response.status().is_success(),
"server request failed for {path}: {}",
response.status()
);
let store = Arc::new(SlateStore::new(object_store, "", Duration::from_millis(1)));
block_on(store.open_run_reader(&run_id)).expect("run store should exist")
response
.json::<T>()
.await
.expect("server response should parse")
}
pub(crate) fn run_state(run_dir: &Path) -> RunProjection {
let store = run_store(run_dir);
block_on(store.state()).expect("run store state should exist")
let run_id = infer_run_id(run_dir);
block_on(get_server_json(run_dir, &format!("/api/v1/runs/{run_id}/state")))
}
pub(crate) fn run_events(run_dir: &Path) -> Vec<EventEnvelope> {
let store = run_store(run_dir);
block_on(store.list_events())
.ok()
.expect("run store events should exist")
let run_id = infer_run_id(run_dir);
let response: serde_json::Value =
block_on(get_server_json(run_dir, &format!("/api/v1/runs/{run_id}/events")));
serde_json::from_value(response["data"].clone()).expect("event list should parse")
}
pub(crate) fn git_stdout(repo_dir: &Path, args: &[&str]) -> String {

10
lib/crates/fabro-server/src/server.rs
View file

@ -2,7 +2,7 @@ use std::collections::HashMap;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;
use std::time::{Duration, Instant};
#[cfg(test)]
use axum::body::to_bytes;
@ -127,6 +127,7 @@ struct ManagedRun {
status: RunStatus,
error: Option<String>,
created_at: chrono::DateTime<chrono::Utc>,
enqueued_at: Instant,
// Populated when running:
interviewer: Option<Arc<WebInterviewer>>,
event_tx: Option<broadcast::Sender<RunEvent>>,
@ -182,6 +183,7 @@ impl AppState {
}
}
/// Build the axum Router with all run endpoints and embedded static assets.
pub fn build_router(state: Arc<AppState>, auth_mode: AuthMode) -> Router {
let middleware_state = Arc::clone(&state);
@ -687,6 +689,7 @@ fn managed_run(
status,
error: None,
created_at,
enqueued_at: Instant::now(),
interviewer: None,
event_tx: None,
checkpoint: None,
@ -912,7 +915,7 @@ async fn start_run(
/// Execute a single run: transitions queued → starting → running → completed/failed/cancelled.
async fn execute_run(state: Arc<AppState>, run_id: RunId) {
// Transition to Starting and set up cancel infrastructure
let (cancel_rx, run_dir, event_tx, cancel_token, execution_mode) = {
let (cancel_rx, run_dir, event_tx, cancel_token, execution_mode, queued_for) = {
let mut runs = state.runs.lock().expect("runs lock poisoned");
let managed_run = match runs.get_mut(&run_id) {
Some(r) if r.status == RunStatus::Queued => r,
@ -937,8 +940,10 @@ async fn execute_run(state: Arc<AppState>, run_id: RunId) {
managed_run.event_tx.clone(),
cancel_token,
managed_run.execution_mode,
managed_run.enqueued_at.elapsed(),
)
};
let _ = queued_for;
// Create interviewer and event plumbing (this is the "provisioning" phase)
let interviewer = Arc::new(WebInterviewer::new());
@ -1351,7 +1356,6 @@ async fn list_run_events(
};
let since_seq = params.since_seq();
let limit = params.limit();
match state.store.open_run_reader(&id).await {
Ok(run_store) => match run_store.list_events_from_with_limit(since_seq, limit).await {
Ok(mut events) => {

109
lib/crates/fabro-store/src/keys.rs
View file

@ -1,45 +1,83 @@
use crate::StageId;
use fabro_types::RunBlobId;
use fabro_types::{RunBlobId, RunId};
pub(crate) const RUNS_PREFIX: &str = "runs/";
pub(crate) const CATALOG_BY_ID_PREFIX: &str = "_catalog/by-id/";
pub(crate) const CATALOG_BY_START_PREFIX: &str = "_catalog/by-start/";
pub(crate) const INIT_KEY: &str = "_init.json";
pub(crate) const EVENTS_PREFIX: &str = "events#";
pub(crate) const BLOBS_PREFIX: &str = "blobs#";
pub(crate) const ARTIFACT_NODES_PREFIX: &str = "artifacts#nodes#";
pub(crate) fn init() -> &'static str {
INIT_KEY
pub(crate) fn run_prefix(run_id: &RunId) -> String {
format!("{RUNS_PREFIX}{run_id}/")
}
pub(crate) fn event_key(seq: u32, epoch_ms: i64) -> String {
format!("{EVENTS_PREFIX}{seq:06}-{epoch_ms}.json")
pub(crate) fn init_key(run_id: &RunId) -> String {
format!("{}{INIT_KEY}", run_prefix(run_id))
}
pub(crate) fn blob_key(id: &RunBlobId) -> String {
format!("{BLOBS_PREFIX}{id}")
pub(crate) fn events_prefix(run_id: &RunId) -> String {
format!("{}{EVENTS_PREFIX}", run_prefix(run_id))
}
pub(crate) fn node_artifact_prefix(node: &StageId) -> String {
pub(crate) fn event_key(run_id: &RunId, seq: u32, epoch_ms: i64) -> String {
format!("{}{seq:06}-{epoch_ms}.json", events_prefix(run_id))
}
pub(crate) fn blobs_prefix(run_id: &RunId) -> String {
format!("{}{BLOBS_PREFIX}", run_prefix(run_id))
}
pub(crate) fn blob_key(run_id: &RunId, id: &RunBlobId) -> String {
format!("{}{id}", blobs_prefix(run_id))
}
pub(crate) fn node_artifact_prefix(run_id: &RunId, node: &StageId) -> String {
format!(
"{ARTIFACT_NODES_PREFIX}{}#visit-{}",
"{}{ARTIFACT_NODES_PREFIX}{}#visit-{}",
run_prefix(run_id),
node.node_id(),
node.visit()
)
}
pub(crate) fn node_artifact(node: &StageId, filename: &str) -> String {
format!("{}#{filename}", node_artifact_prefix(node))
pub(crate) fn node_artifact(run_id: &RunId, node: &StageId, filename: &str) -> String {
format!("{}#{filename}", node_artifact_prefix(run_id, node))
}
pub(crate) fn catalog_by_id_key(run_id: &RunId) -> String {
format!("{CATALOG_BY_ID_PREFIX}{run_id}.json")
}
pub(crate) fn catalog_by_start_prefix() -> &'static str {
CATALOG_BY_START_PREFIX
}
pub(crate) fn catalog_by_start_key(run_id: &RunId) -> String {
format!(
"{CATALOG_BY_START_PREFIX}{}/{run_id}.json",
run_id.created_at().format("%Y-%m-%d-%H-%M")
)
}
pub(crate) fn parse_event_seq(key: &str) -> Option<u32> {
parse_seq(key, EVENTS_PREFIX)
parse_seq(key.rsplit('/').next()?, EVENTS_PREFIX)
}
pub(crate) fn parse_blob_id(key: &str) -> Option<RunBlobId> {
key.strip_prefix(BLOBS_PREFIX)?.parse().ok()
key.rsplit('/').next()?.strip_prefix(BLOBS_PREFIX)?.parse().ok()
}
pub(crate) fn parse_node_artifact_key(key: &str) -> Option<(StageId, String)> {
parse_visit_scoped_key(key, ARTIFACT_NODES_PREFIX)
let artifact_start = key.find(ARTIFACT_NODES_PREFIX)?;
parse_visit_scoped_key(&key[artifact_start..], ARTIFACT_NODES_PREFIX)
}
pub(crate) fn parse_run_id_from_catalog_key(key: &str) -> Option<RunId> {
let filename = key.rsplit('/').next()?;
let run_id = filename.strip_suffix(".json").unwrap_or(filename);
run_id.parse().ok()
}
fn parse_seq(key: &str, prefix: &str) -> Option<u32> {
@ -59,33 +97,50 @@ mod tests {
#[test]
fn top_level_keys_match_spec() {
assert_eq!(init(), "_init.json");
assert_eq!(event_key(7, 123), "events#000007-123.json");
let run_id = "01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap();
assert_eq!(INIT_KEY, "_init.json");
assert_eq!(
event_key(&run_id, 7, 123),
"runs/01JT56VE4Z5NZ814GZN2JZD65A/events#000007-123.json"
);
}
#[test]
fn sequence_keys_are_zero_padded() {
assert_eq!(event_key(7, 123), "events#000007-123.json");
let run_id = "01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap();
assert_eq!(
event_key(&run_id, 7, 123),
"runs/01JT56VE4Z5NZ814GZN2JZD65A/events#000007-123.json"
);
}
#[test]
fn artifact_keys_match_spec() {
let node = StageId::new("code", 2);
let blob_id = RunBlobId::new(&"01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap(), b"summary");
assert_eq!(blob_key(&blob_id), format!("blobs#{blob_id}"));
let run_id = "01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap();
let blob_id = RunBlobId::new(&run_id, b"summary");
assert_eq!(blob_key(&run_id, &blob_id), format!("runs/{run_id}/blobs#{blob_id}"));
assert_eq!(
node_artifact(&node, "src/main.rs"),
"artifacts#nodes#code#visit-2#src/main.rs"
node_artifact(&run_id, &node, "src/main.rs"),
"runs/01JT56VE4Z5NZ814GZN2JZD65A/artifacts#nodes#code#visit-2#src/main.rs"
);
}
#[test]
fn parse_helpers_extract_sequences_and_node_visits() {
assert_eq!(parse_event_seq("events#000007-123.json"), Some(7));
let blob_id = RunBlobId::new(&"01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap(), b"summary");
assert_eq!(parse_blob_id(&format!("blobs#{blob_id}")), Some(blob_id));
assert_eq!(
parse_node_artifact_key("artifacts#nodes#code#visit-2#src/main.rs"),
parse_event_seq("runs/01JT56VE4Z5NZ814GZN2JZD65A/events#000007-123.json"),
Some(7)
);
let blob_id = RunBlobId::new(&"01JT56VE4Z5NZ814GZN2JZD65A".parse().unwrap(), b"summary");
assert_eq!(
parse_blob_id(&format!("runs/01JT56VE4Z5NZ814GZN2JZD65A/blobs#{blob_id}")),
Some(blob_id)
);
assert_eq!(
parse_node_artifact_key(
"runs/01JT56VE4Z5NZ814GZN2JZD65A/artifacts#nodes#code#visit-2#src/main.rs"
),
Some((StageId::new("code", 2), "src/main.rs".to_string()))
);
}
@ -103,7 +158,9 @@ mod tests {
#[test]
fn asset_filename_with_slashes_parses_correctly() {
assert_eq!(
parse_node_artifact_key("artifacts#nodes#build#visit-1#deep/nested/path/file.rs"),
parse_node_artifact_key(
"runs/01JT56VE4Z5NZ814GZN2JZD65A/artifacts#nodes#build#visit-1#deep/nested/path/file.rs"
),
Some((
StageId::new("build", 1),
"deep/nested/path/file.rs".to_string()

169
lib/crates/fabro-store/src/slate/catalog.rs
View file

@ -1,56 +1,35 @@
use std::collections::HashSet;
use std::sync::Arc;
use bytes::Bytes;
use futures::TryStreamExt;
use object_store::ObjectStore;
use object_store::path::Path;
use chrono::{Datelike, Timelike};
use slatedb::Db;
use crate::keys;
use crate::{ListRunsQuery, Result};
use fabro_types::RunId;
pub(crate) async fn write_catalog(
store: Arc<dyn ObjectStore>,
base_prefix: &str,
run_id: &RunId,
) -> Result<()> {
store
.put(&by_id_path(base_prefix, run_id), Bytes::new().into())
.await?;
store
.put(&by_start_path(base_prefix, run_id), Bytes::new().into())
.await?;
pub(crate) async fn write_catalog(db: &Db, run_id: &RunId) -> Result<()> {
db.put(keys::catalog_by_id_key(run_id), []).await?;
db.put(keys::catalog_by_start_key(run_id), []).await?;
Ok(())
}
pub(crate) async fn read_locator(
store: Arc<dyn ObjectStore>,
base_prefix: &str,
run_id: &RunId,
) -> Result<bool> {
match store.head(&by_id_path(base_prefix, run_id)).await {
Ok(_) => Ok(true),
Err(object_store::Error::NotFound { .. }) => Ok(false),
Err(err) => Err(err.into()),
}
pub(crate) async fn read_locator(db: &Db, run_id: &RunId) -> Result<bool> {
Ok(db.get(keys::catalog_by_id_key(run_id)).await?.is_some())
}
pub(crate) async fn list_run_ids(
store: Arc<dyn ObjectStore>,
base_prefix: &str,
query: &ListRunsQuery,
) -> Result<Vec<RunId>> {
let prefix = Path::from(format!("{base_prefix}by-start"));
let metas = store.list(Some(&prefix)).try_collect::<Vec<_>>().await?;
pub(crate) async fn delete_catalog(db: &Db, run_id: &RunId) -> Result<()> {
db.delete(keys::catalog_by_id_key(run_id)).await?;
db.delete(keys::catalog_by_start_key(run_id)).await?;
Ok(())
}
pub(crate) async fn list_run_ids(db: &Db, query: &ListRunsQuery) -> Result<Vec<RunId>> {
let mut iter = db.scan_prefix(keys::catalog_by_start_prefix()).await?;
let mut run_ids = Vec::new();
let mut seen = HashSet::new();
for meta in metas {
let Some(run_id) = parse_run_id_from_path(&meta.location) else {
while let Some(entry) = iter.next().await? {
let key = String::from_utf8(entry.key.to_vec())
.map_err(|err| crate::StoreError::Other(format!("stored key is not valid UTF-8: {err}")))?;
let Some(run_id) = keys::parse_run_id_from_catalog_key(&key) else {
continue;
};
if !seen.insert(run_id) {
continue;
}
let created_at = run_id.created_at();
if let Some(start) = query.start {
if created_at < start {
@ -64,102 +43,16 @@ pub(crate) async fn list_run_ids(
}
run_ids.push(run_id);
}
run_ids.sort_by_key(|run_id| {
let created_at = run_id.created_at();
(
created_at.year(),
created_at.month(),
created_at.day(),
created_at.hour(),
created_at.minute(),
*run_id,
)
});
Ok(run_ids)
}
pub(crate) fn parse_run_id_from_path(path: &Path) -> Option<RunId> {
let filename = path.filename()?;
let run_id = filename.strip_suffix(".json").unwrap_or(filename);
run_id.parse().ok()
}
pub(crate) fn db_prefix(base_prefix: &str, run_id: &RunId) -> String {
format!(
"{base_prefix}db/{}/{run_id}/",
run_id.created_at().format("%Y-%m-%d-%H-%M-%S-%3f")
)
}
pub(crate) fn by_id_path(base_prefix: &str, run_id: &RunId) -> Path {
Path::from(format!("{base_prefix}by-id/{run_id}.json"))
}
pub(crate) fn by_start_path(base_prefix: &str, run_id: &RunId) -> Path {
Path::from(format!(
"{base_prefix}by-start/{}/{run_id}.json",
run_id.created_at().format("%Y-%m-%d-%H-%M")
))
}
#[cfg(test)]
pub(super) mod test_support {
use super::*;
pub(crate) async fn repair_catalog(
store: Arc<dyn ObjectStore>,
base_prefix: &str,
) -> Result<()> {
let by_id_prefix = Path::from(format!("{base_prefix}by-id"));
let by_start_prefix = Path::from(format!("{base_prefix}by-start"));
let by_id_metas = store
.list(Some(&by_id_prefix))
.try_collect::<Vec<_>>()
.await?;
let run_ids = by_id_metas
.iter()
.filter_map(|meta| parse_run_id_from_path(&meta.location))
.collect::<Vec<_>>();
for run_id in &run_ids {
let path = by_start_path(base_prefix, run_id);
if !object_exists(store.clone(), &path).await? {
store.put(&path, Bytes::new().into()).await?;
}
}
let by_start_metas = store
.list(Some(&by_start_prefix))
.try_collect::<Vec<_>>()
.await?;
let canonical = run_ids.into_iter().collect::<HashSet<_>>();
let mut seen = HashSet::new();
for meta in by_start_metas {
let location = meta.location.clone();
let Some(run_id) = parse_run_id_from_path(&location) else {
delete_if_exists(store.clone(), &location).await?;
continue;
};
let expected = by_start_path(base_prefix, &run_id);
if canonical.contains(&run_id) && expected == location {
seen.insert(run_id);
continue;
}
delete_if_exists(store.clone(), &location).await?;
}
for run_id in canonical {
if !seen.contains(&run_id) {
store
.put(&by_start_path(base_prefix, &run_id), Bytes::new().into())
.await?;
}
}
Ok(())
}
async fn object_exists(store: Arc<dyn ObjectStore>, path: &Path) -> Result<bool> {
match store.head(path).await {
Ok(_) => Ok(true),
Err(object_store::Error::NotFound { .. }) => Ok(false),
Err(err) => Err(err.into()),
}
}
async fn delete_if_exists(store: Arc<dyn ObjectStore>, path: &Path) -> Result<()> {
match store.delete(path).await {
Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
Err(err) => Err(err.into()),
}
}
}

2108
lib/crates/fabro-store/src/slate/mod.rs
View file

File diff suppressed because it is too large Load diff

350
lib/crates/fabro-store/src/slate/run_store.rs
View file

@ -1,15 +1,13 @@
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Weak};
use std::time::Duration;
use std::sync::Arc;
use bytes::Bytes;
use chrono::Utc;
use futures::Stream;
use serde::Serialize;
use serde::de::DeserializeOwned;
use slatedb::{CloseReason, DbRead, DbReader, ErrorKind};
use tokio::sync::{Mutex, mpsc};
use tokio::time;
use slatedb::{CloseReason, Db, DbRead, ErrorKind};
use tokio::sync::{Mutex, broadcast, mpsc};
use tokio_stream::wrappers::UnboundedReceiverStream;
use crate::keys;
@ -17,6 +15,8 @@ use crate::run_state::EventProjectionCache;
use crate::{EventEnvelope, EventPayload, Result, RunProjection, RunSummary, StageId, StoreError};
use fabro_types::{RunBlobId, RunId};
const DEFAULT_EVENT_TAIL_LIMIT: usize = 1024;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodeArtifact {
pub node: StageId,
@ -26,108 +26,113 @@ pub struct NodeArtifact {
#[derive(Clone)]
pub struct SlateRunStore {
inner: Arc<SlateRunStoreInner>,
read_only: bool,
}
impl std::fmt::Debug for SlateRunStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SlateRunStore")
.field("run_id", &self.inner.run_id)
.field("db_prefix", &self.inner.db_prefix)
.field("read_only", &self.read_only)
.finish_non_exhaustive()
}
}
pub(crate) struct SlateRunStoreInner {
run_id: RunId,
db_prefix: String,
db: SlateRunDb,
db: Db,
event_seq: AtomicU32,
close_lock: Mutex<()>,
projection_cache: Mutex<EventProjectionCache>,
}
enum SlateRunDb {
Writer(slatedb::Db),
Reader(Box<DbReader>),
recent_events: Mutex<VecDeque<EventEnvelope>>,
recent_event_limit: usize,
event_tx: broadcast::Sender<EventEnvelope>,
}
impl SlateRunStore {
pub(crate) async fn open_writer(
run_id: RunId,
db_prefix: String,
db: slatedb::Db,
) -> Result<Self> {
let event_seq = recover_next_seq(&db, keys::EVENTS_PREFIX, keys::parse_event_seq).await?;
pub(crate) async fn open_writer(run_id: RunId, db: Db) -> Result<Self> {
let event_seq = recover_next_seq(&db, &keys::events_prefix(&run_id), keys::parse_event_seq).await?;
let (event_tx, _) = broadcast::channel(DEFAULT_EVENT_TAIL_LIMIT.max(16));
Ok(Self {
inner: Arc::new(SlateRunStoreInner {
run_id,
db_prefix,
db: SlateRunDb::Writer(db),
db,
event_seq: AtomicU32::new(event_seq),
close_lock: Mutex::new(()),
projection_cache: Mutex::new(EventProjectionCache::default()),
recent_events: Mutex::new(VecDeque::with_capacity(DEFAULT_EVENT_TAIL_LIMIT)),
recent_event_limit: DEFAULT_EVENT_TAIL_LIMIT,
event_tx,
}),
read_only: false,
})
}
pub(crate) async fn open_reader(
run_id: RunId,
db_prefix: String,
db: DbReader,
) -> Result<Self> {
let event_seq = recover_next_seq(&db, keys::EVENTS_PREFIX, keys::parse_event_seq).await?;
pub(crate) async fn open_reader(run_id: RunId, db: Db) -> Result<Self> {
let event_seq = recover_next_seq(&db, &keys::events_prefix(&run_id), keys::parse_event_seq).await?;
let (event_tx, _) = broadcast::channel(DEFAULT_EVENT_TAIL_LIMIT.max(16));
Ok(Self {
inner: Arc::new(SlateRunStoreInner {
run_id,
db_prefix,
db: SlateRunDb::Reader(Box::new(db)),
db,
event_seq: AtomicU32::new(event_seq),
close_lock: Mutex::new(()),
projection_cache: Mutex::new(EventProjectionCache::default()),
recent_events: Mutex::new(VecDeque::with_capacity(DEFAULT_EVENT_TAIL_LIMIT)),
recent_event_limit: DEFAULT_EVENT_TAIL_LIMIT,
event_tx,
}),
read_only: true,
})
}
pub(crate) fn from_inner(inner: Arc<SlateRunStoreInner>) -> Self {
Self { inner }
Self {
inner,
read_only: false,
}
}
pub(crate) fn downgrade(&self) -> Weak<SlateRunStoreInner> {
Arc::downgrade(&self.inner)
pub(crate) fn into_read_only(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
read_only: true,
}
}
pub(crate) fn inner_arc(&self) -> Arc<SlateRunStoreInner> {
Arc::clone(&self.inner)
}
pub(crate) fn run_id(&self) -> RunId {
self.inner.run_id
}
pub(crate) fn matches_run(&self, run_id: &RunId, db_prefix: &str) -> bool {
self.inner.run_id == *run_id && self.inner.db_prefix == db_prefix
pub(crate) fn matches_run(&self, run_id: &RunId) -> bool {
self.inner.run_id == *run_id
}
pub(crate) async fn close(&self) -> Result<()> {
let _guard = self.inner.close_lock.lock().await;
match self.inner.db.close().await {
Ok(()) => Ok(()),
Err(err) if matches!(err.kind(), ErrorKind::Closed(CloseReason::Clean)) => Ok(()),
Err(err) => Err(err.into()),
if Arc::strong_count(&self.inner) <= 1 {
match self.inner.db.close().await {
Ok(()) => Ok(()),
Err(err) if matches!(err.kind(), ErrorKind::Closed(CloseReason::Clean)) => Ok(()),
Err(err) => Err(err.into()),
}
} else {
Ok(())
}
}
pub(crate) async fn snapshot(&self) -> Result<Arc<slatedb::DbSnapshot>> {
match &self.inner.db {
SlateRunDb::Writer(db) => Ok(db.snapshot().await?),
SlateRunDb::Reader(_) => Err(StoreError::ReadOnly),
}
}
pub(crate) async fn validate_init<R>(db: &R, expected: &RunId) -> Result<bool>
pub(crate) async fn validate_init<R>(db: &R, run_id: &RunId) -> Result<bool>
where
R: DbRead + Sync,
{
match get_json::<R, RunId>(db, keys::init()).await? {
Some(existing) if existing == *expected => Ok(true),
match get_json::<R, RunId>(db, &keys::init_key(run_id)).await? {
Some(existing) if existing == *run_id => Ok(true),
Some(existing) => Err(StoreError::Other(format!(
"existing _init.json {existing:?} does not match requested run_id {expected:?}"
"existing init record {existing:?} does not match requested run_id {run_id:?}"
))),
None => Ok(false),
}
@ -137,7 +142,7 @@ impl SlateRunStore {
where
R: DbRead + Sync,
{
let events = list_events_from(db, 1).await?;
let events = list_events_from(db, run_id, 1).await?;
let state = RunProjection::apply_events(&events)?;
Ok(state.build_summary(run_id))
}
@ -147,7 +152,7 @@ impl SlateRunStore {
let cache = self.inner.projection_cache.lock().await;
cache.last_seq.saturating_add(1)
};
let events = self.inner.db.list_events_from(next_seq).await?;
let events = list_events_from(&self.inner.db, &self.inner.run_id, next_seq).await?;
let mut cache = self.inner.projection_cache.lock().await;
for event in &events {
cache.state.apply_event(event)?;
@ -155,24 +160,62 @@ impl SlateRunStore {
}
Ok(cache.state.clone())
}
async fn cache_event(&self, event: &EventEnvelope) -> Result<()> {
{
let mut projection_cache = self.inner.projection_cache.lock().await;
projection_cache.state.apply_event(event)?;
projection_cache.last_seq = event.seq;
}
let mut recent_events = self.inner.recent_events.lock().await;
recent_events.push_back(event.clone());
while recent_events.len() > self.inner.recent_event_limit {
recent_events.pop_front();
}
let _ = self.inner.event_tx.send(event.clone());
Ok(())
}
async fn cached_events_from(&self, start_seq: u32, limit: usize) -> Option<Vec<EventEnvelope>> {
let recent_events = self.inner.recent_events.lock().await;
let oldest_seq = recent_events.front().map(|event| event.seq)?;
if start_seq < oldest_seq {
return None;
}
let mut events = recent_events
.iter()
.filter(|event| event.seq >= start_seq)
.take(limit.saturating_add(1))
.cloned()
.collect::<Vec<_>>();
if events.is_empty() && start_seq <= self.inner.event_seq.load(Ordering::SeqCst) {
events = Vec::new();
}
Some(events)
}
}
impl SlateRunStore {
pub async fn append_event(&self, payload: &EventPayload) -> Result<u32> {
if self.read_only {
return Err(StoreError::ReadOnly);
}
payload.validate(&self.inner.run_id)?;
let seq = self.inner.event_seq.fetch_add(1, Ordering::SeqCst);
self.inner
.db
.put_json(
&keys::event_key(seq, Utc::now().timestamp_millis()),
payload,
)
.await?;
let event = EventEnvelope {
seq,
payload: payload.clone(),
};
self.inner.db.put(
keys::event_key(&self.inner.run_id, seq, Utc::now().timestamp_millis()),
serde_json::to_vec(payload)?,
).await?;
self.cache_event(&event).await?;
Ok(seq)
}
pub async fn list_events(&self) -> Result<Vec<EventEnvelope>> {
self.inner.db.list_events_from(1).await
self.list_events_from_with_limit(1, usize::MAX / 2).await
}
pub async fn list_events_from_with_limit(
@ -180,10 +223,10 @@ impl SlateRunStore {
start_seq: u32,
limit: usize,
) -> Result<Vec<EventEnvelope>> {
self.inner
.db
.list_events_from_with_limit(start_seq, limit)
.await
if let Some(events) = self.cached_events_from(start_seq, limit).await {
return Ok(events);
}
list_events_from_with_limit(&self.inner.db, &self.inner.run_id, start_seq, limit).await
}
pub fn watch_events_from(
@ -192,72 +235,86 @@ impl SlateRunStore {
) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<EventEnvelope>> + Send>>> {
let inner = Arc::clone(&self.inner);
let (sender, receiver) = mpsc::unbounded_channel();
tokio::spawn(async move {
let cached = {
let recent_events = inner.recent_events.lock().await;
recent_events
.iter()
.filter(|event| event.seq >= seq)
.cloned()
.collect::<Vec<_>>()
};
let mut next_seq = seq;
loop {
if sender.is_closed() {
for event in cached {
next_seq = event.seq.saturating_add(1);
if sender.send(Ok(event)).is_err() {
return;
}
}
match inner.db.list_events_from(next_seq).await {
Ok(events) => {
if events.is_empty() {
time::sleep(Duration::from_millis(100)).await;
continue;
}
for event in events {
next_seq = event.seq.saturating_add(1);
if sender.send(Ok(event)).is_err() {
return;
}
}
}
Err(err) => {
let _ = sender.send(Err(err));
return;
}
let mut rx = inner.event_tx.subscribe();
while let Ok(event) = rx.recv().await {
if event.seq < next_seq {
continue;
}
next_seq = event.seq.saturating_add(1);
if sender.send(Ok(event)).is_err() {
return;
}
}
});
Ok(Box::pin(UnboundedReceiverStream::new(receiver)))
}
pub async fn write_blob(&self, data: &[u8]) -> Result<RunBlobId> {
if self.read_only {
return Err(StoreError::ReadOnly);
}
let id = RunBlobId::new(&self.inner.run_id, data);
self.inner.db.put_bytes(&keys::blob_key(&id), data).await?;
self.inner
.db
.put(keys::blob_key(&self.inner.run_id, &id), data)
.await?;
Ok(id)
}
pub async fn read_blob(&self, id: &RunBlobId) -> Result<Option<Bytes>> {
self.inner.db.get_bytes(&keys::blob_key(id)).await
Ok(self
.inner
.db
.get(keys::blob_key(&self.inner.run_id, id))
.await?)
}
pub async fn list_blobs(&self) -> Result<Vec<RunBlobId>> {
self.inner.db.list_blobs().await
list_blobs(&self.inner.db, &self.inner.run_id).await
}
pub async fn put_artifact(&self, node: &StageId, filename: &str, data: &[u8]) -> Result<()> {
if self.read_only {
return Err(StoreError::ReadOnly);
}
self.inner
.db
.put_bytes(&keys::node_artifact(node, filename), data)
.await
.put(keys::node_artifact(&self.inner.run_id, node, filename), data)
.await?;
Ok(())
}
pub async fn get_artifact(&self, node: &StageId, filename: &str) -> Result<Option<Bytes>> {
self.inner
Ok(self
.inner
.db
.get_bytes(&keys::node_artifact(node, filename))
.await
.get(keys::node_artifact(&self.inner.run_id, node, filename))
.await?)
}
pub async fn list_all_artifacts(&self) -> Result<Vec<NodeArtifact>> {
self.inner.db.list_all_artifacts().await
list_all_artifacts(&self.inner.db, &self.inner.run_id).await
}
pub async fn list_artifacts_for_stage(&self, stage_id: &StageId) -> Result<Vec<String>> {
self.inner.db.list_artifacts_for_stage(stage_id).await
list_artifacts_for_stage(&self.inner.db, &self.inner.run_id, stage_id).await
}
pub async fn state(&self) -> Result<RunProjection> {
@ -265,81 +322,6 @@ impl SlateRunStore {
}
}
impl SlateRunDb {
fn writer(&self) -> Result<&slatedb::Db> {
match self {
Self::Writer(db) => Ok(db),
Self::Reader(_) => Err(StoreError::ReadOnly),
}
}
async fn close(&self) -> std::result::Result<(), slatedb::Error> {
match self {
Self::Writer(db) => db.close().await,
Self::Reader(db) => db.close().await,
}
}
async fn put_json<T: Serialize>(&self, key: &str, value: &T) -> Result<()> {
put_json(self.writer()?, key, value).await
}
async fn get_bytes(&self, key: &str) -> Result<Option<Bytes>> {
match self {
Self::Writer(db) => get_bytes(db, key).await,
Self::Reader(db) => db.get(key).await.map_err(Into::into),
}
}
async fn put_bytes(&self, key: &str, value: &[u8]) -> Result<()> {
put_bytes(self.writer()?, key, value).await
}
async fn list_events_from(&self, start_seq: u32) -> Result<Vec<EventEnvelope>> {
match self {
Self::Writer(db) => list_events_from(db, start_seq).await,
Self::Reader(db) => list_events_from(db.as_ref(), start_seq).await,
}
}
async fn list_events_from_with_limit(
&self,
start_seq: u32,
limit: usize,
) -> Result<Vec<EventEnvelope>> {
match self {
Self::Writer(db) => list_events_from_with_limit(db, start_seq, limit).await,
Self::Reader(db) => list_events_from_with_limit(db.as_ref(), start_seq, limit).await,
}
}
async fn list_blobs(&self) -> Result<Vec<RunBlobId>> {
match self {
Self::Writer(db) => list_blobs(db).await,
Self::Reader(db) => list_blobs(db.as_ref()).await,
}
}
async fn list_all_artifacts(&self) -> Result<Vec<NodeArtifact>> {
match self {
Self::Writer(db) => list_all_artifacts(db).await,
Self::Reader(db) => list_all_artifacts(db.as_ref()).await,
}
}
async fn list_artifacts_for_stage(&self, stage_id: &StageId) -> Result<Vec<String>> {
match self {
Self::Writer(db) => list_artifacts_for_stage(db, stage_id).await,
Self::Reader(db) => list_artifacts_for_stage(db.as_ref(), stage_id).await,
}
}
}
async fn put_json<T: Serialize>(db: &slatedb::Db, key: &str, value: &T) -> Result<()> {
db.put(key, serde_json::to_vec(value)?).await?;
Ok(())
}
async fn get_json<R, T>(db: &R, key: &str) -> Result<Option<T>>
where
R: DbRead + Sync,
@ -352,15 +334,6 @@ where
.map_err(Into::into)
}
async fn put_bytes(db: &slatedb::Db, key: &str, value: &[u8]) -> Result<()> {
db.put(key, value).await?;
Ok(())
}
async fn get_bytes(db: &slatedb::Db, key: &str) -> Result<Option<Bytes>> {
Ok(db.get(key).await?)
}
async fn recover_next_seq<R>(db: &R, prefix: &str, parse: fn(&str) -> Option<u32>) -> Result<u32>
where
R: DbRead + Sync,
@ -376,11 +349,11 @@ where
Ok(max_seq.saturating_add(1).max(1))
}
async fn list_events_from<R>(db: &R, start_seq: u32) -> Result<Vec<EventEnvelope>>
async fn list_events_from<R>(db: &R, run_id: &RunId, start_seq: u32) -> Result<Vec<EventEnvelope>>
where
R: DbRead + Sync,
{
let mut iter = db.scan_prefix(keys::EVENTS_PREFIX.as_bytes()).await?;
let mut iter = db.scan_prefix(keys::events_prefix(run_id).as_bytes()).await?;
let mut events = Vec::new();
while let Some(entry) = iter.next().await? {
let key = key_to_string(&entry.key)?;
@ -401,22 +374,23 @@ where
async fn list_events_from_with_limit<R>(
db: &R,
run_id: &RunId,
start_seq: u32,
limit: usize,
) -> Result<Vec<EventEnvelope>>
where
R: DbRead + Sync,
{
let mut events = list_events_from(db, start_seq).await?;
let mut events = list_events_from(db, run_id, start_seq).await?;
events.truncate(limit.saturating_add(1));
Ok(events)
}
async fn list_blobs<R>(db: &R) -> Result<Vec<RunBlobId>>
async fn list_blobs<R>(db: &R, run_id: &RunId) -> Result<Vec<RunBlobId>>
where
R: DbRead + Sync,
{
let mut iter = db.scan_prefix(keys::BLOBS_PREFIX.as_bytes()).await?;
let mut iter = db.scan_prefix(keys::blobs_prefix(run_id).as_bytes()).await?;
let mut blob_ids = Vec::new();
while let Some(entry) = iter.next().await? {
let key = key_to_string(&entry.key)?;
@ -429,12 +403,12 @@ where
Ok(blob_ids)
}
async fn list_all_artifacts<R>(db: &R) -> Result<Vec<NodeArtifact>>
async fn list_all_artifacts<R>(db: &R, run_id: &RunId) -> Result<Vec<NodeArtifact>>
where
R: DbRead + Sync,
{
let mut iter = db
.scan_prefix(keys::ARTIFACT_NODES_PREFIX.as_bytes())
.scan_prefix(keys::run_prefix(run_id).as_bytes())
.await?;
let mut assets = Vec::new();
while let Some(entry) = iter.next().await? {
@ -448,11 +422,11 @@ where
Ok(assets)
}
async fn list_artifacts_for_stage<R>(db: &R, stage_id: &StageId) -> Result<Vec<String>>
async fn list_artifacts_for_stage<R>(db: &R, run_id: &RunId, stage_id: &StageId) -> Result<Vec<String>>
where
R: DbRead + Sync,
{
let prefix = keys::node_artifact_prefix(stage_id);
let prefix = keys::node_artifact_prefix(run_id, stage_id);
let mut iter = db.scan_prefix(prefix.as_bytes()).await?;
let mut filenames = Vec::new();
while let Some(entry) = iter.next().await? {