use fnv::FnvHasher;
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use thiserror::Error;
use uuid::Uuid;
pub struct Exchange<I, O, M>
where
I: Send + Sync,
O: Send + Sync,
M: Send + Sync,
{
uuid: Uuid,
metadata: Option<M>,
input: Option<I>,
output: Option<O>,
input_listeners: Vec<Callback<I>>,
output_listeners: Vec<Callback<O>>,
attachments: Attachments,
}
impl<I, O, M> Exchange<I, O, M>
where
I: Send + Sync,
O: Send + Sync,
M: Send + Sync,
{
pub fn new() -> Self {
Self {
uuid: Uuid::new_v4(),
metadata: None,
input: None,
output: None,
input_listeners: Vec::new(),
output_listeners: Vec::new(),
attachments: Attachments::new(),
}
}
pub fn new_with_uuid(uuid: Uuid) -> Self {
Self {
uuid,
metadata: None,
input: None,
output: None,
input_listeners: Vec::new(),
output_listeners: Vec::new(),
attachments: Attachments::new(),
}
}
pub fn set_metadata(&mut self, metadata: M) {
self.metadata = Some(metadata);
}
pub fn metadata(&self) -> Result<&M, ExchangeError> {
match &self.metadata {
None => Err(ExchangeError::read_error(
&self.uuid,
"Metadata has not been set",
)),
Some(metadata) => Ok(metadata),
}
}
pub fn attachments(&self) -> &Attachments {
&self.attachments
}
pub fn attachments_mut(&mut self) -> &mut Attachments {
&mut self.attachments
}
pub fn add_input_listener(
&mut self,
callback: impl FnMut(&mut I, &mut Attachments) + Send + Sync + 'static,
) {
self.input_listeners.push(Callback::new(callback));
}
pub fn add_output_listener(
&mut self,
callback: impl FnMut(&mut O, &mut Attachments) + Send + Sync + 'static,
) {
self.output_listeners.push(Callback::new(callback));
}
pub fn uuid(&self) -> &Uuid {
&self.uuid
}
pub fn set_input(&mut self, input: I) {
self.input = Some(input);
}
pub async fn input(&self) -> Result<&I, ExchangeError> {
match &self.input {
Some(val) => Ok(val),
None => Err(ExchangeError::read_error(&self.uuid, "No input available")),
}
}
pub async fn input_mut(&mut self) -> Result<&mut I, ExchangeError> {
match &mut self.input {
Some(val) => Ok(val),
None => Err(ExchangeError::read_error(&self.uuid, "No input available"))
}
}
pub async fn take_input(&mut self) -> Result<I, ExchangeError> {
match self.input.take() {
Some(mut val) => {
for mut callback in &mut self.input_listeners.drain(..) {
callback.invoke(&mut val, &mut self.attachments);
}
Ok(val)
}
None => Err(ExchangeError::take_error(
&self.uuid,
"No input available to take",
)),
}
}
pub fn set_output(&mut self, output: O) {
self.output = Some(output);
}
pub async fn output(&self) -> Result<&O, ExchangeError> {
match &self.output {
Some(val) => Ok(val),
None => Err(ExchangeError::read_error(&self.uuid, "No output available")),
}
}
pub async fn output_mut(&mut self) -> Result<&mut O, ExchangeError> {
match &mut self.output {
Some(val) => Ok(val),
None => Err(ExchangeError::read_error(&self.uuid, "No output available")),
}
}
pub async fn take_output(&mut self) -> Result<O, ExchangeError> {
match self.output.take() {
Some(mut val) => {
for mut callback in &mut self.output_listeners.drain(..) {
callback.invoke(&mut val, &mut self.attachments);
}
Ok(val)
}
None => Err(ExchangeError::take_error(
&self.uuid,
"No output available to take",
)),
}
}
pub fn has_input(&self) -> bool {
self.input.is_some()
}
pub fn has_output(&self) -> bool {
self.output.is_some()
}
}
pub struct Attachments {
attachments: HashMap<AttachmentKey, Box<dyn Any + Send + Sync>, fnv::FnvBuildHasher>,
}
impl Attachments {
pub fn new() -> Self {
Self {
attachments: HashMap::with_hasher(fnv::FnvBuildHasher::default()),
}
}
pub fn add<K>(&mut self, key: impl AsRef<str>, value: K)
where
K: Send + Sync + 'static,
{
let type_id = TypeId::of::<K>();
self.attachments
.insert(AttachmentKey::new(key, type_id), Box::new(value));
}
pub fn get<K>(&self, key: impl AsRef<str>) -> Option<&K>
where
K: Send + 'static,
{
let type_id = TypeId::of::<K>();
if let Some(option_any) = self.attachments.get(&AttachmentKey::new(key, type_id)) {
option_any.downcast_ref::<K>()
} else {
None
}
}
pub fn get_mut<K>(&mut self, key: impl AsRef<str>) -> Option<&mut K>
where
K: Send + 'static,
{
let type_id = TypeId::of::<K>();
if let Some(option_any) = self.attachments.get_mut(&AttachmentKey::new(key, type_id)) {
option_any.downcast_mut::<K>()
} else {
None
}
}
}
#[derive(Error, Debug)]
pub enum ExchangeError {
#[error("Exchange '{uuid}' has already been completed")]
Completed { uuid: Uuid },
#[error("Read error occurred for exchange '{uuid}'. {message}")]
Read { uuid: Uuid, message: String },
#[error("Take error occurred for exchange '{uuid}'. {message}")]
Take { uuid: Uuid, message: String },
#[error("Callback error occurred for exchange '{uuid}'. {message}")]
Callback { uuid: Uuid, message: String },
}
impl ExchangeError {
#[inline]
pub const fn exchange_completed(uuid: &Uuid) -> Self {
ExchangeError::Completed { uuid: *uuid }
}
#[inline]
pub fn read_error(uuid: &Uuid, msg: impl Into<String>) -> Self {
ExchangeError::Read {
uuid: *uuid,
message: msg.into(),
}
}
#[inline]
pub(crate) fn take_error(uuid: &Uuid, msg: impl Into<String>) -> Self {
ExchangeError::Take {
uuid: *uuid,
message: msg.into(),
}
}
#[inline]
pub(crate) fn callback_error(uuid: &Uuid, msg: impl Into<String>) -> Self {
ExchangeError::Callback {
uuid: *uuid,
message: msg.into(),
}
}
}
#[derive(PartialOrd, PartialEq, Hash, Eq)]
pub struct AttachmentKey {
key_hash: u64,
type_hash: u64,
}
impl AttachmentKey {
pub fn new(key: impl AsRef<str>, type_id: TypeId) -> Self {
let key = key.as_ref();
let mut key_hasher = FnvHasher::default();
key.hash(&mut key_hasher);
let key_hash = key_hasher.finish();
let mut type_hasher = FnvHasher::default();
type_id.hash(&mut type_hasher);
let type_hash = type_hasher.finish();
Self {
key_hash,
type_hash,
}
}
}
pub struct Callback<T> {
callback: Box<dyn FnMut(&mut T, &mut Attachments) + Send + Sync>,
}
impl<T> Callback<T>
where
T: Send,
{
pub fn new(callback: impl FnMut(&mut T, &mut Attachments) + Send + Sync + 'static) -> Self {
Self {
callback: Box::new(callback),
}
}
pub fn invoke(&mut self, write: &mut T, attachments: &mut Attachments) {
(self.callback)(write, attachments);
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::exchange::Attachments;
use hyper::body::Bytes;
struct TestStruct;
#[test]
fn test_attachments() {
let mut attachments = Attachments::new();
let key1 = "test_key1";
let key2 = "test_key2";
let key3 = "test_key3";
let key4 = "test_key4";
{
attachments.add::<u64>(key1, 1);
attachments.add::<String>(key2, String::from("test"));
attachments.add::<bool>(key3, true);
let test_struct = TestStruct;
attachments.add::<TestStruct>(key4, test_struct);
}
{
assert!(attachments.get::<u64>(key1).is_some());
assert!(attachments.get::<String>(key2).is_some());
assert!(attachments.get::<bool>(key3).is_some());
assert!(attachments.get::<TestStruct>(key4).is_some());
}
}
#[tokio::test]
async fn test_unified_exchange_buffered() {
let mut exchange: Exchange<Bytes, Bytes, ()> = Exchange::new();
exchange.set_input(Bytes::from("test input"));
assert!(exchange.has_input());
exchange.set_output(Bytes::from("test output"));
assert!(exchange.has_output());
let output = exchange.take_output().await.unwrap();
assert_eq!(output, Bytes::from("test output"));
}
#[tokio::test]
async fn test_unified_exchange_callbacks() {
let mut exchange: Exchange<Bytes, Bytes, ()> = Exchange::new();
exchange.add_input_listener(|input: &mut Bytes, _| {
let mut new_data = b"prefix:".to_vec();
new_data.extend_from_slice(input);
*input = Bytes::from(new_data);
});
exchange.set_input(Bytes::from("test"));
let input = exchange.take_input().await.unwrap();
assert_eq!(input, Bytes::from("prefix:test"));
}
#[tokio::test]
async fn test_unified_exchange_metadata() {
let mut exchange: Exchange<(), (), String> = Exchange::new();
exchange.set_metadata("test metadata".to_string());
let metadata = exchange.metadata().unwrap();
assert_eq!(metadata, "test metadata");
}
}