use crate::Error;
use commonware_formatting::hex;
#[cfg(test)]
use commonware_utils::channel::oneshot::Sender as OneshotSender;
use commonware_utils::{channel::watch, sync::Mutex};
#[cfg(not(target_os = "linux"))]
use std::collections::HashSet;
#[cfg(test)]
use std::sync::atomic::{AtomicU64, Ordering};
#[cfg(test)]
use std::sync::mpsc::{Receiver as MpscReceiver, Sender as MpscSender};
use std::{
collections::HashMap,
ptr,
sync::{Arc, Weak},
};
#[derive(Default)]
pub(crate) struct Pending {
pub(super) entries: Mutex<HashMap<(String, Vec<u8>), Entry>>,
#[cfg(not(target_os = "linux"))]
pub(super) flushed: Mutex<HashSet<(String, Vec<u8>)>>,
#[cfg(test)]
pub(super) test: Hooks,
}
#[cfg(test)]
#[derive(Default)]
pub(super) struct Hooks {
pub(super) completions: AtomicU64,
pub(super) before_complete: Mutex<Option<(OneshotSender<()>, MpscReceiver<()>)>>,
pub(super) fail_flush: Mutex<Option<Error>>,
pub(super) after_dispatch: Mutex<Option<(OneshotSender<()>, MpscReceiver<()>)>>,
pub(super) fail_creation_after: Mutex<Option<usize>>,
pub(super) after_identity_observation: Mutex<Option<(MpscSender<usize>, MpscReceiver<()>)>>,
pub(super) before_metadata: Mutex<Option<(OneshotSender<()>, MpscReceiver<()>)>>,
pub(super) before_admit: Mutex<Option<(OneshotSender<()>, MpscReceiver<()>)>>,
}
#[derive(Default)]
pub(super) struct Entry {
pub(super) identity: Weak<Generation>,
pub(super) settle: Option<Receiver>,
pub(super) dirty: bool,
pub(super) failed: Option<Error>,
}
impl Entry {
const fn is_clear(&self) -> bool {
self.settle.is_none() && !self.dirty && self.failed.is_none()
}
}
pub(crate) struct Generation {
pub(super) pending: Arc<Pending>,
pub(super) key: (String, Vec<u8>),
}
impl Generation {
pub(crate) fn release(&self) -> Option<Sender> {
self.pending.start(self)
}
}
impl Drop for Generation {
fn drop(&mut self) {
let mut entries = self.pending.entries.lock();
if entries
.get(&self.key)
.is_some_and(|entry| ptr::eq(entry.identity.as_ptr(), self) && entry.is_clear())
{
entries.remove(&self.key);
}
}
}
type Receiver = watch::Receiver<bool>;
pub(crate) type Sender = watch::Sender<bool>;
impl Pending {
pub(crate) fn admit(
self: &Arc<Self>,
partition: &str,
name: &[u8],
existing: bool,
) -> Result<(Arc<Generation>, Option<Receiver>, bool), Error> {
if !existing {
self.forget(partition, Some(name));
}
#[cfg(test)]
if let Some((entered, released)) = self.test.before_admit.lock().take() {
let _ = entered.send(());
let _ = released.recv();
}
let key = (partition.to_owned(), name.to_vec());
let mut entries = self.entries.lock();
let entry = entries.entry(key.clone()).or_default();
let live = entry.identity.strong_count() != 0;
#[cfg(test)]
self.observe_identity(&entry.identity);
if live {
return Err(Error::BlobAlreadyOpen(partition.to_owned(), hex(name)));
}
if let Some(failed) = &entry.failed {
return Err(failed.clone());
}
let generation = Arc::new(Generation {
pending: self.clone(),
key,
});
entry.identity = Arc::downgrade(&generation);
#[cfg(not(target_os = "linux"))]
if self.flushed.lock().insert(generation.key.clone()) && existing {
entry.dirty = true;
}
let owed = entry.settle.is_some() || entry.dirty;
Ok((generation, entry.settle.clone(), owed))
}
fn start(&self, generation: &Generation) -> Option<Sender> {
let mut entries = self.entries.lock();
let entry = entries.get_mut(&generation.key)?;
if !ptr::eq(entry.identity.as_ptr(), generation) || entry.settle.is_some() {
return None;
}
let (sender, receiver) = watch::channel(false);
entry.settle = Some(receiver);
Some(sender)
}
pub(crate) fn settle(
&self,
key: &(String, Vec<u8>),
sender: Sender,
dirty: bool,
failure: Option<Error>,
) {
{
let mut entries = self.entries.lock();
if let Some(entry) = entries.get_mut(key)
&& entry
.settle
.as_ref()
.is_some_and(|receiver| receiver.same_channel(&sender.subscribe()))
{
entry.settle = None;
entry.dirty |= dirty;
if failure.is_some() {
entry.failed = failure;
}
let live = entry.identity.strong_count() != 0;
#[cfg(test)]
self.observe_identity(&entry.identity);
if !live && entry.is_clear() {
entries.remove(key);
}
}
}
let _ = sender.send(true);
}
pub(crate) async fn wait(receiver: Option<Receiver>) -> Result<(), Error> {
let Some(mut receiver) = receiver else {
return Ok(());
};
receiver
.wait_for(|settled| *settled)
.await
.map(|_| ())
.map_err(|_| Error::Closed)
}
pub(crate) fn debt(
&self,
key: &(String, Vec<u8>),
identity: &Weak<Generation>,
) -> Result<bool, Error> {
let entries = self.entries.lock();
let Some(entry) = entries.get(key) else {
return Ok(false);
};
if !Weak::ptr_eq(&entry.identity, identity) {
return Ok(false);
}
if let Some(failed) = &entry.failed {
return Err(failed.clone());
}
Ok(entry.dirty)
}
pub(crate) fn clear(
&self,
key: &(String, Vec<u8>),
identity: &Weak<Generation>,
result: &Result<(), Error>,
) {
let mut entries = self.entries.lock();
let Some(entry) = entries.get_mut(key) else {
return;
};
if !Weak::ptr_eq(&entry.identity, identity) {
return;
}
match result {
Ok(()) => entry.dirty = false,
Err(error) => entry.failed = Some(error.clone()),
}
}
pub(crate) fn fail(&self, generation: &Generation, error: Error) {
let mut entries = self.entries.lock();
if let Some(entry) = entries.get_mut(&generation.key)
&& ptr::eq(entry.identity.as_ptr(), generation)
{
entry.failed = Some(error);
}
}
pub(crate) fn forget(&self, partition: &str, name: Option<&[u8]>) {
if let Some(name) = name {
let key = (partition.to_owned(), name.to_vec());
self.entries.lock().remove(&key);
#[cfg(not(target_os = "linux"))]
self.flushed.lock().remove(&key);
} else {
self.entries
.lock()
.retain(|(stored, _), _| stored != partition);
#[cfg(not(target_os = "linux"))]
self.flushed
.lock()
.retain(|(stored, _)| stored != partition);
}
}
}
#[cfg(test)]
impl Pending {
fn observe_identity(&self, identity: &Weak<Generation>) {
let hook = self.test.after_identity_observation.lock().take();
if let Some((entered, release)) = hook {
entered.send(identity.strong_count()).unwrap();
let _ = release.recv();
}
}
pub(crate) fn take_flush_failure(&self) -> Option<Error> {
self.test.fail_flush.lock().take()
}
pub(crate) fn before_complete(&self) {
let hook = self.test.before_complete.lock().take();
if let Some((entered, released)) = hook {
let _ = entered.send(());
let _ = released.recv();
}
}
pub(crate) fn completed(&self) {
self.test.completions.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn outstanding(&self) -> usize {
self.entries
.lock()
.values()
.filter(|entry| entry.settle.is_some())
.count()
}
pub(crate) fn owes(&self, partition: &str, name: &[u8]) -> bool {
self.entries
.lock()
.get(&(partition.to_owned(), name.to_vec()))
.is_some_and(|entry| entry.dirty || entry.failed.is_some())
}
pub(crate) fn completions(&self) -> u64 {
self.test.completions.load(Ordering::Acquire)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::{sync::mpsc, thread};
fn key() -> (String, Vec<u8>) {
("a".to_owned(), b"1".to_vec())
}
fn check_last_generation_drop_during_observation(settle: bool) {
let pending = Arc::new(Pending::default());
let (generation, _, _) = pending.admit("a", b"1", false).unwrap();
let identity = Arc::downgrade(&generation);
let sender = settle.then(|| generation.release().unwrap());
let outcome = sender.as_ref().map(|sender| sender.subscribe());
let (entered, entering) = mpsc::channel();
let (release, released) = mpsc::channel();
*pending.test.after_identity_observation.lock() = Some((entered, released));
let observer = {
let pending = pending.clone();
thread::spawn(move || {
if let Some(sender) = sender {
pending.settle(&key(), sender, false, None);
} else {
assert!(matches!(
pending.admit("a", b"1", true),
Err(Error::BlobAlreadyOpen(partition, name))
if partition == "a" && name == "31"
));
}
})
};
let owners = entering.recv().unwrap();
let dropper = thread::spawn(move || drop(generation));
while identity.strong_count() == owners {
thread::yield_now();
}
release.send(()).unwrap();
let independent = {
let pending = pending.clone();
thread::spawn(move || {
drop(pending.admit("independent", b"2", false).unwrap());
})
};
observer.join().unwrap();
dropper.join().unwrap();
independent.join().unwrap();
if let Some(outcome) = outcome {
assert!(*outcome.borrow());
}
assert!(pending.entries.lock().is_empty());
}
#[test]
fn test_last_generation_drop_during_settle() {
check_last_generation_drop_during_observation(true);
}
#[test]
fn test_last_generation_drop_during_admit() {
check_last_generation_drop_during_observation(false);
}
#[tokio::test]
async fn test_wait_observes_settling_predecessor() {
let pending = Arc::new(Pending::default());
let (generation, wait, owed) = pending.admit("a", b"1", false).unwrap();
assert!(!owed);
Pending::wait(wait).await.unwrap();
let sender = generation.release().unwrap();
drop(generation);
let (generation, wait, owed) = pending.admit("a", b"1", true).unwrap();
assert!(owed);
let waiter = tokio::spawn(Pending::wait(wait));
tokio::task::yield_now().await;
assert!(!waiter.is_finished());
pending.settle(&key(), sender, false, None);
waiter.await.unwrap().unwrap();
assert_eq!(pending.outstanding(), 0);
assert!(!pending.debt(&key(), &Arc::downgrade(&generation)).unwrap());
drop(generation);
assert!(pending.entries.lock().is_empty());
}
#[test]
fn test_debt_is_established_by_the_next_open() {
let pending = Arc::new(Pending::default());
let (first, _, _) = pending.admit("a", b"1", false).unwrap();
let sender = first.release().unwrap();
drop(first);
pending.settle(&key(), sender, true, None);
assert!(pending.owes("a", b"1"));
let (second, wait, owed) = pending.admit("a", b"1", true).unwrap();
assert!(wait.is_none());
assert!(owed);
assert!(pending.debt(&key(), &Arc::downgrade(&second)).unwrap());
let sender = second.release().unwrap();
drop(second);
pending.settle(&key(), sender, false, None);
let (third, _, owed) = pending.admit("a", b"1", true).unwrap();
assert!(owed);
assert!(pending.debt(&key(), &Arc::downgrade(&third)).unwrap());
pending.clear(&key(), &Arc::downgrade(&third), &Ok(()));
assert!(!pending.debt(&key(), &Arc::downgrade(&third)).unwrap());
assert!(!pending.owes("a", b"1"));
drop(third);
assert!(pending.entries.lock().is_empty());
}
#[test]
fn test_failures_stay_until_forgotten() {
let pending = Arc::new(Pending::default());
let (generation, _, _) = pending.admit("a", b"1", false).unwrap();
let sender = generation.release().unwrap();
pending.settle(&key(), sender, false, Some(Error::Closed));
drop(generation);
for _ in 0..2 {
assert!(matches!(pending.admit("a", b"1", true), Err(Error::Closed)));
}
pending.forget("a", Some(b"1"));
let (generation, _, _) = pending.admit("a", b"1", false).unwrap();
pending.clear(
&key(),
&Arc::downgrade(&generation),
&Err(Error::ReadFailed),
);
assert!(matches!(
pending.debt(&key(), &Arc::downgrade(&generation)),
Err(Error::ReadFailed)
));
pending.clear(&key(), &Arc::downgrade(&generation), &Ok(()));
assert!(matches!(
pending.debt(&key(), &Arc::downgrade(&generation)),
Err(Error::ReadFailed)
));
drop(generation);
assert!(matches!(
pending.admit("a", b"1", true),
Err(Error::ReadFailed)
));
pending.forget("a", Some(b"1"));
let (generation, _, _) = pending.admit("a", b"1", false).unwrap();
pending.fail(&generation, Error::WriteFailed);
drop(generation);
assert!(matches!(
pending.admit("a", b"1", true),
Err(Error::WriteFailed)
));
pending.forget("a", Some(b"1"));
drop(pending.admit("a", b"1", false).unwrap());
assert!(pending.entries.lock().is_empty());
}
#[test]
fn test_stale_settlement_leaves_a_recreated_name_alone() {
let pending = Arc::new(Pending::default());
let (old, _, _) = pending.admit("a", b"1", false).unwrap();
let stale = old.release().unwrap();
pending.forget("a", Some(b"1"));
let (current, _, owed) = pending.admit("a", b"1", false).unwrap();
assert!(!owed);
pending.settle(&key(), stale, true, Some(Error::Closed));
assert!(old.release().is_none());
pending.clear(&key(), &Arc::downgrade(&old), &Err(Error::Closed));
assert!(!pending.owes("a", b"1"));
assert!(!pending.debt(&key(), &Arc::downgrade(¤t)).unwrap());
drop(old);
let sender = current.release().unwrap();
drop(current);
pending.settle(&key(), sender, false, None);
assert!(pending.entries.lock().is_empty());
}
#[test]
fn test_live_open_refuses_a_second_admit() {
let pending = Arc::new(Pending::default());
let (first, _, _) = pending.admit("a", b"1", false).unwrap();
assert!(matches!(
pending.admit("a", b"1", true),
Err(Error::BlobAlreadyOpen(partition, name)) if partition == "a" && name == "31"
));
drop(first);
drop(pending.admit("a", b"1", true).unwrap());
}
#[test]
fn test_generations_retire_and_release_clean_entries() {
let pending = Arc::new(Pending::default());
let (first, _, _) = pending.admit("a", b"1", false).unwrap();
let sender = first.release().unwrap();
drop(first);
assert_eq!(pending.entries.lock().len(), 1);
pending.settle(&key(), sender, false, None);
assert!(pending.entries.lock().is_empty());
for name in 0..128u64 {
drop(pending.admit("clean", &name.to_be_bytes(), false).unwrap());
assert!(pending.entries.lock().is_empty());
}
}
}