pub mod postgres;
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::SystemTime;
use async_trait::async_trait;
use super::catalog::{
CatalogContentId, RawPayload, Refusable, Refusal, RefusalReason, SchemaVersion, SourceSnapshot,
SourceValidators,
};
use super::models_dev::{ModelsDevAdapter, ModelsDevError};
use super::{BackendFailure, Capabilities, FailureCategory};
use crate::backends::catalog::CatalogSnapshot;
use crate::desired_state::BlobError;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedCatalog {
pub source: SourceSnapshot,
pub payload: RawPayload,
}
impl RetainedCatalog {
pub fn content_id(&self) -> CatalogContentId {
self.source.content_id
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct StoredCatalogState {
pub active: Option<RetainedCatalog>,
pub consecutive_refusals: u32,
pub last_refusal: Option<RefusalReason>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Retention {
Retained,
AlreadyRetained,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum CatalogStoreError {
#[error("catalogue store `{backend}` unavailable: {message}")]
Unavailable {
backend: &'static str,
message: String,
},
#[error("catalogue store `{backend}` holds a record it cannot answer with: {message}")]
Corrupt {
backend: &'static str,
message: String,
},
#[error("catalogue store `{backend}` refused the operation: {message}")]
Denied {
backend: &'static str,
message: String,
},
}
impl CatalogStoreError {
pub fn unavailable(backend: &'static str, message: impl Into<String>) -> Self {
Self::Unavailable {
backend,
message: message.into(),
}
}
pub fn corrupt(backend: &'static str, message: impl Into<String>) -> Self {
Self::Corrupt {
backend,
message: message.into(),
}
}
pub fn denied(backend: &'static str, message: impl Into<String>) -> Self {
Self::Denied {
backend,
message: message.into(),
}
}
}
impl BackendFailure for CatalogStoreError {
fn category(&self) -> FailureCategory {
match self {
Self::Unavailable { .. } => FailureCategory::Unavailable,
Self::Corrupt { .. } => FailureCategory::Corrupt,
Self::Denied { .. } => FailureCategory::Denied,
}
}
}
impl Refusable for CatalogStoreError {
fn refusal(&self) -> Refusal {
Refusal::new(RefusalReason::NotRetained)
}
}
#[async_trait]
pub trait CatalogStore: Send + Sync {
fn name(&self) -> &'static str;
fn capabilities(&self) -> Capabilities;
async fn load(&self) -> Result<StoredCatalogState, CatalogStoreError>;
async fn retained(
&self,
content_id: CatalogContentId,
) -> Result<Option<RetainedCatalog>, CatalogStoreError>;
async fn activate(
&self,
import: &RetainedCatalog,
activated_at: SystemTime,
) -> Result<Retention, CatalogStoreError>;
async fn confirm(
&self,
content_id: CatalogContentId,
validators: &SourceValidators,
confirmed_at: SystemTime,
) -> Result<bool, CatalogStoreError>;
async fn refuse(
&self,
reason: RefusalReason,
refused_at: SystemTime,
) -> Result<(), CatalogStoreError>;
}
#[async_trait]
impl<T: CatalogStore + ?Sized> CatalogStore for &T {
fn name(&self) -> &'static str {
(**self).name()
}
fn capabilities(&self) -> Capabilities {
(**self).capabilities()
}
async fn load(&self) -> Result<StoredCatalogState, CatalogStoreError> {
(**self).load().await
}
async fn retained(
&self,
content_id: CatalogContentId,
) -> Result<Option<RetainedCatalog>, CatalogStoreError> {
(**self).retained(content_id).await
}
async fn activate(
&self,
import: &RetainedCatalog,
activated_at: SystemTime,
) -> Result<Retention, CatalogStoreError> {
(**self).activate(import, activated_at).await
}
async fn confirm(
&self,
content_id: CatalogContentId,
validators: &SourceValidators,
confirmed_at: SystemTime,
) -> Result<bool, CatalogStoreError> {
(**self).confirm(content_id, validators, confirmed_at).await
}
async fn refuse(
&self,
reason: RefusalReason,
refused_at: SystemTime,
) -> Result<(), CatalogStoreError> {
(**self).refuse(reason, refused_at).await
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum HydrationError {
#[error(
"catalogue {content_id} was imported under schema `{}`, which this build does not read",
schema.as_str()
)]
UnknownSchema {
content_id: CatalogContentId,
schema: SchemaVersion,
},
#[error("the stored payload for catalogue {content_id} is not the bytes it names: {source}")]
Payload {
content_id: CatalogContentId,
#[source]
source: BlobError,
},
#[error("the stored payload for catalogue {content_id} no longer parses: {source}")]
Parse {
content_id: CatalogContentId,
#[source]
source: ModelsDevError,
},
#[error("the stored payload for catalogue {content_id} now normalizes to {recomputed}")]
Drift {
content_id: CatalogContentId,
recomputed: CatalogContentId,
},
}
impl Refusable for HydrationError {
fn refusal(&self) -> Refusal {
match self {
Self::Parse { source, .. } => source.refusal(),
Self::UnknownSchema { .. } | Self::Payload { .. } | Self::Drift { .. } => {
Refusal::new(RefusalReason::NotRetained)
}
}
}
}
pub fn hydrate(retained: &RetainedCatalog) -> Result<CatalogSnapshot, HydrationError> {
let content_id = retained.source.content_id;
if retained.source.schema_version != SchemaVersion::MODELS_DEV_CATALOG_V1 {
return Err(HydrationError::UnknownSchema {
content_id,
schema: retained.source.schema_version,
});
}
retained
.source
.raw
.verify(retained.payload.as_bytes())
.map_err(|source| HydrationError::Payload { content_id, source })?;
let adapter = ModelsDevAdapter::new(retained.source.source_url.clone())
.map_err(|source| HydrationError::Parse { content_id, source })?;
let snapshot = adapter
.parse(
retained.payload.as_bytes(),
retained.source.validators.clone(),
retained.source.fetched_at,
)
.map_err(|source| HydrationError::Parse { content_id, source })?;
if snapshot.source.content_id != content_id {
return Err(HydrationError::Drift {
content_id,
recomputed: snapshot.source.content_id,
});
}
Ok(snapshot)
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ActivePointer {
content_id: CatalogContentId,
validators: SourceValidators,
confirmed_at: SystemTime,
}
#[derive(Debug, Default)]
struct InMemoryState {
retained: BTreeMap<CatalogContentId, RetainedCatalog>,
active: Option<ActivePointer>,
consecutive_refusals: u32,
last_refusal: Option<RefusalReason>,
}
#[derive(Debug, Default)]
pub struct InMemoryCatalogStore {
state: Mutex<InMemoryState>,
}
const IN_MEMORY: &str = "in-memory";
impl InMemoryCatalogStore {
pub fn new() -> Self {
Self::default()
}
pub fn retained_count(&self) -> usize {
self.state
.lock()
.expect("catalogue store lock")
.retained
.len()
}
fn locked(&self) -> std::sync::MutexGuard<'_, InMemoryState> {
self.state.lock().expect("catalogue store lock")
}
}
#[async_trait]
impl CatalogStore for InMemoryCatalogStore {
fn name(&self) -> &'static str {
IN_MEMORY
}
fn capabilities(&self) -> Capabilities {
Capabilities::NONE
}
async fn load(&self) -> Result<StoredCatalogState, CatalogStoreError> {
let state = self.locked();
let active = state
.active
.as_ref()
.map(|pointer| {
let retained = state.retained.get(&pointer.content_id).ok_or_else(|| {
CatalogStoreError::corrupt(
IN_MEMORY,
format!("active catalogue {} is not retained", pointer.content_id),
)
})?;
let mut retained = retained.clone();
retained.source.validators = pointer.validators.clone();
retained.source.fetched_at = pointer.confirmed_at;
Ok::<_, CatalogStoreError>(retained)
})
.transpose()?;
Ok(StoredCatalogState {
active,
consecutive_refusals: state.consecutive_refusals,
last_refusal: state.last_refusal,
})
}
async fn retained(
&self,
content_id: CatalogContentId,
) -> Result<Option<RetainedCatalog>, CatalogStoreError> {
Ok(self.locked().retained.get(&content_id).cloned())
}
async fn activate(
&self,
import: &RetainedCatalog,
activated_at: SystemTime,
) -> Result<Retention, CatalogStoreError> {
let mut state = self.locked();
let content_id = import.content_id();
let retention = match state.retained.entry(content_id) {
std::collections::btree_map::Entry::Occupied(_) => Retention::AlreadyRetained,
std::collections::btree_map::Entry::Vacant(slot) => {
slot.insert(import.clone());
Retention::Retained
}
};
let mut validators = import.source.validators.clone();
if let Some(active) = state
.active
.as_ref()
.filter(|active| active.content_id == content_id)
{
let mut held = active.validators.clone();
held.carry_over(validators);
validators = held;
}
state.active = Some(ActivePointer {
content_id,
validators,
confirmed_at: activated_at,
});
state.consecutive_refusals = 0;
state.last_refusal = None;
Ok(retention)
}
async fn confirm(
&self,
content_id: CatalogContentId,
validators: &SourceValidators,
confirmed_at: SystemTime,
) -> Result<bool, CatalogStoreError> {
let mut state = self.locked();
let Some(pointer) = state
.active
.as_mut()
.filter(|pointer| pointer.content_id == content_id)
else {
return Ok(false);
};
pointer.validators.carry_over(validators.clone());
pointer.confirmed_at = confirmed_at;
state.consecutive_refusals = 0;
state.last_refusal = None;
Ok(true)
}
async fn refuse(
&self,
reason: RefusalReason,
_refused_at: SystemTime,
) -> Result<(), CatalogStoreError> {
let mut state = self.locked();
state.consecutive_refusals = state.consecutive_refusals.saturating_add(1);
state.last_refusal = Some(reason);
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::super::models_dev::{SEED_PAYLOAD, seed_snapshot};
use super::*;
fn seed_import() -> RetainedCatalog {
RetainedCatalog {
source: seed_snapshot().source,
payload: RawPayload::new(SEED_PAYLOAD.as_bytes()),
}
}
fn at(seconds: u64) -> SystemTime {
SystemTime::UNIX_EPOCH + Duration::from_secs(seconds)
}
#[tokio::test]
async fn a_retained_import_rehydrates_into_the_snapshot_it_was() {
let import = seed_import();
let hydrated = hydrate(&import).expect("the seed rehydrates");
assert_eq!(hydrated, seed_snapshot());
}
#[tokio::test]
async fn a_payload_that_is_not_the_bytes_it_names_is_refused() {
let mut import = seed_import();
import.payload = RawPayload::new(&b"{\"models\":{},\"providers\":{}}"[..]);
let error = hydrate(&import).expect_err("the digest no longer matches");
assert!(matches!(error, HydrationError::Payload { .. }));
assert_eq!(error.refusal().reason(), RefusalReason::NotRetained);
}
#[tokio::test]
async fn a_record_stored_under_an_identity_its_bytes_do_not_produce_is_refused() {
let mut import = seed_import();
let elsewhere = crate::desired_state::Checksum::of(b"another catalogue");
import.source.content_id = CatalogContentId::from_checksum(elsewhere);
let error = hydrate(&import).expect_err("the content id does not match the bytes");
let HydrationError::Drift { recomputed, .. } = error else {
panic!("a mismatched identity is drift, not a parse failure: {error}");
};
assert_eq!(recomputed, seed_snapshot().source.content_id);
}
#[tokio::test]
async fn an_import_this_build_cannot_parse_is_named_by_its_schema() {
let mut import = seed_import();
import.source.source_url = "https://models.dev/api.json".to_owned();
let error = hydrate(&import).expect_err("the URL is not a catalogue document");
assert!(matches!(error, HydrationError::Parse { .. }));
}
#[tokio::test]
async fn re_importing_unchanged_content_retains_nothing_new() {
let store = InMemoryCatalogStore::new();
let import = seed_import();
assert_eq!(
store.activate(&import, at(10)).await.expect("activate"),
Retention::Retained
);
assert_eq!(
store.activate(&import, at(20)).await.expect("re-activate"),
Retention::AlreadyRetained
);
assert_eq!(store.retained_count(), 1);
let state = store.load().await.expect("load");
let active = state.active.expect("an active catalogue");
assert_eq!(active.content_id(), import.content_id());
assert_eq!(
active.source.fetched_at,
at(20),
"the second import is when the content was last confirmed"
);
}
#[tokio::test]
async fn confirming_moves_the_check_time_without_touching_the_import() {
let store = InMemoryCatalogStore::new();
let import = seed_import();
store.activate(&import, at(10)).await.expect("activate");
let confirmed = store
.confirm(
import.content_id(),
&SourceValidators::etag("\"later\""),
at(600),
)
.await
.expect("confirm");
assert!(confirmed);
let state = store.load().await.expect("load");
let active = state.active.expect("an active catalogue");
assert_eq!(active.source.fetched_at, at(600));
assert_eq!(
active.source.validators,
SourceValidators::etag("\"later\"")
);
assert_eq!(
store
.retained(import.content_id())
.await
.expect("retained")
.expect("the import itself")
.source
.validators,
import.source.validators,
"the immutable import keeps the validators it arrived with"
);
}
#[tokio::test]
async fn a_validator_the_answer_does_not_state_is_kept() {
let store = InMemoryCatalogStore::new();
let mut import = seed_import();
import.source.validators = SourceValidators::etag("\"held\"");
store.activate(&import, at(10)).await.expect("activate");
store
.confirm(import.content_id(), &SourceValidators::default(), at(20))
.await
.expect("confirm");
let state = store.load().await.expect("load");
assert_eq!(
state.active.expect("active").source.validators,
SourceValidators::etag("\"held\""),
"an unstated validator is not a withdrawn one"
);
}
#[tokio::test]
async fn re_activating_the_active_content_without_a_validator_keeps_the_held_one() {
let store = InMemoryCatalogStore::new();
let mut import = seed_import();
import.source.validators = SourceValidators::etag("\"held\"");
store.activate(&import, at(10)).await.expect("activate");
let mut stripped = import.clone();
stripped.source.validators = SourceValidators::default();
store
.activate(&stripped, at(20))
.await
.expect("re-activate");
let state = store.load().await.expect("load");
assert_eq!(
state.active.expect("active").source.validators,
SourceValidators::etag("\"held\""),
);
}
#[tokio::test]
async fn activating_new_content_does_not_inherit_the_previous_validator() {
let store = InMemoryCatalogStore::new();
let mut first = seed_import();
first.source.validators = SourceValidators::etag("\"held\"");
store.activate(&first, at(10)).await.expect("activate");
let mut second = seed_import();
second.source.content_id =
CatalogContentId::from_checksum(crate::desired_state::Checksum::of(b"other"));
second.source.validators = SourceValidators::default();
store.activate(&second, at(20)).await.expect("activate");
let state = store.load().await.expect("load");
let active = state.active.expect("active");
assert_eq!(active.content_id(), second.content_id());
assert_eq!(active.source.validators, SourceValidators::default());
}
#[tokio::test]
async fn confirming_content_that_is_not_active_records_nothing() {
let store = InMemoryCatalogStore::new();
let confirmed = store
.confirm(
seed_import().content_id(),
&SourceValidators::etag("\"any\""),
at(10),
)
.await
.expect("confirm");
assert!(!confirmed, "there was no active pointer to move");
assert_eq!(
store.load().await.expect("load"),
StoredCatalogState::default()
);
}
#[tokio::test]
async fn refusals_are_counted_durably_and_cleared_by_an_import() {
let store = InMemoryCatalogStore::new();
store
.refuse(RefusalReason::Unreachable, at(10))
.await
.expect("refuse");
store
.refuse(RefusalReason::Schema, at(20))
.await
.expect("refuse");
let state = store.load().await.expect("load");
assert_eq!(state.consecutive_refusals, 2);
assert_eq!(state.last_refusal, Some(RefusalReason::Schema));
assert!(state.active.is_none(), "nothing was ever imported");
store
.activate(&seed_import(), at(30))
.await
.expect("activate");
let state = store.load().await.expect("load");
assert_eq!(state.consecutive_refusals, 0);
assert_eq!(state.last_refusal, None);
}
#[test]
fn every_storage_failure_refuses_the_import_without_blaming_the_payload() {
for error in [
CatalogStoreError::unavailable(IN_MEMORY, "connection refused"),
CatalogStoreError::corrupt(IN_MEMORY, "active row names an absent snapshot"),
CatalogStoreError::denied(IN_MEMORY, "no privilege on axond_catalog_snapshot"),
] {
assert_eq!(error.refusal().reason(), RefusalReason::NotRetained);
}
assert!(CatalogStoreError::unavailable(IN_MEMORY, "down").retryable());
assert!(!CatalogStoreError::corrupt(IN_MEMORY, "damaged").retryable());
assert!(!CatalogStoreError::denied(IN_MEMORY, "refused").retryable());
}
}