mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-06 08:18:58 +00:00
Fix Slate store fencing and restore status shape
This commit is contained in:
Bryan Helmkamp
parent
7b5a1ecb0b
commit
46737b7728
5 changed files with 412 additions and 96 deletions
|
|
@ -130,7 +130,7 @@ pub async fn repair_catalog(store: Arc<dyn ObjectStore>, base_prefix: &str) -> R
|
|||
pub(crate) fn db_prefix(base_prefix: &str, created_at: DateTime<Utc>, run_id: &str) -> String {
|
||||
format!(
|
||||
"{base_prefix}db/{}/{run_id}/",
|
||||
created_at.format("%Y-%m-%d-%H-%M")
|
||||
created_at.format("%Y-%m-%d-%H-%M-%S-%3f")
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
mod catalog;
|
||||
mod run_store;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
|
@ -8,15 +9,18 @@ use chrono::{DateTime, Utc};
|
|||
use futures::TryStreamExt;
|
||||
use object_store::path::Path;
|
||||
use object_store::ObjectStore;
|
||||
use slatedb::DbReader;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::keys;
|
||||
use crate::{CatalogRecord, ListRunsQuery, Result, RunStore, RunSummary, Store, StoreError};
|
||||
use run_store::SlateRunStore;
|
||||
use run_store::{SlateRunStore, SlateRunStoreInner};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SlateStore {
|
||||
object_store: Arc<dyn ObjectStore>,
|
||||
base_prefix: String,
|
||||
active_runs: Arc<Mutex<HashMap<String, std::sync::Weak<SlateRunStoreInner>>>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for SlateStore {
|
||||
|
|
@ -32,6 +36,7 @@ impl SlateStore {
|
|||
Self {
|
||||
object_store,
|
||||
base_prefix: normalize_base_prefix(base_prefix.into()),
|
||||
active_runs: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -43,15 +48,74 @@ impl SlateStore {
|
|||
Ok(slatedb::Db::open(db_prefix.to_string(), self.object_store.clone()).await?)
|
||||
}
|
||||
|
||||
async fn open_reader(&self, db_prefix: &str) -> Result<DbReader> {
|
||||
Ok(DbReader::open(
|
||||
db_prefix.to_string(),
|
||||
self.object_store.clone(),
|
||||
None,
|
||||
slatedb::config::DbReaderOptions::default(),
|
||||
)
|
||||
.await?)
|
||||
}
|
||||
|
||||
async fn db_prefix_has_objects(&self, db_prefix: &str) -> Result<bool> {
|
||||
let prefix = Path::from(db_prefix.to_string());
|
||||
let mut items = self.object_store.list(Some(&prefix));
|
||||
Ok(items.try_next().await?.is_some())
|
||||
}
|
||||
|
||||
async fn get_active_run(&self, run_id: &str) -> Option<SlateRunStore> {
|
||||
let mut active_runs = self.active_runs.lock().await;
|
||||
let weak = active_runs.get(run_id).cloned()?;
|
||||
match weak.upgrade() {
|
||||
Some(inner) => Some(SlateRunStore::from_inner(inner)),
|
||||
None => {
|
||||
active_runs.remove(run_id);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn cache_active_run(&self, run_store: &SlateRunStore) {
|
||||
self.active_runs
|
||||
.lock()
|
||||
.await
|
||||
.insert(run_store.record().run_id.clone(), run_store.downgrade());
|
||||
}
|
||||
|
||||
async fn remove_active_run(&self, run_id: &str) -> Option<SlateRunStore> {
|
||||
let weak = self.active_runs.lock().await.remove(run_id)?;
|
||||
weak.upgrade().map(SlateRunStore::from_inner)
|
||||
}
|
||||
|
||||
async fn open_run_store(&self, record: &CatalogRecord) -> Result<Option<SlateRunStore>> {
|
||||
if let Some(active) = self.get_active_run(&record.run_id).await {
|
||||
if active.matches_record(record) {
|
||||
return Ok(Some(active));
|
||||
}
|
||||
return Err(StoreError::Other(format!(
|
||||
"active run cache mismatch for run_id {:?}",
|
||||
record.run_id
|
||||
)));
|
||||
}
|
||||
if !self.db_prefix_has_objects(&record.db_prefix).await? {
|
||||
return Ok(None);
|
||||
}
|
||||
let db = self.open_db(&record.db_prefix).await?;
|
||||
if !SlateRunStore::has_init(&db).await? {
|
||||
let has_init = match SlateRunStore::validate_init(&db, record).await {
|
||||
Ok(has_init) => has_init,
|
||||
Err(err) => {
|
||||
let _ = db.close().await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
if !has_init {
|
||||
let _ = db.close().await;
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(
|
||||
SlateRunStore::open(record.run_id.clone(), record.created_at, db).await?,
|
||||
))
|
||||
let run_store = SlateRunStore::open(record.clone(), db).await?;
|
||||
self.cache_active_run(&run_store).await;
|
||||
Ok(Some(run_store))
|
||||
}
|
||||
|
||||
async fn delete_db_prefix(&self, db_prefix: &str) -> Result<()> {
|
||||
|
|
@ -77,6 +141,26 @@ impl Store for SlateStore {
|
|||
) -> Result<Box<dyn RunStore>> {
|
||||
let locator =
|
||||
catalog::read_locator(self.object_store.clone(), &self.base_prefix, run_id).await?;
|
||||
if let Some(active) = self.get_active_run(run_id).await {
|
||||
if active.created_at() != created_at
|
||||
|| locator
|
||||
.as_ref()
|
||||
.is_some_and(|existing| existing.created_at != created_at)
|
||||
{
|
||||
return Err(StoreError::RunAlreadyExists(run_id.to_string()));
|
||||
}
|
||||
let record = active.record();
|
||||
catalog::write_catalog(
|
||||
self.object_store.clone(),
|
||||
&self.base_prefix,
|
||||
run_id,
|
||||
created_at,
|
||||
&record.db_prefix,
|
||||
)
|
||||
.await?;
|
||||
return Ok(Box::new(active));
|
||||
}
|
||||
|
||||
let db_prefix = match locator {
|
||||
Some(existing) if existing.created_at != created_at => {
|
||||
return Err(StoreError::RunAlreadyExists(run_id.to_string()));
|
||||
|
|
@ -92,7 +176,10 @@ impl Store for SlateStore {
|
|||
};
|
||||
|
||||
let db = self.open_db(&db_prefix).await?;
|
||||
SlateRunStore::validate_init(&db, &record).await?;
|
||||
db.put(keys::init(), serde_json::to_vec(&record)?).await?;
|
||||
let run_store = SlateRunStore::open(record.clone(), db).await?;
|
||||
self.cache_active_run(&run_store).await;
|
||||
catalog::write_catalog(
|
||||
self.object_store.clone(),
|
||||
&self.base_prefix,
|
||||
|
|
@ -101,9 +188,7 @@ impl Store for SlateStore {
|
|||
&db_prefix,
|
||||
)
|
||||
.await?;
|
||||
Ok(Box::new(
|
||||
SlateRunStore::open(run_id.to_string(), created_at, db).await?,
|
||||
))
|
||||
Ok(Box::new(run_store))
|
||||
}
|
||||
|
||||
async fn open_run(&self, run_id: &str) -> Result<Option<Box<dyn RunStore>>> {
|
||||
|
|
@ -124,13 +209,28 @@ impl Store for SlateStore {
|
|||
catalog::list_catalogs(self.object_store.clone(), &self.base_prefix, query).await?;
|
||||
let mut summaries = Vec::new();
|
||||
for record in catalogs {
|
||||
let db = self.open_db(&record.db_prefix).await?;
|
||||
if !SlateRunStore::has_init(&db).await? {
|
||||
let _ = db.close().await;
|
||||
if let Some(active) = self.get_active_run(&record.run_id).await {
|
||||
if !active.matches_record(&record) {
|
||||
return Err(StoreError::Other(format!(
|
||||
"active run cache mismatch for run_id {:?}",
|
||||
record.run_id
|
||||
)));
|
||||
}
|
||||
let snapshot = active.snapshot().await?;
|
||||
summaries.push(SlateRunStore::build_summary(snapshot.as_ref(), &record).await?);
|
||||
continue;
|
||||
}
|
||||
let summary = SlateRunStore::build_summary(&db, &record).await?;
|
||||
let _ = db.close().await;
|
||||
if !self.db_prefix_has_objects(&record.db_prefix).await? {
|
||||
continue;
|
||||
}
|
||||
let reader = self.open_reader(&record.db_prefix).await?;
|
||||
if !SlateRunStore::validate_init(&reader, &record).await? {
|
||||
let _ = reader.close().await;
|
||||
continue;
|
||||
}
|
||||
let summary = SlateRunStore::build_summary(&reader, &record).await;
|
||||
let _ = reader.close().await;
|
||||
let summary = summary?;
|
||||
summaries.push(summary);
|
||||
}
|
||||
summaries.sort_by(|a, b| b.created_at.cmp(&a.created_at));
|
||||
|
|
@ -138,6 +238,12 @@ impl Store for SlateStore {
|
|||
}
|
||||
|
||||
async fn delete_run(&self, run_id: &str) -> Result<()> {
|
||||
let active = self.remove_active_run(run_id).await;
|
||||
let active_record = active.as_ref().map(SlateRunStore::record);
|
||||
if let Some(active) = &active {
|
||||
active.close().await?;
|
||||
}
|
||||
|
||||
if let Some(locator) =
|
||||
catalog::read_locator(self.object_store.clone(), &self.base_prefix, run_id).await?
|
||||
{
|
||||
|
|
@ -155,6 +261,21 @@ impl Store for SlateStore {
|
|||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some(record) = active_record {
|
||||
delete_path(
|
||||
self.object_store.clone(),
|
||||
&catalog::by_start_path(&self.base_prefix, record.created_at, run_id),
|
||||
)
|
||||
.await?;
|
||||
self.delete_db_prefix(&record.db_prefix).await?;
|
||||
delete_path(
|
||||
self.object_store.clone(),
|
||||
&catalog::by_id_path(&self.base_prefix, run_id),
|
||||
)
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let by_start_prefix = Path::from(format!("{}by-start", self.base_prefix));
|
||||
let metas = self
|
||||
.object_store
|
||||
|
|
@ -211,6 +332,7 @@ mod tests {
|
|||
RunStatus, RunStatusRecord, StageStatus, StartRecord, StatusReason,
|
||||
};
|
||||
use object_store::memory::InMemory;
|
||||
use slatedb::{CloseReason, ErrorKind};
|
||||
|
||||
use crate::{EventPayload, NodeVisitRef};
|
||||
|
||||
|
|
@ -521,6 +643,39 @@ mod tests {
|
|||
assert!(matches!(conflict, Err(StoreError::RunAlreadyExists(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_runs_and_open_run_reuse_active_handle_without_fencing() {
|
||||
let (_object_store, store) = make_store();
|
||||
let created_at = dt("2026-03-27T12:00:00Z");
|
||||
let run = store.create_run("run-1", created_at).await.unwrap();
|
||||
run.put_run(&sample_run_record("run-1", created_at))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let listed = store.list_runs(&ListRunsQuery::default()).await.unwrap();
|
||||
assert_eq!(listed.len(), 1);
|
||||
|
||||
let reopened = store.open_run("run-1").await.unwrap().unwrap();
|
||||
let first_event = run
|
||||
.append_event(&event_payload("run-1", "2026-03-27T12:00:00Z", "Started"))
|
||||
.await
|
||||
.unwrap();
|
||||
let second_event = reopened
|
||||
.append_event(&event_payload("run-1", "2026-03-27T12:00:01Z", "Continued"))
|
||||
.await
|
||||
.unwrap();
|
||||
let first_checkpoint = run.append_checkpoint(&sample_checkpoint()).await.unwrap();
|
||||
let second_checkpoint = reopened
|
||||
.append_checkpoint(&sample_checkpoint())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first_event, 1);
|
||||
assert_eq!(second_event, 2);
|
||||
assert_eq!(first_checkpoint, 1);
|
||||
assert_eq!(second_checkpoint, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn watch_events_from_polls_new_events() {
|
||||
let (_object_store, store) = make_store();
|
||||
|
|
@ -573,6 +728,26 @@ mod tests {
|
|||
assert!(list_paths(object_store, "runs").await.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_run_closes_active_handles() {
|
||||
let (object_store, store) = make_store();
|
||||
let created_at = dt("2026-03-27T12:00:00Z");
|
||||
let run = store.create_run("run-1", created_at).await.unwrap();
|
||||
run.put_run(&sample_run_record("run-1", created_at))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
store.delete_run("run-1").await.unwrap();
|
||||
|
||||
let err = run.put_graph("digraph night_sky {}").await.unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
StoreError::Slate(err) if matches!(err.kind(), ErrorKind::Closed(CloseReason::Clean))
|
||||
));
|
||||
assert!(store.open_run("run-1").await.unwrap().is_none());
|
||||
assert!(list_paths(object_store, "runs").await.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn repair_catalog_removes_stale_wrong_time_prefixes() {
|
||||
let (object_store, store) = make_store();
|
||||
|
|
@ -601,6 +776,62 @@ mod tests {
|
|||
assert!(paths[0].contains("2026-03-27-12-00/run-1.json"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_run_uses_distinct_db_prefix_for_same_minute_orphan() {
|
||||
let (object_store, store) = make_store();
|
||||
let old_created_at = dt("2026-03-27T12:00:00Z");
|
||||
let new_created_at = dt("2026-03-27T12:00:30Z");
|
||||
let orphan = CatalogRecord {
|
||||
run_id: "run-1".to_string(),
|
||||
created_at: old_created_at,
|
||||
db_prefix: catalog::db_prefix("runs/", old_created_at, "run-1"),
|
||||
};
|
||||
let new_prefix = catalog::db_prefix("runs/", new_created_at, "run-1");
|
||||
assert_ne!(orphan.db_prefix, new_prefix);
|
||||
|
||||
let db = seed_db(object_store.clone(), &orphan, true).await;
|
||||
db.put(keys::graph(), b"stale graph").await.unwrap();
|
||||
db.close().await.unwrap();
|
||||
|
||||
let run = store.create_run("run-1", new_created_at).await.unwrap();
|
||||
assert_eq!(run.get_graph().await.unwrap(), None);
|
||||
|
||||
let locator = catalog::read_locator(object_store, "runs/", "run-1")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(locator.created_at, new_created_at);
|
||||
assert_eq!(locator.db_prefix, new_prefix);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_run_rejects_mismatched_init_for_existing_prefix() {
|
||||
let (object_store, store) = make_store();
|
||||
let created_at = dt("2026-03-27T12:00:00Z");
|
||||
let db_prefix = catalog::db_prefix("runs/", created_at, "run-1");
|
||||
let db = slatedb::Db::open(db_prefix.clone(), object_store)
|
||||
.await
|
||||
.unwrap();
|
||||
let mismatched = CatalogRecord {
|
||||
run_id: "other-run".to_string(),
|
||||
created_at,
|
||||
db_prefix,
|
||||
};
|
||||
db.put(keys::init(), serde_json::to_vec(&mismatched).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
db.close().await.unwrap();
|
||||
|
||||
let err = match store.create_run("run-1", created_at).await {
|
||||
Ok(_) => panic!("expected create_run to reject mismatched _init.json"),
|
||||
Err(err) => err,
|
||||
};
|
||||
assert!(matches!(
|
||||
err,
|
||||
StoreError::Other(message) if message.contains("_init.json")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn slate_run_store_round_trips_node_data_and_assets() {
|
||||
let (_object_store, store) = make_store();
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use std::collections::BTreeSet;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::sync::{Arc, Weak};
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
|
@ -8,7 +9,8 @@ use chrono::{DateTime, Utc};
|
|||
use futures::Stream;
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::mpsc;
|
||||
use slatedb::{CloseReason, DbRead, ErrorKind};
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
use tokio_stream::wrappers::UnboundedReceiverStream;
|
||||
|
||||
use crate::keys;
|
||||
|
|
@ -21,44 +23,99 @@ use fabro_types::{
|
|||
StartRecord,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SlateRunStore {
|
||||
inner: Arc<SlateRunStoreInner>,
|
||||
}
|
||||
|
||||
pub(crate) struct SlateRunStoreInner {
|
||||
run_id: String,
|
||||
created_at: DateTime<Utc>,
|
||||
db_prefix: String,
|
||||
db: slatedb::Db,
|
||||
event_seq: AtomicU32,
|
||||
checkpoint_seq: AtomicU32,
|
||||
close_lock: Mutex<()>,
|
||||
}
|
||||
|
||||
impl SlateRunStore {
|
||||
pub(crate) async fn open(
|
||||
run_id: String,
|
||||
created_at: DateTime<Utc>,
|
||||
db: slatedb::Db,
|
||||
) -> Result<Self> {
|
||||
pub(crate) async fn open(record: CatalogRecord, db: slatedb::Db) -> Result<Self> {
|
||||
let event_seq = recover_next_seq(&db, keys::EVENTS_PREFIX, keys::parse_event_seq).await?;
|
||||
let checkpoint_seq =
|
||||
recover_next_seq(&db, keys::CHECKPOINTS_PREFIX, keys::parse_checkpoint_seq).await?;
|
||||
Ok(Self {
|
||||
run_id,
|
||||
created_at,
|
||||
db,
|
||||
event_seq: AtomicU32::new(event_seq),
|
||||
checkpoint_seq: AtomicU32::new(checkpoint_seq),
|
||||
inner: Arc::new(SlateRunStoreInner {
|
||||
run_id: record.run_id,
|
||||
created_at: record.created_at,
|
||||
db_prefix: record.db_prefix,
|
||||
db,
|
||||
event_seq: AtomicU32::new(event_seq),
|
||||
checkpoint_seq: AtomicU32::new(checkpoint_seq),
|
||||
close_lock: Mutex::new(()),
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn has_init(db: &slatedb::Db) -> Result<bool> {
|
||||
Ok(db.get(keys::init()).await?.is_some())
|
||||
pub(crate) fn from_inner(inner: Arc<SlateRunStoreInner>) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub(crate) async fn build_summary(
|
||||
db: &slatedb::Db,
|
||||
catalog: &CatalogRecord,
|
||||
) -> Result<RunSummary> {
|
||||
let run = get_json::<RunRecord>(db, keys::run()).await?;
|
||||
let start = get_json::<StartRecord>(db, keys::start()).await?;
|
||||
let status = get_json::<RunStatusRecord>(db, keys::status()).await?;
|
||||
let conclusion = get_json::<Conclusion>(db, keys::conclusion()).await?;
|
||||
pub(crate) fn downgrade(&self) -> Weak<SlateRunStoreInner> {
|
||||
Arc::downgrade(&self.inner)
|
||||
}
|
||||
|
||||
pub(crate) fn record(&self) -> CatalogRecord {
|
||||
CatalogRecord {
|
||||
run_id: self.inner.run_id.clone(),
|
||||
created_at: self.inner.created_at,
|
||||
db_prefix: self.inner.db_prefix.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn matches_record(&self, record: &CatalogRecord) -> bool {
|
||||
self.inner.run_id == record.run_id
|
||||
&& self.inner.created_at == record.created_at
|
||||
&& self.inner.db_prefix == record.db_prefix
|
||||
}
|
||||
|
||||
pub(crate) fn created_at(&self) -> DateTime<Utc> {
|
||||
self.inner.created_at
|
||||
}
|
||||
|
||||
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()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn snapshot(&self) -> Result<Arc<slatedb::DbSnapshot>> {
|
||||
Ok(self.inner.db.snapshot().await?)
|
||||
}
|
||||
|
||||
pub(crate) async fn validate_init<R>(db: &R, expected: &CatalogRecord) -> Result<bool>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
match get_json::<R, CatalogRecord>(db, keys::init()).await? {
|
||||
Some(existing) if existing == *expected => Ok(true),
|
||||
Some(existing) => Err(StoreError::Other(format!(
|
||||
"existing _init.json {existing:?} does not match requested catalog {expected:?}"
|
||||
))),
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn build_summary<R>(db: &R, catalog: &CatalogRecord) -> Result<RunSummary>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
let run = get_json::<_, RunRecord>(db, keys::run()).await?;
|
||||
let start = get_json::<_, StartRecord>(db, keys::start()).await?;
|
||||
let status = get_json::<_, RunStatusRecord>(db, keys::status()).await?;
|
||||
let conclusion = get_json::<_, Conclusion>(db, keys::conclusion()).await?;
|
||||
|
||||
let workflow_name = run.as_ref().map(|run| {
|
||||
if run.graph.name.is_empty() {
|
||||
|
|
@ -93,16 +150,16 @@ impl SlateRunStore {
|
|||
}
|
||||
|
||||
fn validate_run_record(&self, record: &RunRecord) -> Result<()> {
|
||||
if record.created_at != self.created_at {
|
||||
if record.created_at != self.inner.created_at {
|
||||
return Err(StoreError::Other(format!(
|
||||
"run record created_at {:?} does not match store created_at {:?}",
|
||||
record.created_at, self.created_at
|
||||
record.created_at, self.inner.created_at
|
||||
)));
|
||||
}
|
||||
if record.run_id != self.run_id {
|
||||
if record.run_id != self.inner.run_id {
|
||||
return Err(StoreError::Other(format!(
|
||||
"run record run_id {:?} does not match store run_id {:?}",
|
||||
record.run_id, self.run_id
|
||||
record.run_id, self.inner.run_id
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
|
|
@ -112,11 +169,11 @@ impl SlateRunStore {
|
|||
Ok(NodeSnapshot {
|
||||
node_id: node.node_id.to_string(),
|
||||
visit: node.visit,
|
||||
prompt: get_text(&self.db, &keys::node_prompt(node)).await?,
|
||||
response: get_text(&self.db, &keys::node_response(node)).await?,
|
||||
status: get_json(&self.db, &keys::node_status(node)).await?,
|
||||
stdout: get_text(&self.db, &keys::node_stdout(node)).await?,
|
||||
stderr: get_text(&self.db, &keys::node_stderr(node)).await?,
|
||||
prompt: get_text(&self.inner.db, &keys::node_prompt(node)).await?,
|
||||
response: get_text(&self.inner.db, &keys::node_response(node)).await?,
|
||||
status: get_json(&self.inner.db, &keys::node_status(node)).await?,
|
||||
stdout: get_text(&self.inner.db, &keys::node_stdout(node)).await?,
|
||||
stderr: get_text(&self.inner.db, &keys::node_stderr(node)).await?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -125,42 +182,42 @@ impl SlateRunStore {
|
|||
impl RunStore for SlateRunStore {
|
||||
async fn put_run(&self, record: &RunRecord) -> Result<()> {
|
||||
self.validate_run_record(record)?;
|
||||
put_json(&self.db, keys::run(), record).await
|
||||
put_json(&self.inner.db, keys::run(), record).await
|
||||
}
|
||||
|
||||
async fn get_run(&self) -> Result<Option<RunRecord>> {
|
||||
get_json(&self.db, keys::run()).await
|
||||
get_json(&self.inner.db, keys::run()).await
|
||||
}
|
||||
|
||||
async fn put_start(&self, record: &StartRecord) -> Result<()> {
|
||||
put_json(&self.db, keys::start(), record).await
|
||||
put_json(&self.inner.db, keys::start(), record).await
|
||||
}
|
||||
|
||||
async fn get_start(&self) -> Result<Option<StartRecord>> {
|
||||
get_json(&self.db, keys::start()).await
|
||||
get_json(&self.inner.db, keys::start()).await
|
||||
}
|
||||
|
||||
async fn put_status(&self, record: &RunStatusRecord) -> Result<()> {
|
||||
put_json(&self.db, keys::status(), record).await
|
||||
put_json(&self.inner.db, keys::status(), record).await
|
||||
}
|
||||
|
||||
async fn get_status(&self) -> Result<Option<RunStatusRecord>> {
|
||||
get_json(&self.db, keys::status()).await
|
||||
get_json(&self.inner.db, keys::status()).await
|
||||
}
|
||||
|
||||
async fn put_checkpoint(&self, record: &Checkpoint) -> Result<()> {
|
||||
put_json(&self.db, keys::checkpoint(), record).await
|
||||
put_json(&self.inner.db, keys::checkpoint(), record).await
|
||||
}
|
||||
|
||||
async fn get_checkpoint(&self) -> Result<Option<Checkpoint>> {
|
||||
get_json(&self.db, keys::checkpoint()).await
|
||||
get_json(&self.inner.db, keys::checkpoint()).await
|
||||
}
|
||||
|
||||
async fn append_checkpoint(&self, record: &Checkpoint) -> Result<u32> {
|
||||
let seq = self.checkpoint_seq.fetch_add(1, Ordering::SeqCst);
|
||||
let seq = self.inner.checkpoint_seq.fetch_add(1, Ordering::SeqCst);
|
||||
self.put_checkpoint(record).await?;
|
||||
put_json(
|
||||
&self.db,
|
||||
&self.inner.db,
|
||||
&keys::checkpoint_history_key(seq, Utc::now().timestamp_millis()),
|
||||
record,
|
||||
)
|
||||
|
|
@ -169,47 +226,47 @@ impl RunStore for SlateRunStore {
|
|||
}
|
||||
|
||||
async fn list_checkpoints(&self) -> Result<Vec<(u32, Checkpoint)>> {
|
||||
list_checkpoints(&self.db).await
|
||||
list_checkpoints(&self.inner.db).await
|
||||
}
|
||||
|
||||
async fn put_conclusion(&self, record: &Conclusion) -> Result<()> {
|
||||
put_json(&self.db, keys::conclusion(), record).await
|
||||
put_json(&self.inner.db, keys::conclusion(), record).await
|
||||
}
|
||||
|
||||
async fn get_conclusion(&self) -> Result<Option<Conclusion>> {
|
||||
get_json(&self.db, keys::conclusion()).await
|
||||
get_json(&self.inner.db, keys::conclusion()).await
|
||||
}
|
||||
|
||||
async fn put_retro(&self, retro: &Retro) -> Result<()> {
|
||||
put_json(&self.db, keys::retro(), retro).await
|
||||
put_json(&self.inner.db, keys::retro(), retro).await
|
||||
}
|
||||
|
||||
async fn get_retro(&self) -> Result<Option<Retro>> {
|
||||
get_json(&self.db, keys::retro()).await
|
||||
get_json(&self.inner.db, keys::retro()).await
|
||||
}
|
||||
|
||||
async fn put_graph(&self, dot_source: &str) -> Result<()> {
|
||||
put_text(&self.db, keys::graph(), dot_source).await
|
||||
put_text(&self.inner.db, keys::graph(), dot_source).await
|
||||
}
|
||||
|
||||
async fn get_graph(&self) -> Result<Option<String>> {
|
||||
get_text(&self.db, keys::graph()).await
|
||||
get_text(&self.inner.db, keys::graph()).await
|
||||
}
|
||||
|
||||
async fn put_sandbox(&self, record: &SandboxRecord) -> Result<()> {
|
||||
put_json(&self.db, keys::sandbox(), record).await
|
||||
put_json(&self.inner.db, keys::sandbox(), record).await
|
||||
}
|
||||
|
||||
async fn get_sandbox(&self) -> Result<Option<SandboxRecord>> {
|
||||
get_json(&self.db, keys::sandbox()).await
|
||||
get_json(&self.inner.db, keys::sandbox()).await
|
||||
}
|
||||
|
||||
async fn put_node_prompt(&self, node: &NodeVisitRef<'_>, prompt: &str) -> Result<()> {
|
||||
put_text(&self.db, &keys::node_prompt(node), prompt).await
|
||||
put_text(&self.inner.db, &keys::node_prompt(node), prompt).await
|
||||
}
|
||||
|
||||
async fn put_node_response(&self, node: &NodeVisitRef<'_>, response: &str) -> Result<()> {
|
||||
put_text(&self.db, &keys::node_response(node), response).await
|
||||
put_text(&self.inner.db, &keys::node_response(node), response).await
|
||||
}
|
||||
|
||||
async fn put_node_status(
|
||||
|
|
@ -217,15 +274,15 @@ impl RunStore for SlateRunStore {
|
|||
node: &NodeVisitRef<'_>,
|
||||
status: &NodeStatusRecord,
|
||||
) -> Result<()> {
|
||||
put_json(&self.db, &keys::node_status(node), status).await
|
||||
put_json(&self.inner.db, &keys::node_status(node), status).await
|
||||
}
|
||||
|
||||
async fn put_node_stdout(&self, node: &NodeVisitRef<'_>, log: &str) -> Result<()> {
|
||||
put_text(&self.db, &keys::node_stdout(node), log).await
|
||||
put_text(&self.inner.db, &keys::node_stdout(node), log).await
|
||||
}
|
||||
|
||||
async fn put_node_stderr(&self, node: &NodeVisitRef<'_>, log: &str) -> Result<()> {
|
||||
put_text(&self.db, &keys::node_stderr(node), log).await
|
||||
put_text(&self.inner.db, &keys::node_stderr(node), log).await
|
||||
}
|
||||
|
||||
async fn get_node(&self, node: &NodeVisitRef<'_>) -> Result<NodeSnapshot> {
|
||||
|
|
@ -234,7 +291,7 @@ impl RunStore for SlateRunStore {
|
|||
|
||||
async fn list_node_visits(&self, node_id: &str) -> Result<Vec<u32>> {
|
||||
let prefix = format!("nodes/{node_id}/visit-");
|
||||
let mut iter = self.db.scan_prefix(prefix.as_bytes()).await?;
|
||||
let mut iter = self.inner.db.scan_prefix(prefix.as_bytes()).await?;
|
||||
let mut visits = BTreeSet::new();
|
||||
while let Some(entry) = iter.next().await? {
|
||||
let key = key_to_string(entry.key)?;
|
||||
|
|
@ -248,10 +305,10 @@ impl RunStore for SlateRunStore {
|
|||
}
|
||||
|
||||
async fn append_event(&self, payload: &EventPayload) -> Result<u32> {
|
||||
payload.validate(&self.run_id)?;
|
||||
let seq = self.event_seq.fetch_add(1, Ordering::SeqCst);
|
||||
payload.validate(&self.inner.run_id)?;
|
||||
let seq = self.inner.event_seq.fetch_add(1, Ordering::SeqCst);
|
||||
put_json(
|
||||
&self.db,
|
||||
&self.inner.db,
|
||||
&keys::event_key(seq, Utc::now().timestamp_millis()),
|
||||
payload,
|
||||
)
|
||||
|
|
@ -260,18 +317,18 @@ impl RunStore for SlateRunStore {
|
|||
}
|
||||
|
||||
async fn list_events(&self) -> Result<Vec<EventEnvelope>> {
|
||||
list_events_from(&self.db, 1).await
|
||||
list_events_from(&self.inner.db, 1).await
|
||||
}
|
||||
|
||||
async fn list_events_from(&self, seq: u32) -> Result<Vec<EventEnvelope>> {
|
||||
list_events_from(&self.db, seq).await
|
||||
list_events_from(&self.inner.db, seq).await
|
||||
}
|
||||
|
||||
async fn watch_events_from(
|
||||
&self,
|
||||
seq: u32,
|
||||
) -> Result<std::pin::Pin<Box<dyn Stream<Item = Result<EventEnvelope>> + Send>>> {
|
||||
let db = self.db.clone();
|
||||
let db = self.inner.db.clone();
|
||||
let (sender, receiver) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
|
|
@ -306,40 +363,40 @@ impl RunStore for SlateRunStore {
|
|||
}
|
||||
|
||||
async fn put_retro_prompt(&self, text: &str) -> Result<()> {
|
||||
put_text(&self.db, keys::retro_prompt(), text).await
|
||||
put_text(&self.inner.db, keys::retro_prompt(), text).await
|
||||
}
|
||||
|
||||
async fn get_retro_prompt(&self) -> Result<Option<String>> {
|
||||
get_text(&self.db, keys::retro_prompt()).await
|
||||
get_text(&self.inner.db, keys::retro_prompt()).await
|
||||
}
|
||||
|
||||
async fn put_retro_response(&self, text: &str) -> Result<()> {
|
||||
put_text(&self.db, keys::retro_response(), text).await
|
||||
put_text(&self.inner.db, keys::retro_response(), text).await
|
||||
}
|
||||
|
||||
async fn get_retro_response(&self) -> Result<Option<String>> {
|
||||
get_text(&self.db, keys::retro_response()).await
|
||||
get_text(&self.inner.db, keys::retro_response()).await
|
||||
}
|
||||
|
||||
async fn put_artifact_value(&self, artifact_id: &str, value: &serde_json::Value) -> Result<()> {
|
||||
put_json(&self.db, &keys::artifact_value(artifact_id), value).await
|
||||
put_json(&self.inner.db, &keys::artifact_value(artifact_id), value).await
|
||||
}
|
||||
|
||||
async fn get_artifact_value(&self, artifact_id: &str) -> Result<Option<serde_json::Value>> {
|
||||
get_json(&self.db, &keys::artifact_value(artifact_id)).await
|
||||
get_json(&self.inner.db, &keys::artifact_value(artifact_id)).await
|
||||
}
|
||||
|
||||
async fn put_asset(&self, node: &NodeVisitRef<'_>, filename: &str, data: &[u8]) -> Result<()> {
|
||||
put_bytes(&self.db, &keys::node_asset(node, filename), data).await
|
||||
put_bytes(&self.inner.db, &keys::node_asset(node, filename), data).await
|
||||
}
|
||||
|
||||
async fn get_asset(&self, node: &NodeVisitRef<'_>, filename: &str) -> Result<Option<Bytes>> {
|
||||
get_bytes(&self.db, &keys::node_asset(node, filename)).await
|
||||
get_bytes(&self.inner.db, &keys::node_asset(node, filename)).await
|
||||
}
|
||||
|
||||
async fn list_assets(&self, node: &NodeVisitRef<'_>) -> Result<Vec<String>> {
|
||||
let prefix = format!("{}/", keys::node_asset_prefix(node));
|
||||
let mut iter = self.db.scan_prefix(prefix.as_bytes()).await?;
|
||||
let mut iter = self.inner.db.scan_prefix(prefix.as_bytes()).await?;
|
||||
let mut assets = Vec::new();
|
||||
while let Some(entry) = iter.next().await? {
|
||||
let key = key_to_string(entry.key)?;
|
||||
|
|
@ -356,7 +413,7 @@ impl RunStore for SlateRunStore {
|
|||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut iter = self.db.scan_prefix(b"nodes/").await?;
|
||||
let mut iter = self.inner.db.scan_prefix(b"nodes/").await?;
|
||||
let mut visits = BTreeSet::new();
|
||||
while let Some(entry) = iter.next().await? {
|
||||
let key = key_to_string(entry.key)?;
|
||||
|
|
@ -393,7 +450,11 @@ async fn put_json<T: Serialize>(db: &slatedb::Db, key: &str, value: &T) -> Resul
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_json<T: DeserializeOwned>(db: &slatedb::Db, key: &str) -> Result<Option<T>> {
|
||||
async fn get_json<R, T>(db: &R, key: &str) -> Result<Option<T>>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
T: DeserializeOwned,
|
||||
{
|
||||
db.get(key)
|
||||
.await?
|
||||
.map(|value| serde_json::from_slice(&value))
|
||||
|
|
@ -406,7 +467,10 @@ async fn put_text(db: &slatedb::Db, key: &str, value: &str) -> Result<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_text(db: &slatedb::Db, key: &str) -> Result<Option<String>> {
|
||||
async fn get_text<R>(db: &R, key: &str) -> Result<Option<String>>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
db.get(key)
|
||||
.await?
|
||||
.map(|value| {
|
||||
|
|
@ -425,11 +489,10 @@ async fn get_bytes(db: &slatedb::Db, key: &str) -> Result<Option<Bytes>> {
|
|||
Ok(db.get(key).await?)
|
||||
}
|
||||
|
||||
async fn recover_next_seq(
|
||||
db: &slatedb::Db,
|
||||
prefix: &str,
|
||||
parse: fn(&str) -> Option<u32>,
|
||||
) -> Result<u32> {
|
||||
async fn recover_next_seq<R>(db: &R, prefix: &str, parse: fn(&str) -> Option<u32>) -> Result<u32>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
let mut iter = db.scan_prefix(prefix.as_bytes()).await?;
|
||||
let mut max_seq = 0;
|
||||
while let Some(entry) = iter.next().await? {
|
||||
|
|
@ -441,7 +504,10 @@ async fn recover_next_seq(
|
|||
Ok(max_seq.saturating_add(1).max(1))
|
||||
}
|
||||
|
||||
async fn list_events_from(db: &slatedb::Db, start_seq: u32) -> Result<Vec<EventEnvelope>> {
|
||||
async fn list_events_from<R>(db: &R, start_seq: u32) -> Result<Vec<EventEnvelope>>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
let mut iter = db.scan_prefix(keys::EVENTS_PREFIX.as_bytes()).await?;
|
||||
let mut events = Vec::new();
|
||||
while let Some(entry) = iter.next().await? {
|
||||
|
|
@ -461,7 +527,10 @@ async fn list_events_from(db: &slatedb::Db, start_seq: u32) -> Result<Vec<EventE
|
|||
Ok(events)
|
||||
}
|
||||
|
||||
async fn list_checkpoints(db: &slatedb::Db) -> Result<Vec<(u32, Checkpoint)>> {
|
||||
async fn list_checkpoints<R>(db: &R) -> Result<Vec<(u32, Checkpoint)>>
|
||||
where
|
||||
R: DbRead + Sync,
|
||||
{
|
||||
let mut iter = db.scan_prefix(keys::CHECKPOINTS_PREFIX.as_bytes()).await?;
|
||||
let mut checkpoints = Vec::new();
|
||||
while let Some(entry) = iter.next().await? {
|
||||
|
|
|
|||
|
|
@ -6,9 +6,9 @@ use crate::outcome::StageStatus;
|
|||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct NodeStatusRecord {
|
||||
pub status: StageStatus,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
#[serde(default)]
|
||||
pub notes: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
#[serde(default)]
|
||||
pub failure_reason: Option<String>,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -137,4 +137,20 @@ mod tests {
|
|||
);
|
||||
assert!(value.get("timestamp").and_then(|v| v.as_str()).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_node_status_preserves_null_optional_fields() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let outcome = Outcome {
|
||||
status: StageStatus::Success,
|
||||
..Outcome::default()
|
||||
};
|
||||
|
||||
write_node_status(temp.path(), "work", 1, &outcome);
|
||||
|
||||
let data = std::fs::read_to_string(temp.path().join("nodes/work/status.json")).unwrap();
|
||||
let value: serde_json::Value = serde_json::from_str(&data).unwrap();
|
||||
assert_eq!(value.get("notes"), Some(&serde_json::Value::Null));
|
||||
assert_eq!(value.get("failure_reason"), Some(&serde_json::Value::Null));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue