use std::{path::PathBuf, sync::Arc};
use crate::object::{
Action, ActionId, Blob, ContentHash, State, StateAttachment, StateAttachmentId, StateId, Tree,
};
pub mod actor_presence;
pub mod agent_task;
pub mod codec;
pub mod fs;
pub mod liveness;
#[cfg(any(test, feature = "memory-backend"))]
pub mod memory;
pub mod pack;
pub mod shallow;
pub mod source;
pub mod store_compliance;
pub mod writer_lease;
pub use actor_presence::{
ActorChainNode, ActorPresence, ActorPresenceStatus, ActorPresenceStore, AgentUsageSummary,
ContextQueryEntry, generate_actor_session_id,
};
pub use agent_task::{
AGENT_TASK_SCHEMA_VERSION, AgentTaskRecord, AgentTaskStatus, AgentTaskStore,
generate_agent_task_id, validate_task_id,
};
pub use fs::{
DEFAULT_PACK_INSTALL_INTENT_TTL_SECS, FsStore, PackInstallIntent, PackInstallMetricsSnapshot,
PackInstallPhase, PackInstallRecoverReport, install_pack_bytes_journaled,
pack_install_metrics_reset, pack_install_metrics_snapshot, recover_pack_install_intents,
recover_pack_install_intents_with_ttl,
};
pub use heddle_format::compression::{CompressionConfig, CompressionError, compress, decompress};
pub use liveness::{
AGENT_LEASE_DURATION, Liveness, current_boot_id, process_alive, reservation_liveness_at,
};
#[cfg(any(test, feature = "memory-backend"))]
pub use memory::InMemoryStore;
pub use pack::{PackBuilder, PackObjectId, PackReader, PackStats, StreamingPackBuilder, SyncData};
pub use shallow::ShallowInfo;
#[cfg(feature = "async-source")]
pub use source::AsyncObjectSource;
pub use source::ObjectSource;
pub use writer_lease::{
WriterLease, WriterLeaseAuthOutcome, WriterLeaseDraft, WriterLeaseGrant,
WriterLeaseReserveOutcome, WriterLeaseStatus, WriterLeaseStore, generate_writer_lease_id,
generate_writer_lease_token,
};
pub trait ExternalObjectSource: Send + Sync {
fn get_blob(&self, hash: &ContentHash) -> Result<Option<Blob>>;
fn get_tree(&self, hash: &ContentHash) -> Result<Option<Tree>>;
fn get_state(&self, id: &StateId) -> Result<Option<State>>;
fn list_states(&self) -> Result<Vec<StateId>>;
}
pub use crate::error::{HeddleError as StoreError, HeddleError, Result};
impl From<CompressionError> for HeddleError {
fn from(e: CompressionError) -> Self {
HeddleError::Compression(e.to_string())
}
}
#[derive(Clone)]
pub enum AnyStore {
Fs(FsStore),
}
macro_rules! any_store_dispatch {
($self:ident, $method:ident ( $($arg:expr),* )) => {
match $self {
AnyStore::Fs(inner) => inner.$method($($arg),*),
}
};
}
impl ObjectStore for AnyStore {
fn get_blob(&self, hash: &ContentHash) -> Result<Option<Blob>> {
match self {
AnyStore::Fs(inner) => ObjectStore::get_blob(inner, hash),
}
}
fn put_blob(&self, blob: &Blob) -> Result<ContentHash> {
any_store_dispatch!(self, put_blob(blob))
}
fn get_blob_bytes(&self, hash: &ContentHash) -> Result<Option<bytes::Bytes>> {
match self {
AnyStore::Fs(inner) => ObjectStore::get_blob_bytes(inner, hash),
}
}
fn blob_size(&self, hash: &ContentHash) -> Result<Option<u64>> {
any_store_dispatch!(self, blob_size(hash))
}
fn loose_blob_path(&self, hash: &ContentHash) -> Option<PathBuf> {
any_store_dispatch!(self, loose_blob_path(hash))
}
fn promote_to_loose_uncompressed(&self, hash: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, promote_to_loose_uncompressed(hash))
}
fn clear_recent_caches(&self) {
any_store_dispatch!(self, clear_recent_caches())
}
fn put_blob_with_hash(&self, blob: &Blob, hash: ContentHash) -> Result<ContentHash> {
any_store_dispatch!(self, put_blob_with_hash(blob, hash))
}
fn has_blob(&self, hash: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, has_blob(hash))
}
fn has_blob_locally(&self, hash: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, has_blob_locally(hash))
}
fn get_tree(&self, hash: &ContentHash) -> Result<Option<Tree>> {
match self {
AnyStore::Fs(inner) => ObjectStore::get_tree(inner, hash),
}
}
fn get_tree_serialized(&self, hash: &ContentHash) -> Result<Option<Vec<u8>>> {
match self {
AnyStore::Fs(inner) => ObjectStore::get_tree_serialized(inner, hash),
}
}
fn put_tree(&self, tree: &Tree) -> Result<ContentHash> {
any_store_dispatch!(self, put_tree(tree))
}
fn has_tree(&self, hash: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, has_tree(hash))
}
fn has_tree_locally(&self, hash: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, has_tree_locally(hash))
}
fn get_state(&self, id: &StateId) -> Result<Option<State>> {
match self {
AnyStore::Fs(inner) => ObjectStore::get_state(inner, id),
}
}
fn put_state(&self, state: &State) -> Result<()> {
any_store_dispatch!(self, put_state(state))
}
fn has_state(&self, id: &StateId) -> Result<bool> {
any_store_dispatch!(self, has_state(id))
}
fn list_states(&self) -> Result<Vec<StateId>> {
any_store_dispatch!(self, list_states())
}
fn get_state_attachment(
&self,
state: &StateId,
id: &StateAttachmentId,
) -> Result<Option<StateAttachment>> {
any_store_dispatch!(self, get_state_attachment(state, id))
}
fn put_state_attachment(&self, attachment: &StateAttachment) -> Result<StateAttachmentId> {
any_store_dispatch!(self, put_state_attachment(attachment))
}
fn list_state_attachments(&self, state: &StateId) -> Result<Vec<StateAttachment>> {
any_store_dispatch!(self, list_state_attachments(state))
}
fn get_action(&self, id: &ActionId) -> Result<Option<Action>> {
any_store_dispatch!(self, get_action(id))
}
fn put_action(&self, action: &mut Action) -> Result<ActionId> {
any_store_dispatch!(self, put_action(action))
}
fn list_actions(&self) -> Result<Vec<ActionId>> {
any_store_dispatch!(self, list_actions())
}
fn list_blobs(&self) -> Result<Vec<ContentHash>> {
any_store_dispatch!(self, list_blobs())
}
fn list_trees(&self) -> Result<Vec<ContentHash>> {
any_store_dispatch!(self, list_trees())
}
fn put_blob_bytes_with_hash(&self, data: &[u8], hash: ContentHash) -> Result<ContentHash> {
any_store_dispatch!(self, put_blob_bytes_with_hash(data, hash))
}
fn put_tree_serialized(&self, data: &[u8], hash: ContentHash) -> Result<ContentHash> {
match self {
AnyStore::Fs(inner) => ObjectStore::put_tree_serialized(inner, data, hash),
}
}
fn put_state_serialized(&self, data: &[u8], id: StateId) -> Result<()> {
any_store_dispatch!(self, put_state_serialized(data, id))
}
fn put_action_serialized(&self, data: &[u8], id: ActionId) -> Result<()> {
any_store_dispatch!(self, put_action_serialized(data, id))
}
fn get_pack_object(
&self,
id: &pack::PackObjectId,
) -> Result<Option<(pack::ObjectType, Vec<u8>)>> {
any_store_dispatch!(self, get_pack_object(id))
}
fn put_blobs_packed(&self, blobs: Vec<(ContentHash, Vec<u8>)>) -> Result<()> {
any_store_dispatch!(self, put_blobs_packed(blobs))
}
fn install_pack(&self, pack_data: &[u8], index_data: &[u8]) -> Result<Vec<pack::PackObjectId>> {
any_store_dispatch!(self, install_pack(pack_data, index_data))
}
fn install_pack_streaming(
&self,
pack_path: &std::path::Path,
index_path: &std::path::Path,
) -> Result<Vec<pack::PackObjectId>> {
any_store_dispatch!(self, install_pack_streaming(pack_path, index_path))
}
fn pack_objects(&self, aggressive: bool) -> Result<(u64, u64)> {
any_store_dispatch!(self, pack_objects(aggressive))
}
fn prune_loose_objects(&self) -> Result<(u64, u64)> {
any_store_dispatch!(self, prune_loose_objects())
}
fn begin_snapshot_write_batch(&self) -> Result<()> {
any_store_dispatch!(self, begin_snapshot_write_batch())
}
fn flush_snapshot_write_batch(&self) -> Result<()> {
any_store_dispatch!(self, flush_snapshot_write_batch())
}
fn abort_snapshot_write_batch(&self) {
any_store_dispatch!(self, abort_snapshot_write_batch())
}
fn has_redactions_for_blob(&self, blob: &ContentHash) -> Result<bool> {
any_store_dispatch!(self, has_redactions_for_blob(blob))
}
fn get_redactions_bytes_for_blob(&self, blob: &ContentHash) -> Result<Option<Vec<u8>>> {
any_store_dispatch!(self, get_redactions_bytes_for_blob(blob))
}
fn put_redactions_bytes_for_blob(&self, blob: &ContentHash, bytes: &[u8]) -> Result<()> {
any_store_dispatch!(self, put_redactions_bytes_for_blob(blob, bytes))
}
fn list_blobs_with_redactions(&self) -> Result<Vec<ContentHash>> {
any_store_dispatch!(self, list_blobs_with_redactions())
}
fn has_state_visibility_for_state(&self, state: &StateId) -> Result<bool> {
any_store_dispatch!(self, has_state_visibility_for_state(state))
}
fn get_state_visibility_bytes_for_state(&self, state: &StateId) -> Result<Option<Vec<u8>>> {
any_store_dispatch!(self, get_state_visibility_bytes_for_state(state))
}
fn put_state_visibility_bytes_for_state(&self, state: &StateId, bytes: &[u8]) -> Result<()> {
any_store_dispatch!(self, put_state_visibility_bytes_for_state(state, bytes))
}
fn list_states_with_visibility(&self) -> Result<Vec<StateId>> {
any_store_dispatch!(self, list_states_with_visibility())
}
}
impl AnyStore {
pub fn set_external_source(&mut self, source: Arc<dyn ExternalObjectSource>) {
match self {
Self::Fs(store) => store.set_external_source(source),
}
}
}
pub trait ObjectStore: Send + Sync {
fn get_blob(&self, hash: &ContentHash) -> Result<Option<Blob>>;
fn put_blob(&self, blob: &Blob) -> Result<ContentHash>;
fn get_blob_bytes(&self, hash: &ContentHash) -> Result<Option<bytes::Bytes>> {
Ok(self
.get_blob(hash)?
.map(|blob| bytes::Bytes::from(blob.into_content())))
}
fn blob_size(&self, hash: &ContentHash) -> Result<Option<u64>> {
Ok(self.get_blob(hash)?.map(|blob| blob.content().len() as u64))
}
fn loose_blob_path(&self, _hash: &ContentHash) -> Option<PathBuf> {
None
}
fn promote_to_loose_uncompressed(&self, _hash: &ContentHash) -> Result<bool> {
Ok(false)
}
fn clear_recent_caches(&self) {}
fn put_blob_with_hash(&self, blob: &Blob, hash: ContentHash) -> Result<ContentHash> {
if blob.hash() != hash {
return Err(HeddleError::InvalidObject("blob hash mismatch".to_string()));
}
self.put_blob(blob)
}
fn has_blob(&self, hash: &ContentHash) -> Result<bool>;
fn has_blob_locally(&self, hash: &ContentHash) -> Result<bool> {
self.has_blob(hash)
}
fn get_tree(&self, hash: &ContentHash) -> Result<Option<Tree>>;
fn put_tree(&self, tree: &Tree) -> Result<ContentHash>;
fn has_tree(&self, hash: &ContentHash) -> Result<bool>;
fn has_tree_locally(&self, hash: &ContentHash) -> Result<bool> {
self.has_tree(hash)
}
fn get_state(&self, id: &StateId) -> Result<Option<State>>;
fn put_state(&self, state: &State) -> Result<()>;
fn has_state(&self, id: &StateId) -> Result<bool>;
fn list_states(&self) -> Result<Vec<StateId>>;
fn get_state_attachment(
&self,
_state: &StateId,
_id: &StateAttachmentId,
) -> Result<Option<StateAttachment>> {
Ok(None)
}
fn put_state_attachment(&self, _attachment: &StateAttachment) -> Result<StateAttachmentId> {
Err(HeddleError::InvalidObject(
"object store does not support state attachments".to_string(),
))
}
fn list_state_attachments(&self, _state: &StateId) -> Result<Vec<StateAttachment>> {
Ok(Vec::new())
}
fn get_action(&self, id: &ActionId) -> Result<Option<Action>>;
fn put_action(&self, action: &mut Action) -> Result<ActionId>;
fn list_actions(&self) -> Result<Vec<ActionId>>;
fn list_blobs(&self) -> Result<Vec<ContentHash>>;
fn list_trees(&self) -> Result<Vec<ContentHash>>;
fn put_blob_bytes_with_hash(&self, data: &[u8], hash: ContentHash) -> Result<ContentHash> {
self.put_blob_with_hash(&Blob::from_slice(data), hash)
}
fn get_tree_serialized(&self, hash: &ContentHash) -> Result<Option<Vec<u8>>> {
Ok(self
.get_tree(hash)?
.map(|tree| rmp_serde::to_vec(&tree))
.transpose()?)
}
fn put_tree_serialized(&self, data: &[u8], hash: ContentHash) -> Result<ContentHash> {
let tree: Tree = rmp_serde::from_slice(data)?;
tree.validate()?;
if tree.hash() != hash {
return Err(HeddleError::Corruption {
expected: hash,
found: tree.hash(),
});
}
self.put_tree(&tree)
}
fn put_state_serialized(&self, data: &[u8], id: StateId) -> Result<()> {
let state: State = rmp_serde::from_slice(data)?;
let found = state.id();
if found != id {
return Err(HeddleError::InvalidObject(format!(
"state id mismatch: expected {id}, computed {found}"
)));
}
self.put_state(&state)
}
fn put_action_serialized(&self, data: &[u8], id: ActionId) -> Result<()> {
let mut action: Action = rmp_serde::from_slice(data)?;
let found_id = action.compute_id();
if found_id != id {
return Err(HeddleError::InvalidObject(format!(
"action id mismatch: expected {}, found {}",
id, found_id
)));
}
let stored_id = self.put_action(&mut action)?;
if stored_id != id {
return Err(HeddleError::InvalidObject(format!(
"action id mismatch after write: expected {}, found {}",
id, stored_id
)));
}
Ok(())
}
fn get_pack_object(
&self,
id: &pack::PackObjectId,
) -> Result<Option<(pack::ObjectType, Vec<u8>)>> {
match id {
pack::PackObjectId::Hash(hash) => {
if let Some(blob) = self.get_blob(hash)? {
return Ok(Some((pack::ObjectType::Blob, blob.content().to_vec())));
}
if let Some(tree) = self.get_tree(hash)? {
return Ok(Some((
pack::ObjectType::Tree,
rmp_serde::to_vec_named(&tree)?,
)));
}
if let Some(action) = self.get_action(&ActionId::from_hash(*hash))? {
return Ok(Some((
pack::ObjectType::Action,
rmp_serde::to_vec_named(&action)?,
)));
}
Ok(None)
}
pack::PackObjectId::StateId(change_id) => {
if let Some(state) = self.get_state(change_id)? {
Ok(Some((
pack::ObjectType::State,
rmp_serde::to_vec_named(&state)?,
)))
} else {
Ok(None)
}
}
}
}
fn put_blobs_packed(&self, blobs: Vec<(ContentHash, Vec<u8>)>) -> Result<()> {
for (hash, data) in blobs {
if !self.has_blob(&hash)? {
self.put_blob_bytes_with_hash(&data, hash)?;
}
}
Ok(())
}
fn install_pack(&self, pack_data: &[u8], index_data: &[u8]) -> Result<Vec<pack::PackObjectId>> {
let reader = pack::PackReader::from_slice(pack_data, index_data)?;
let ids = reader.list_ids();
for id in &ids {
let Some((obj_type, data)) = reader.get_object(id)? else {
continue;
};
match (id, obj_type) {
(pack::PackObjectId::Hash(hash), pack::ObjectType::Blob) => {
self.put_blob_bytes_with_hash(&data, *hash)?;
}
(pack::PackObjectId::Hash(hash), pack::ObjectType::Tree) => {
self.put_tree_serialized(&data, *hash)?;
}
(pack::PackObjectId::Hash(hash), pack::ObjectType::Action) => {
self.put_action_serialized(&data, ActionId::from_hash(*hash))?;
}
(pack::PackObjectId::StateId(change_id), pack::ObjectType::State) => {
self.put_state_serialized(&data, *change_id)?;
}
_ => {
return Err(HeddleError::InvalidObject(format!(
"unsupported native pack object: {:?} {:?}",
id, obj_type
)));
}
}
}
Ok(ids)
}
fn install_pack_streaming(
&self,
pack_path: &std::path::Path,
index_path: &std::path::Path,
) -> Result<Vec<pack::PackObjectId>> {
let pack_data = std::fs::read(pack_path).map_err(StoreError::from)?;
let index_data = std::fs::read(index_path).map_err(StoreError::from)?;
let ids = self.install_pack(&pack_data, &index_data)?;
let _ = std::fs::remove_file(pack_path);
let _ = std::fs::remove_file(index_path);
Ok(ids)
}
fn pack_objects(&self, aggressive: bool) -> Result<(u64, u64)> {
let _ = aggressive;
Ok((0, 0))
}
fn prune_loose_objects(&self) -> Result<(u64, u64)> {
Ok((0, 0))
}
fn begin_snapshot_write_batch(&self) -> Result<()> {
Ok(())
}
fn flush_snapshot_write_batch(&self) -> Result<()> {
Ok(())
}
fn abort_snapshot_write_batch(&self) {}
fn has_redactions_for_blob(&self, _blob: &ContentHash) -> Result<bool> {
Ok(false)
}
fn get_redactions_bytes_for_blob(&self, _blob: &ContentHash) -> Result<Option<Vec<u8>>> {
Ok(None)
}
fn put_redactions_bytes_for_blob(&self, _blob: &ContentHash, _bytes: &[u8]) -> Result<()> {
Err(HeddleError::InvalidObject(
"this object store does not support persisting redactions".to_string(),
))
}
fn list_blobs_with_redactions(&self) -> Result<Vec<ContentHash>> {
Ok(Vec::new())
}
fn has_state_visibility_for_state(&self, _state: &StateId) -> Result<bool> {
Ok(false)
}
fn get_state_visibility_bytes_for_state(&self, _state: &StateId) -> Result<Option<Vec<u8>>> {
Ok(None)
}
fn put_state_visibility_bytes_for_state(&self, _state: &StateId, _bytes: &[u8]) -> Result<()> {
Err(HeddleError::InvalidObject(
"this object store does not support persisting state visibility".to_string(),
))
}
fn list_states_with_visibility(&self) -> Result<Vec<StateId>> {
Ok(Vec::new())
}
}
#[cfg(test)]
mod any_store_tests {
use tempfile::TempDir;
use super::*;
use crate::object::{Attribution, Operation, Principal};
fn fs_any_store() -> (TempDir, AnyStore) {
let temp = TempDir::new().unwrap();
let store = FsStore::new(temp.path().join(".heddle"));
store.init().unwrap();
(temp, AnyStore::Fs(store))
}
#[test]
fn fs_variant_dispatches_every_object_store_method() {
let (_temp, store) = fs_any_store();
let blob = Blob::from("any-store dispatch blob");
let blob_hash = store.put_blob(&blob).unwrap();
assert_eq!(
ObjectStore::get_blob(&store, &blob_hash)
.unwrap()
.unwrap()
.content(),
blob.content()
);
assert!(store.has_blob(&blob_hash).unwrap());
assert_eq!(
ObjectStore::get_blob_bytes(&store, &blob_hash)
.unwrap()
.unwrap()
.as_ref(),
blob.content()
);
assert_eq!(
store.blob_size(&blob_hash).unwrap().unwrap(),
blob.content().len() as u64
);
assert!(store.loose_blob_path(&blob_hash).is_some());
store.promote_to_loose_uncompressed(&blob_hash).unwrap();
assert!(store.list_blobs().unwrap().contains(&blob_hash));
let bytes_blob = Blob::from("put-with-hash blob");
let bytes_hash = bytes_blob.hash();
assert_eq!(
store.put_blob_with_hash(&bytes_blob, bytes_hash).unwrap(),
bytes_hash
);
let raw_blob = Blob::from("raw bytes blob");
let raw_hash = raw_blob.hash();
assert_eq!(
store
.put_blob_bytes_with_hash(raw_blob.content(), raw_hash)
.unwrap(),
raw_hash
);
let tree = Tree::new();
let tree_hash = store.put_tree(&tree).unwrap();
assert!(ObjectStore::get_tree(&store, &tree_hash).unwrap().is_some());
assert!(store.has_tree(&tree_hash).unwrap());
assert!(store.list_trees().unwrap().contains(&tree_hash));
let tree2 = Tree::new();
let tree2_bytes = rmp_serde::to_vec_named(&tree2).unwrap();
assert_eq!(
store
.put_tree_serialized(&tree2_bytes, tree2.hash())
.unwrap(),
tree2.hash()
);
let attribution =
Attribution::human(Principal::new("AnyStore Test", "anystore@example.com"));
let state = State::new(tree_hash, vec![], attribution.clone());
let state_id = state.id();
store.put_state(&state).unwrap();
assert!(ObjectStore::get_state(&store, &state_id).unwrap().is_some());
assert!(store.has_state(&state_id).unwrap());
assert!(store.list_states().unwrap().contains(&state_id));
let state2 = State::new(tree2.hash(), vec![], attribution.clone());
let state2_bytes = rmp_serde::to_vec_named(&state2).unwrap();
store
.put_state_serialized(&state2_bytes, state2.id())
.unwrap();
let mut action = Action::new(
None,
StateId::from_bytes([3; 32]),
Operation::Snapshot,
"any-store action",
attribution,
);
let action_id = store.put_action(&mut action).unwrap();
assert!(store.get_action(&action_id).unwrap().is_some());
assert!(store.list_actions().unwrap().contains(&action_id));
let action_bytes = rmp_serde::to_vec_named(&action).unwrap();
store
.put_action_serialized(&action_bytes, action_id)
.unwrap();
let packed = Blob::from("packed-via-any-store");
let packed_hash = packed.hash();
store
.put_blobs_packed(vec![(packed_hash, packed.into_content())])
.unwrap();
assert!(
store
.get_pack_object(&pack::PackObjectId::Hash(packed_hash))
.unwrap()
.is_some()
);
store.pack_objects(false).unwrap();
store.prune_loose_objects().unwrap();
let _ = store.install_pack(&[], &[]);
let _ = store.install_pack_streaming(
std::path::Path::new("/nonexistent/pack"),
std::path::Path::new("/nonexistent/idx"),
);
store.begin_snapshot_write_batch().unwrap();
store.flush_snapshot_write_batch().unwrap();
store.begin_snapshot_write_batch().unwrap();
store.abort_snapshot_write_batch();
let redaction = b"any-store redaction bytes";
store
.put_redactions_bytes_for_blob(&blob_hash, redaction)
.unwrap();
assert!(store.has_redactions_for_blob(&blob_hash).unwrap());
assert_eq!(
store
.get_redactions_bytes_for_blob(&blob_hash)
.unwrap()
.as_deref(),
Some(redaction.as_slice())
);
assert!(
store
.list_blobs_with_redactions()
.unwrap()
.contains(&blob_hash)
);
let state_visibility = b"any-store state visibility bytes";
store
.put_state_visibility_bytes_for_state(&state_id, state_visibility)
.unwrap();
assert!(store.has_state_visibility_for_state(&state_id).unwrap());
assert_eq!(
store
.get_state_visibility_bytes_for_state(&state_id)
.unwrap()
.as_deref(),
Some(state_visibility.as_slice())
);
assert!(
store
.list_states_with_visibility()
.unwrap()
.contains(&state_id)
);
store.clear_recent_caches();
}
}