use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use structfs_core_store::{
DetachedFuture, DetachedReader, DetachedWriter, Error, NoCodec, Path, Record, Value,
};
use crate::gate::CancelToken;
pub struct HandleCx {
pub id: u64,
pub cancel: CancelToken,
}
pub trait HandleProtocol: Send + Sync + 'static {
type Handle: Send + Sync + 'static;
fn open(&self, cx: HandleCx, request: Value) -> Result<Self::Handle, Error>;
fn read(&self, handle: Arc<Self::Handle>, sub: Path) -> DetachedFuture<Option<Record>>;
fn write(&self, handle: Arc<Self::Handle>, sub: Path, data: Record) -> DetachedFuture<Path>;
fn close(&self, handle: Arc<Self::Handle>) {
let _ = handle;
}
fn docs(&self) -> Option<Value> {
None
}
}
struct Entry<H> {
handle: Arc<H>,
cancel: CancelToken,
}
struct Inner<P: HandleProtocol> {
protocol: Arc<P>,
next_id: AtomicU64,
entries: Mutex<BTreeMap<u64, Entry<P::Handle>>>,
}
impl<P: HandleProtocol> Drop for Inner<P> {
fn drop(&mut self) {
let entries = std::mem::take(self.entries.get_mut().unwrap_or_else(|e| e.into_inner()));
for (_, entry) in entries {
entry.cancel.cancel();
self.protocol.close(entry.handle);
}
}
}
pub struct HandleStore<P: HandleProtocol> {
inner: Arc<Inner<P>>,
}
impl<P: HandleProtocol> Clone for HandleStore<P> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
const OUTSTANDING: &str = "outstanding";
impl<P: HandleProtocol> HandleStore<P> {
pub fn new(protocol: P) -> Self {
Self {
inner: Arc::new(Inner {
protocol: Arc::new(protocol),
next_id: AtomicU64::new(0),
entries: Mutex::new(BTreeMap::new()),
}),
}
}
fn lock_entries(&self) -> std::sync::MutexGuard<'_, BTreeMap<u64, Entry<P::Handle>>> {
self.inner.entries.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn handle_path(id: u64) -> Path {
Path::from_components(vec![OUTSTANDING.to_string(), id.to_string()])
}
pub fn live_handles(&self) -> usize {
self.lock_entries().len()
}
pub fn handle_ids(&self) -> Vec<u64> {
self.lock_entries().keys().copied().collect()
}
pub fn get_handle(&self, id: u64) -> Option<Arc<P::Handle>> {
self.get(id)
}
fn mint(&self, request: Value) -> Result<Path, Error> {
let id = self.inner.next_id.fetch_add(1, Ordering::SeqCst);
let cancel = CancelToken::new();
let handle = self.inner.protocol.open(
HandleCx {
id,
cancel: cancel.clone(),
},
request,
)?;
self.lock_entries().insert(
id,
Entry {
handle: Arc::new(handle),
cancel,
},
);
Ok(Self::handle_path(id))
}
fn release(&self, id: u64) {
let entry = self.lock_entries().remove(&id);
if let Some(entry) = entry {
entry.cancel.cancel();
self.inner.protocol.close(entry.handle);
}
}
fn get(&self, id: u64) -> Option<Arc<P::Handle>> {
self.lock_entries().get(&id).map(|e| e.handle.clone())
}
fn listing(&self) -> Value {
let items: Vec<Value> = self
.lock_entries()
.keys()
.map(|id| Value::String(Self::handle_path(*id).to_string()))
.collect();
let mut map = BTreeMap::new();
map.insert("items".to_string(), Value::Array(items));
Value::Map(map)
}
fn parse_handle(path: &Path) -> Option<(u64, Path)> {
if path.len() < 2 || &path[0] != OUTSTANDING {
return None;
}
let id: u64 = path[1].parse().ok()?;
Some((id, path.slice(2, path.len())))
}
}
impl<P: HandleProtocol> DetachedReader for HandleStore<P> {
fn read_detached(&mut self, from: &Path) -> DetachedFuture<Option<Record>> {
if from.is_empty() || (from.len() == 1 && &from[0] == OUTSTANDING) {
let listing = self.listing();
return Box::pin(async move { Ok(Some(Record::parsed(listing))) });
}
if from.len() == 1 && &from[0] == "docs" {
let docs = self.inner.protocol.docs();
return Box::pin(async move { Ok(docs.map(Record::parsed)) });
}
let Some((id, sub)) = Self::parse_handle(from) else {
return Box::pin(async move { Ok(None) });
};
let Some(handle) = self.get(id) else {
return Box::pin(async move { Ok(None) });
};
self.inner.protocol.read(handle, sub)
}
}
impl<P: HandleProtocol> DetachedWriter for HandleStore<P> {
fn write_detached(&mut self, to: &Path, data: Record) -> DetachedFuture<Path> {
if to.is_empty() {
let result = data.into_value(&NoCodec).and_then(|value| self.mint(value));
return Box::pin(async move { result });
}
let Some((id, sub)) = Self::parse_handle(to) else {
let path = to.clone();
return Box::pin(async move {
Err(Error::store(
"handle_store",
"write",
format!("no such path: {}", path),
))
});
};
if sub.is_empty() {
let result = match data.into_value(&NoCodec) {
Err(e) => Err(e),
Ok(value) if value.is_null() => {
self.release(id);
Ok(to.clone())
}
Ok(_) => Err(Error::conflict(format!(
"cannot overwrite outstanding handle {}; write Null to release it",
Self::handle_path(id)
))),
};
return Box::pin(async move { result });
}
let Some(handle) = self.get(id) else {
let path = to.clone();
return Box::pin(async move { Err(Error::not_found(path)) });
};
let fut = self.inner.protocol.write(handle, sub, data);
Box::pin(async move {
let rel = fut.await?;
Ok(Self::handle_path(id).join(&rel))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tail::TailLog;
struct StreamProtocol;
struct StreamHandle {
log: TailLog<Value>,
cancel: CancelToken,
}
impl HandleProtocol for StreamProtocol {
type Handle = StreamHandle;
fn open(&self, cx: HandleCx, _request: Value) -> Result<Self::Handle, Error> {
Ok(StreamHandle {
log: TailLog::new(),
cancel: cx.cancel,
})
}
fn read(&self, handle: Arc<Self::Handle>, sub: Path) -> DetachedFuture<Option<Record>> {
Box::pin(async move {
if sub.len() == 3 && &sub[0] == "events" && &sub[1] == "from" {
let seq: u64 = sub[2]
.parse()
.map_err(|_| Error::store("stream", "read", "bad cursor"))?;
let page = handle
.log
.read_from_cancellable(seq, &handle.cancel)
.await
.map_err(|c| c.into_error("stream handle released"))?;
return Ok(Some(Record::parsed(page.into_value())));
}
if sub.len() == 1 && &sub[0] == "status" {
let status = if handle.log.is_done() { "done" } else { "open" };
return Ok(Some(Record::parsed(Value::from(status))));
}
Ok(None)
})
}
fn write(
&self,
handle: Arc<Self::Handle>,
sub: Path,
data: Record,
) -> DetachedFuture<Path> {
Box::pin(async move {
if sub.len() == 1 && &sub[0] == "push" {
let value = data.into_value(&NoCodec)?;
handle.log.push(value);
return Ok(sub);
}
if sub.len() == 1 && &sub[0] == "done" {
handle.log.finish();
return Ok(sub);
}
Err(Error::store("stream", "write", "unknown sub-path"))
})
}
fn close(&self, handle: Arc<Self::Handle>) {
handle.log.finish();
}
}
fn store() -> HandleStore<StreamProtocol> {
HandleStore::new(StreamProtocol)
}
fn parsed(v: Value) -> Record {
Record::parsed(v)
}
#[tokio::test]
async fn mint_returns_handle_path() {
let mut s = store();
let path = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("req")))
.await
.unwrap();
assert_eq!(path.to_string(), "outstanding/0");
let second = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("req")))
.await
.unwrap();
assert_eq!(second.to_string(), "outstanding/1");
}
#[tokio::test]
async fn overwrite_is_conflict() {
let mut s = store();
let path = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap();
let err = s
.write_detached(&path, parsed(Value::from("clobber")))
.await
.unwrap_err();
assert!(matches!(err, Error::Conflict { .. }));
}
#[tokio::test]
async fn write_result_is_in_caller_namespace() {
let mut s = store();
let path = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap();
let result = s
.write_detached(
&path.join(&Path::parse("push").unwrap()),
parsed(Value::from(1i64)),
)
.await
.unwrap();
assert_eq!(result.to_string(), format!("{}/push", path));
}
#[tokio::test]
async fn tail_read_through_store() {
let mut s = store();
let handle = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap();
s.write_detached(
&handle.join(&Path::parse("push").unwrap()),
parsed(Value::from(1i64)),
)
.await
.unwrap();
s.write_detached(
&handle.join(&Path::parse("done").unwrap()),
parsed(Value::Null),
)
.await
.unwrap();
let record = s
.read_detached(&handle.join(&Path::parse("events/from/0").unwrap()))
.await
.unwrap()
.unwrap();
let map = match record.as_value().unwrap() {
Value::Map(m) => m.clone(),
_ => panic!("expected envelope"),
};
assert_eq!(map.get("status"), Some(&Value::from("done")));
assert!(matches!(map.get("items"), Some(Value::Array(a)) if a.len() == 1));
}
#[tokio::test]
async fn release_cancels_parked_reads() {
let mut s = store();
let handle = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap();
let mut reader = s.clone();
let tail_path = handle.join(&Path::parse("events/from/0").unwrap());
let parked = tokio::spawn(async move { reader.read_detached(&tail_path).await });
tokio::task::yield_now().await;
s.write_detached(&handle, parsed(Value::Null))
.await
.unwrap();
let err = parked.await.unwrap().unwrap_err();
assert!(err.is_cancelled());
assert!(s.read_detached(&handle).await.unwrap().is_none());
}
#[tokio::test]
async fn dropping_the_store_releases_live_handles() {
let s = store();
let handle_path = {
let mut s = s.clone();
s.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap()
};
let tail = handle_path.join(&Path::parse("events/from/0").unwrap());
let fut = {
let mut reader = s.clone();
reader.read_detached(&tail)
};
let parked = tokio::spawn(fut);
tokio::task::yield_now().await;
drop(s);
let err = tokio::time::timeout(std::time::Duration::from_secs(5), parked)
.await
.expect("parked read never resolved after store drop")
.unwrap()
.unwrap_err();
assert!(err.is_cancelled());
}
#[tokio::test]
async fn release_is_idempotent() {
let mut s = store();
let handle = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("r")))
.await
.unwrap();
s.write_detached(&handle, parsed(Value::Null))
.await
.unwrap();
s.write_detached(&handle, parsed(Value::Null))
.await
.unwrap();
s.write_detached(
&Path::parse("outstanding/999").unwrap(),
parsed(Value::Null),
)
.await
.unwrap();
}
#[tokio::test]
async fn listing_tracks_live_handles() {
let mut s = store();
let a = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("a")))
.await
.unwrap();
let _b = s
.write_detached(&Path::parse("").unwrap(), parsed(Value::from("b")))
.await
.unwrap();
let listing = s
.read_detached(&Path::parse("outstanding").unwrap())
.await
.unwrap()
.unwrap();
let items = match listing.as_value().unwrap() {
Value::Map(m) => match m.get("items").unwrap() {
Value::Array(a) => a.len(),
_ => panic!(),
},
_ => panic!(),
};
assert_eq!(items, 2);
s.write_detached(&a, parsed(Value::Null)).await.unwrap();
assert_eq!(s.live_handles(), 1);
}
#[tokio::test]
async fn unknown_paths_absent() {
let mut s = store();
assert!(s
.read_detached(&Path::parse("outstanding/42").unwrap())
.await
.unwrap()
.is_none());
assert!(s
.read_detached(&Path::parse("something/else").unwrap())
.await
.unwrap()
.is_none());
let err = s
.write_detached(
&Path::parse("outstanding/42/push").unwrap(),
parsed(Value::from(1i64)),
)
.await
.unwrap_err();
assert!(err.is_not_found());
}
#[test]
fn handle_path_component_is_valid() {
let p = HandleStore::<StreamProtocol>::handle_path(7);
assert_eq!(p.to_string(), "outstanding/7");
let _ = structfs_core_store::PathComponent::try_new("outstanding").unwrap();
}
}