use crate::journal::Error;
use commonware_formatting::hex;
use commonware_runtime::{
buffer::{
paged::{CacheRef, Writer},
Write,
},
telemetry::metrics::{Counter, Gauge, GaugeExt, MetricsExt as _},
Blob, BufferPool, Error as RError, Handle, Metrics, Storage,
};
use futures::future::{join_all, try_join_all};
use std::{
collections::{BTreeMap, BTreeSet},
future::Future,
mem::take,
num::NonZeroUsize,
};
use tracing::debug;
pub trait SectionBuffer: Send + Sync {
fn size(&self) -> u64;
fn sync(&mut self) -> impl Future<Output = Result<(), RError>> + Send;
fn start_sync(&mut self) -> impl Future<Output = Handle<()>> + Send;
fn wait_for_sync(&mut self) -> impl Future<Output = Result<(), RError>> + Send;
fn resize(&mut self, len: u64) -> impl Future<Output = Result<(), RError>> + Send;
}
impl<B: Blob> SectionBuffer for Writer<B> {
fn size(&self) -> u64 {
Self::size(self)
}
async fn sync(&mut self) -> Result<(), RError> {
Self::sync(self).await
}
async fn start_sync(&mut self) -> Handle<()> {
Self::start_sync(self).await
}
async fn wait_for_sync(&mut self) -> Result<(), RError> {
Self::wait_for_sync(self).await
}
async fn resize(&mut self, len: u64) -> Result<(), RError> {
Self::resize(self, len).await
}
}
impl<B: Blob> SectionBuffer for Write<B> {
fn size(&self) -> u64 {
Self::size(self)
}
async fn sync(&mut self) -> Result<(), RError> {
Self::sync(self).await
}
async fn start_sync(&mut self) -> Handle<()> {
Self::start_sync(self).await
}
async fn wait_for_sync(&mut self) -> Result<(), RError> {
Self::wait_for_sync(self).await
}
async fn resize(&mut self, len: u64) -> Result<(), RError> {
Self::resize(self, len).await
}
}
pub trait BufferFactory<B: Blob>: Clone + Send + Sync {
type Buffer: SectionBuffer;
fn create(
&self,
blob: B,
size: u64,
) -> impl Future<Output = Result<Self::Buffer, RError>> + Send;
}
#[derive(Clone)]
pub struct AppendFactory {
pub write_buffer: NonZeroUsize,
pub page_cache_ref: CacheRef,
}
impl<B: Blob> BufferFactory<B> for AppendFactory {
type Buffer = Writer<B>;
async fn create(&self, blob: B, size: u64) -> Result<Self::Buffer, RError> {
Writer::new(
blob,
size,
self.write_buffer.get(),
self.page_cache_ref.clone(),
)
.await
}
}
#[derive(Clone)]
pub struct WriteFactory {
pub capacity: NonZeroUsize,
pub pool: BufferPool,
}
impl<B: Blob> BufferFactory<B> for WriteFactory {
type Buffer = Write<B>;
async fn create(&self, blob: B, size: u64) -> Result<Self::Buffer, RError> {
Ok(Write::new(blob, size, self.capacity, self.pool.clone()))
}
}
#[derive(Clone)]
pub struct Config<F> {
pub partition: String,
pub factory: F,
}
pub struct Manager<E: Storage + Metrics, F: BufferFactory<E::Blob>> {
context: E,
partition: String,
factory: F,
pub(crate) blobs: BTreeMap<u64, F::Buffer>,
oldest_retained_section: u64,
tracked: Gauge,
synced: Counter,
pruned: Counter,
}
impl<E: Storage + Metrics, F: BufferFactory<E::Blob>> Manager<E, F> {
async fn wait_for_syncs<'a>(
blobs: impl IntoIterator<Item = &'a mut F::Buffer>,
) -> Result<(), Error>
where
F::Buffer: 'a,
{
try_join_all(blobs.into_iter().map(|blob| blob.wait_for_sync()))
.await
.map(|_| ())
.map_err(Error::Runtime)
}
pub async fn init(context: E, cfg: Config<F>) -> Result<Self, Error> {
let mut blobs = BTreeMap::new();
let stored_blobs = match context.scan(&cfg.partition).await {
Ok(blobs) => blobs,
Err(RError::PartitionMissing(_)) => Vec::new(),
Err(err) => return Err(Error::Runtime(err)),
};
for name in stored_blobs {
let (blob, size) = context.open(&cfg.partition, &name).await?;
let hex_name = hex(&name);
let section = match name.try_into() {
Ok(section) => u64::from_be_bytes(section),
Err(_) => return Err(Error::InvalidBlobName(hex_name)),
};
debug!(section, blob = hex_name, size, "loaded section");
let buffer = cfg.factory.create(blob, size).await?;
blobs.insert(section, buffer);
}
let tracked = context.gauge("tracked", "Number of blobs");
let synced = context.counter("synced", "Number of syncs");
let pruned = context.counter("pruned", "Number of blobs pruned");
let _ = tracked.try_set(blobs.len());
Ok(Self {
context,
partition: cfg.partition,
factory: cfg.factory,
blobs,
oldest_retained_section: 0,
tracked,
synced,
pruned,
})
}
pub const fn prune_guard(&self, section: u64) -> Result<(), Error> {
if section < self.oldest_retained_section {
Err(Error::AlreadyPrunedToSection(self.oldest_retained_section))
} else {
Ok(())
}
}
pub fn get(&self, section: u64) -> Result<Option<&F::Buffer>, Error> {
self.prune_guard(section)?;
Ok(self.blobs.get(§ion))
}
pub async fn get_or_create(&mut self, section: u64) -> Result<&mut F::Buffer, Error> {
self.prune_guard(section)?;
if !self.blobs.contains_key(§ion) {
let name = section.to_be_bytes();
let (blob, size) = self.context.open(&self.partition, &name).await?;
let buffer = self.factory.create(blob, size).await?;
self.tracked.inc();
self.blobs.insert(section, buffer);
}
Ok(self.blobs.get_mut(§ion).unwrap())
}
pub async fn sync(&mut self, sections: impl crate::Sections) -> Result<(), Error> {
let sections = sections.sections().collect::<BTreeSet<_>>();
for §ion in §ions {
self.prune_guard(section)?;
}
let futures: Vec<_> = self
.blobs
.iter_mut()
.filter(|(section, _)| sections.contains(section))
.map(|(_, blob)| blob.sync())
.collect();
let count = futures.len() as u64;
try_join_all(futures).await.map_err(Error::Runtime)?;
self.synced.inc_by(count);
Ok(())
}
pub async fn start_sync(
&mut self,
sections: impl crate::Sections,
) -> Result<Handle<()>, Error> {
let sections = sections.sections().collect::<BTreeSet<_>>();
for §ion in §ions {
self.prune_guard(section)?;
}
let futures: Vec<_> = self
.blobs
.iter_mut()
.filter(|(section, _)| sections.contains(section))
.map(|(_, blob)| blob.start_sync())
.collect();
self.synced.inc_by(futures.len() as u64);
let handles = join_all(futures).await;
Ok(Handle::from_future(async move {
try_join_all(handles).await.map(|_| ())
}))
}
pub async fn sync_all(&mut self) -> Result<(), Error> {
let count = self.blobs.len() as u64;
try_join_all(self.blobs.values_mut().map(|b| b.sync()))
.await
.map_err(Error::Runtime)?;
self.synced.inc_by(count);
Ok(())
}
pub async fn prune(&mut self, min: u64) -> Result<bool, Error> {
let mut pruned = false;
while let Some((§ion, _)) = self.blobs.first_key_value() {
if section >= min {
break;
}
let mut blob = self.blobs.remove(§ion).unwrap();
blob.wait_for_sync().await?;
let size = blob.size();
drop(blob);
self.context
.remove(&self.partition, Some(§ion.to_be_bytes()))
.await?;
pruned = true;
debug!(section, size, "pruned blob");
self.tracked.dec();
self.pruned.inc();
}
if pruned {
self.oldest_retained_section = min;
}
Ok(pruned)
}
pub fn oldest_section(&self) -> Option<u64> {
self.blobs.first_key_value().map(|(&s, _)| s)
}
pub fn newest_section(&self) -> Option<u64> {
self.blobs.last_key_value().map(|(&s, _)| s)
}
pub fn is_empty(&self) -> bool {
self.blobs.is_empty()
}
pub fn num_sections(&self) -> usize {
self.blobs.len()
}
pub fn sections_from(
&mut self,
start_section: u64,
) -> impl Iterator<Item = (&u64, &mut F::Buffer)> {
self.blobs.range_mut(start_section..)
}
pub fn sections(&self) -> impl Iterator<Item = u64> + '_ {
self.blobs.keys().copied()
}
pub async fn remove_section(&mut self, section: u64) -> Result<bool, Error> {
self.prune_guard(section)?;
if let Some(mut blob) = self.blobs.remove(§ion) {
blob.wait_for_sync().await?;
let size = blob.size();
drop(blob);
self.context
.remove(&self.partition, Some(§ion.to_be_bytes()))
.await?;
self.tracked.dec();
debug!(section, size, "removed section");
Ok(true)
} else {
Ok(false)
}
}
pub async fn destroy(mut self) -> Result<(), Error> {
Self::wait_for_syncs(self.blobs.values_mut()).await?;
for (section, blob) in self.blobs.into_iter() {
let size = blob.size();
drop(blob);
debug!(section, size, "destroyed blob");
self.context
.remove(&self.partition, Some(§ion.to_be_bytes()))
.await?;
}
match self.context.remove(&self.partition, None).await {
Ok(()) => {}
Err(RError::PartitionMissing(_)) => {}
Err(err) => return Err(Error::Runtime(err)),
}
Ok(())
}
pub async fn clear(&mut self) -> Result<(), Error> {
Self::wait_for_syncs(self.blobs.values_mut()).await?;
let blobs = take(&mut self.blobs);
for (section, blob) in blobs {
let size = blob.size();
drop(blob);
debug!(section, size, "cleared blob");
self.context
.remove(&self.partition, Some(§ion.to_be_bytes()))
.await?;
}
let _ = self.tracked.try_set(0);
self.oldest_retained_section = 0;
Ok(())
}
pub async fn rewind(&mut self, section: u64, size: u64) -> Result<(), Error> {
self.prune_guard(section)?;
let sections_to_remove: Vec<u64> = match section.checked_add(1) {
Some(next) => self.blobs.range(next..).rev().map(|(&s, _)| s).collect(),
None => Vec::new(),
};
for s in sections_to_remove {
let mut blob = self.blobs.remove(&s).unwrap();
blob.wait_for_sync().await?;
drop(blob);
self.context
.remove(&self.partition, Some(&s.to_be_bytes()))
.await?;
self.tracked.dec();
debug!(section = s, "removed blob during rewind");
}
if let Some(blob) = self.blobs.get_mut(§ion) {
let current_size = blob.size();
if size < current_size {
blob.resize(size).await?;
debug!(
section,
old_size = current_size,
new_size = size,
"rewound blob"
);
}
}
Ok(())
}
pub async fn rewind_section(&mut self, section: u64, size: u64) -> Result<(), Error> {
self.prune_guard(section)?;
if let Some(blob) = self.blobs.get_mut(§ion) {
let current = blob.size();
if size < current {
blob.resize(size).await?;
debug!(section, from = current, to = size, "rewound section");
}
}
Ok(())
}
pub fn size(&self, section: u64) -> Result<u64, Error> {
self.prune_guard(section)?;
Ok(self.blobs.get(§ion).map_or(0, |blob| blob.size()))
}
}
#[cfg(test)]
mod tests {
use super::*;
use commonware_runtime::{deterministic, Runner as _, Spawner as _, Supervisor as _};
use commonware_utils::{channel::oneshot, sync::Mutex};
use futures::{
future::{BoxFuture, Shared},
FutureExt as _,
};
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
type SyncSender = oneshot::Sender<Result<(), RError>>;
type PendingSyncs = Arc<Mutex<Vec<SyncSender>>>;
type SharedSync = Shared<BoxFuture<'static, Result<(), RError>>>;
#[derive(Clone)]
struct TestFactory {
pending: PendingSyncs,
wait_for_syncs: Arc<AtomicUsize>,
}
struct TestBuffer {
pending: PendingSyncs,
wait_for_syncs: Arc<AtomicUsize>,
syncing: Option<SharedSync>,
}
impl SectionBuffer for TestBuffer {
fn size(&self) -> u64 {
0
}
async fn sync(&mut self) -> Result<(), RError> {
Ok(())
}
async fn start_sync(&mut self) -> Handle<()> {
if let Some(syncing) = &self.syncing {
return Handle::from_future(syncing.clone());
}
let (sender, receiver) = oneshot::channel();
self.pending.lock().push(sender);
let sync = async move {
receiver.await.map_err(|_| RError::Closed)??;
Ok(())
}
.boxed()
.shared();
self.syncing = Some(sync.clone());
Handle::from_future(sync)
}
async fn wait_for_sync(&mut self) -> Result<(), RError> {
if let Some(syncing) = self.syncing.take() {
self.wait_for_syncs.fetch_add(1, Ordering::Relaxed);
syncing.await?;
}
Ok(())
}
async fn resize(&mut self, _len: u64) -> Result<(), RError> {
Ok(())
}
}
impl<B: Blob> BufferFactory<B> for TestFactory {
type Buffer = TestBuffer;
async fn create(&self, _blob: B, _size: u64) -> Result<Self::Buffer, RError> {
Ok(TestBuffer {
pending: self.pending.clone(),
wait_for_syncs: self.wait_for_syncs.clone(),
syncing: None,
})
}
}
fn test_config(pending: PendingSyncs, wait_for_syncs: Arc<AtomicUsize>) -> Config<TestFactory> {
Config {
partition: "test".into(),
factory: TestFactory {
pending,
wait_for_syncs,
},
}
}
fn release_pending_syncs(pending: &PendingSyncs) {
for sender in std::mem::take(&mut *pending.lock()) {
let _ = sender.send(Ok(()));
}
}
fn complete_next_pending_sync(pending: &PendingSyncs, result: Result<(), RError>) {
let sender = {
let mut pending = pending.lock();
assert!(!pending.is_empty(), "no pending sync to complete");
pending.remove(0)
};
let _ = sender.send(result);
}
#[test]
fn test_start_sync_multiple_sections_returns_combined_handle() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let pending = Arc::new(Mutex::new(Vec::new()));
let wait_for_syncs = Arc::new(AtomicUsize::new(0));
let cfg = test_config(pending.clone(), wait_for_syncs);
let mut manager = Manager::init(context.child("manager"), cfg)
.await
.expect("failed to initialize manager");
manager
.get_or_create(1)
.await
.expect("failed to create first section");
manager
.get_or_create(2)
.await
.expect("failed to create second section");
let handle = manager
.start_sync([1, 2])
.await
.expect("failed to start sync");
assert_eq!(pending.lock().len(), 2);
futures::pin_mut!(handle);
complete_next_pending_sync(&pending, Ok(()));
assert!(
futures::poll!(handle.as_mut()).is_pending(),
"combined sync handle must wait for every selected section"
);
complete_next_pending_sync(&pending, Ok(()));
handle.await.expect("sync handle should complete");
manager.destroy().await.expect("destroy failed");
});
}
#[test]
fn test_start_sync_reuses_in_flight_section_handle_without_waiting() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let pending = Arc::new(Mutex::new(Vec::new()));
let wait_for_syncs = Arc::new(AtomicUsize::new(0));
let cfg = test_config(pending.clone(), wait_for_syncs);
let mut manager = Manager::init(context.child("manager"), cfg)
.await
.expect("failed to initialize manager");
manager
.get_or_create(1)
.await
.expect("failed to create section");
let first = manager.start_sync(1).await.expect("failed to start sync");
assert_eq!(pending.lock().len(), 1);
let second = manager
.start_sync(1)
.await
.expect("failed to observe in-flight sync");
assert_eq!(
pending.lock().len(),
1,
"repeated start_sync should observe the in-flight section sync"
);
futures::pin_mut!(second);
assert!(
futures::poll!(second.as_mut()).is_pending(),
"reused start_sync handle must wait for the in-flight sync"
);
release_pending_syncs(&pending);
first.await.expect("first sync handle should complete");
second.await.expect("reused sync handle should complete");
manager.destroy().await.expect("destroy failed");
});
}
#[test]
fn test_prune_waits_for_in_flight_start_sync() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let pending = Arc::new(Mutex::new(Vec::new()));
let wait_for_syncs = Arc::new(AtomicUsize::new(0));
let cfg = test_config(pending.clone(), wait_for_syncs.clone());
let mut manager = Manager::init(context.child("manager"), cfg)
.await
.expect("failed to initialize manager");
manager
.get_or_create(1)
.await
.expect("failed to create section");
let handle = manager.start_sync(1).await.expect("failed to start sync");
assert_eq!(pending.lock().len(), 1);
let completed = Arc::new(AtomicUsize::new(0));
let completed_clone = completed.clone();
let waiter = context.child("prune").spawn(|_| async move {
assert!(manager.prune(2).await.expect("prune failed"));
completed_clone.fetch_add(1, Ordering::Relaxed);
manager
});
while wait_for_syncs.load(Ordering::Relaxed) == 0 {
commonware_runtime::reschedule().await;
}
commonware_runtime::reschedule().await;
assert_eq!(
completed.load(Ordering::Relaxed),
0,
"prune must wait for the in-flight start_sync handle"
);
release_pending_syncs(&pending);
handle.await.expect("sync handle should complete");
while completed.load(Ordering::Relaxed) == 0 {
commonware_runtime::reschedule().await;
}
let manager = waiter.await.expect("prune task failed");
assert!(manager.is_empty());
});
}
#[test]
fn test_destroy_waits_for_in_flight_start_sync() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let pending = Arc::new(Mutex::new(Vec::new()));
let wait_for_syncs = Arc::new(AtomicUsize::new(0));
let cfg = test_config(pending.clone(), wait_for_syncs.clone());
let mut manager = Manager::init(context.child("manager"), cfg)
.await
.expect("failed to initialize manager");
manager
.get_or_create(1)
.await
.expect("failed to create section");
let handle = manager.start_sync(1).await.expect("failed to start sync");
assert_eq!(pending.lock().len(), 1);
let completed = Arc::new(AtomicUsize::new(0));
let completed_clone = completed.clone();
let waiter = context.child("destroy").spawn(|_| async move {
manager.destroy().await.expect("destroy failed");
completed_clone.fetch_add(1, Ordering::Relaxed);
});
while wait_for_syncs.load(Ordering::Relaxed) == 0 {
commonware_runtime::reschedule().await;
}
commonware_runtime::reschedule().await;
assert_eq!(
completed.load(Ordering::Relaxed),
0,
"destroy must wait for the in-flight start_sync handle"
);
release_pending_syncs(&pending);
handle.await.expect("sync handle should complete");
while completed.load(Ordering::Relaxed) == 0 {
commonware_runtime::reschedule().await;
}
waiter.await.expect("destroy task failed");
});
}
#[test]
fn test_destroy_surfaces_failed_in_flight_start_sync() {
let executor = deterministic::Runner::default();
executor.start(|context| async move {
let pending = Arc::new(Mutex::new(Vec::new()));
let wait_for_syncs = Arc::new(AtomicUsize::new(0));
let cfg = test_config(pending.clone(), wait_for_syncs);
let mut manager = Manager::init(context.child("manager"), cfg)
.await
.expect("failed to initialize manager");
manager
.get_or_create(1)
.await
.expect("failed to create section");
let handle = manager.start_sync(1).await.expect("failed to start sync");
complete_next_pending_sync(&pending, Err(RError::Closed));
let err = manager
.destroy()
.await
.expect_err("destroy should surface the sync failure");
assert!(matches!(err, Error::Runtime(RError::Closed)));
assert!(matches!(
handle.await.expect_err("sync handle should fail"),
RError::Closed
));
});
}
}