#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
use std::collections::HashMap;
use std::fs;
use std::future::{Future, ready};
use std::mem;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, LazyLock, Weak};
#[cfg(feature = "testing")]
use ruststream::RawMessage;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
#[cfg(feature = "testing")]
use ruststream::testing::{Coordinator, InProcess, TestableBroker};
use ruststream::{
AddressedCopies, Broker, ConnectedBroker, DefaultPublish, DescribeServer, Lend,
OutgoingMessage, PairError, PublishPolicy, Publisher, ServerSpec, Subscribe,
};
use sea_streamer_file::{
AutoStreamReset, FileConnectOptions, FileConsumerOptions, FileErr, FileId, FileProducer,
FileProducerOptions, FileStreamer, MessageSource, StreamMode,
};
use sea_streamer_types::{
ConsumerMode, ConsumerOptions as _, Producer as _, StreamErr, StreamKey, Streamer as _,
};
use tokio::runtime::Handle;
use tokio::sync::{Mutex, OnceCell};
use crate::error::{SeaFileError, box_err};
#[cfg(feature = "testing")]
use crate::in_process::MemoryFile;
use crate::stream::FileStream;
#[cfg(feature = "asyncapi")]
use crate::stream::channel_extension;
use crate::subscriber::{FileSubscriber, Reader};
use crate::wire;
pub(crate) struct Core {
pub(crate) link: Link,
pub(crate) path: String,
pub(crate) closed: AtomicBool,
}
pub(crate) struct Disk {
pub(crate) streamer: FileStreamer,
pub(crate) producer: Arc<Mutex<Appender>>,
pub(crate) runtime: Handle,
}
pub(crate) enum Appender {
Open(FileProducer),
Failed(FileProducer, FileErr),
Closed,
}
static APPENDERS: LazyLock<Mutex<HashMap<String, Weak<Mutex<Appender>>>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
async fn open_appender(
appenders: &HashMap<String, Weak<Mutex<Appender>>>,
path: &str,
) -> Option<Arc<Mutex<Appender>>> {
let shared = appenders.get(path)?.upgrade()?;
let open = !matches!(*shared.lock().await, Appender::Closed);
open.then_some(shared)
}
pub(crate) enum Link {
File(Disk),
#[cfg(feature = "testing")]
InProcess(Arc<MemoryFile>),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Link>() == size_of::<Disk>());
impl Core {
pub(crate) fn ensure_open(&self) -> Result<(), SeaFileError> {
if self.closed.load(Ordering::Acquire) {
return Err(SeaFileError::NotConnected);
}
Ok(())
}
}
impl std::fmt::Debug for Core {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Core")
.field("path", &self.path)
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
pub(crate) type CoreCell = Arc<OnceCell<Arc<Core>>>;
const PRIME_STREAM: &str = "ruststream-internal";
#[derive(Debug, Clone)]
#[must_use]
pub struct FileBroker {
path: String,
create: bool,
end_with_eos: bool,
beacon_interval: Option<u32>,
#[cfg(test)]
size_limit: Option<u64>,
cell: CoreCell,
}
impl FileBroker {
pub fn new(path: impl Into<String>) -> Self {
Self {
path: path.into(),
create: true,
end_with_eos: false,
beacon_interval: None,
#[cfg(test)]
size_limit: None,
cell: Arc::new(OnceCell::new()),
}
}
#[cfg(test)]
const fn size_limit(mut self, bytes: u64) -> Self {
self.size_limit = Some(bytes);
self
}
pub fn existing_only(mut self) -> Self {
self.create = false;
self
}
pub fn end_with_eos(mut self) -> Self {
self.end_with_eos = true;
self
}
pub fn beacon_interval(mut self, bytes: u32) -> Self {
self.beacon_interval = Some(bytes);
self
}
#[must_use]
pub fn publisher(&self) -> FilePublisher {
FilePublisher {
cell: Arc::clone(&self.cell),
}
}
}
impl Broker for FileBroker {
type Error = SeaFileError;
type Connected = ConnectedFileBroker;
async fn connect(self) -> Result<Self::Connected, Self::Error> {
let core = self
.cell
.get_or_try_init(async || {
let mut options = FileConnectOptions::default();
if self.create {
options.set_create_if_not_exists(true);
}
if self.end_with_eos {
options.set_end_with_eos(true);
}
#[cfg(test)]
if let Some(bytes) = self.size_limit {
options.set_file_size_limit(bytes);
}
if let Some(interval) = self.beacon_interval {
options.set_beacon_interval(interval).map_err(|e| {
SeaFileError::Invalid(format!("invalid beacon interval: {e}"))
})?;
}
let file_id = FileId::new(self.path.clone());
let uri = file_id
.to_streamer_uri()
.map_err(|e| SeaFileError::Connect {
target: self.path.clone(),
source: box_err(e),
})?;
let streamer = FileStreamer::connect(uri, options).await.map_err(|e| {
SeaFileError::Connect {
target: self.path.clone(),
source: box_err(e),
}
})?;
let connect_err = |e: StreamErr<FileErr>| SeaFileError::Connect {
target: self.path.clone(),
source: box_err(e),
};
let mut appenders = APPENDERS.lock().await;
if let Some(shared) = open_appender(&appenders, &self.path).await {
return Ok(Arc::new(Core {
link: Link::File(Disk {
streamer,
producer: shared,
runtime: Handle::current(),
}),
path: self.path.clone(),
closed: AtomicBool::new(false),
}));
}
let producer = streamer
.create_generic_producer(FileProducerOptions::default())
.await
.map_err(connect_err)?;
let fresh = fs::metadata(&self.path).map_or(true, |meta| meta.len() <= 128);
if fresh {
let prime_key = StreamKey::new(PRIME_STREAM)
.map_err(|e| SeaFileError::Invalid(e.to_string()))?;
producer
.send_to(&prime_key, b"1".as_slice())
.map_err(connect_err)?
.await
.map_err(connect_err)?;
let mut flusher = producer.clone();
flusher.flush().await.map_err(connect_err)?;
}
let shared = Arc::new(Mutex::new(Appender::Open(producer)));
appenders.retain(|_, appender| appender.strong_count() > 0);
appenders.insert(self.path.clone(), Arc::downgrade(&shared));
drop(appenders);
Ok::<_, SeaFileError>(Arc::new(Core {
link: Link::File(Disk {
streamer,
producer: shared,
runtime: Handle::current(),
}),
path: self.path.clone(),
closed: AtomicBool::new(false),
}))
})
.await?
.clone();
Ok(ConnectedFileBroker {
core,
cell: self.cell,
})
}
}
#[cfg(feature = "testing")]
impl InProcess for FileBroker {
async fn connect_in_process(self) -> Result<Self::Connected, Self::Error> {
let core = self
.cell
.get_or_try_init(async || {
if let Some(interval) = self.beacon_interval {
FileConnectOptions::default()
.set_beacon_interval(interval)
.map_err(|e| {
SeaFileError::Invalid(format!("invalid beacon interval: {e}"))
})?;
}
FileId::new(self.path.clone())
.to_streamer_uri()
.map_err(|e| SeaFileError::Connect {
target: self.path.clone(),
source: box_err(e),
})?;
Ok::<_, SeaFileError>(Arc::new(Core {
link: Link::InProcess(MemoryFile::new(self.end_with_eos)),
path: self.path.clone(),
closed: AtomicBool::new(false),
}))
})
.await?
.clone();
Ok(ConnectedFileBroker {
core,
cell: self.cell,
})
}
}
#[cfg(feature = "testing")]
ruststream::register_testable_broker!(FileBroker);
impl DescribeServer for FileBroker {
fn describe_server(&self) -> ServerSpec {
ServerSpec::in_process("file").description(self.path.clone())
}
}
#[derive(Debug)]
pub struct ConnectedFileBroker {
pub(crate) core: Arc<Core>,
cell: CoreCell,
}
impl ConnectedFileBroker {
#[must_use]
pub fn publisher(&self) -> FilePublisher {
FilePublisher {
cell: Arc::clone(&self.cell),
}
}
pub async fn subscribe_stream(
&self,
descriptor: FileStream,
) -> Result<FileSubscriber, SeaFileError> {
descriptor.validate()?;
self.core.ensure_open()?;
let key = StreamKey::new(descriptor.stream())
.map_err(|e| SeaFileError::Invalid(format!("'{}': {e}", descriptor.stream())))?;
let disk = match &self.core.link {
Link::File(disk) => disk,
#[cfg(feature = "testing")]
Link::InProcess(file) => {
return Ok(FileSubscriber::in_process(
file.subscribe(key.name(), descriptor.replay_value()),
));
}
};
let stream = descriptor.stream().to_owned();
let subscribe_err = |source| SeaFileError::Subscribe {
stream: stream.clone(),
source,
};
let reader = if descriptor.replay_value() {
let file_id = FileId::new(self.core.path.clone());
let source = MessageSource::new(file_id, StreamMode::Replay)
.await
.map_err(|e| subscribe_err(box_err(e)))?;
Reader::Replay {
source: Box::new(source),
key,
}
} else {
let mut options = FileConsumerOptions::new(ConsumerMode::RealTime);
options.set_auto_stream_reset(AutoStreamReset::Latest);
options.set_live_streaming(true);
let streamer = disk.streamer.clone();
let consumer = disk
.runtime
.spawn(async move { streamer.create_consumer(&[key], options).await })
.await
.map_err(|e| subscribe_err(box_err(e)))?
.map_err(|e| subscribe_err(box_err(e)))?;
Reader::Tail(consumer)
};
let core = Arc::clone(&self.core);
Ok(FileSubscriber::spawn(
&disk.runtime,
stream,
reader,
move || core.closed.load(Ordering::Acquire),
))
}
}
impl ConnectedBroker for ConnectedFileBroker {
type Error = SeaFileError;
type Closed = ();
async fn shutdown(self) -> Result<(), Self::Error> {
self.core.closed.store(true, Ordering::Release);
let disk = match &self.core.link {
Link::File(disk) => disk,
#[cfg(feature = "testing")]
Link::InProcess(file) => {
file.finish();
return Ok(());
}
};
let taken = mem::replace(&mut *disk.producer.lock().await, Appender::Closed);
let producer = match taken {
Appender::Open(producer) | Appender::Failed(producer, _) => producer,
Appender::Closed => return Ok(()),
};
let ended = disk.streamer.clone().disconnect().await;
drop(producer);
let _settled = disk.streamer.clone().disconnect().await;
match ended {
Ok(()) | Err(StreamErr::Backend(FileErr::ProducerEnded)) => Ok(()),
Err(e) => Err(SeaFileError::Connect {
target: self.core.path.clone(),
source: box_err(e),
}),
}
}
}
impl Subscribe for ConnectedFileBroker {
type Subscriber = FileSubscriber;
type Copies = AddressedCopies;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
self.subscribe_stream(FileStream::new(name)).await
}
}
#[derive(Clone)]
pub struct FilePublisher {
cell: CoreCell,
}
impl std::fmt::Debug for FilePublisher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FilePublisher").finish_non_exhaustive()
}
}
impl Publisher for FilePublisher {
type Payload = Lend;
type Error = SeaFileError;
type Options = ();
async fn publish(
&self,
msg: OutgoingMessage<'_, &[u8]>,
_options: Option<&()>,
) -> Result<(), Self::Error> {
let core = self.cell.get().ok_or(SeaFileError::NotConnected)?;
core.ensure_open()?;
let key = StreamKey::new(msg.name())
.map_err(|e| SeaFileError::Invalid(format!("'{}': {e}", msg.name())))?;
let payload = wire::encode(msg.headers(), msg.payload(), false);
match &core.link {
Link::File(disk) => {
let mut appender = disk.producer.lock().await;
appender.append(msg.name(), &key, &payload).await
}
.map_err(|err| match err {
SeaFileError::NotConnected if !core.closed.load(Ordering::Acquire) => {
SeaFileError::Publish {
stream: msg.name().to_owned(),
source: box_err(FileErr::ProducerEnded),
}
}
other => other,
}),
#[cfg(feature = "testing")]
Link::InProcess(file) => {
file.append(key, payload);
Ok(())
}
}
}
}
impl Appender {
async fn append(
&mut self,
stream: &str,
key: &StreamKey,
payload: &[u8],
) -> Result<(), SeaFileError> {
let failed = |source| SeaFileError::Publish {
stream: stream.to_owned(),
source,
};
let producer = match self {
Self::Open(producer) => producer,
Self::Failed(_, err) => return Err(failed(box_err(err.take()))),
Self::Closed => return Err(SeaFileError::NotConnected),
};
let mut cause = match append(producer, key, payload).await {
Ok(()) => return Ok(()),
Err(StreamErr::Backend(cause)) => cause,
Err(other) => return Err(failed(box_err(other))),
};
let reported = cause.take();
if let Self::Open(producer) = mem::replace(self, Self::Closed) {
*self = Self::Failed(producer, cause);
}
Err(failed(box_err(reported)))
}
}
async fn append(
producer: &mut FileProducer,
key: &StreamKey,
payload: &[u8],
) -> Result<(), StreamErr<FileErr>> {
producer.send_to(key, payload)?.await?;
producer.flush().await
}
#[derive(Debug, Clone, Copy, Default)]
#[must_use]
pub struct FilePublish;
impl PublishPolicy<ConnectedFileBroker> for FilePublish {
type Live = FilePublisher;
fn pair(
self,
connected: &ConnectedFileBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.publisher()))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
channel_extension(channel)
}
}
impl DefaultPublish for ConnectedFileBroker {
type Policy = FilePublish;
}
#[cfg(feature = "testing")]
impl TestableBroker for ConnectedFileBroker {
fn install_coordinator(&self, coordinator: Coordinator) {
if let Link::InProcess(file) = &self.core.link {
file.install(coordinator);
}
}
fn inject(&self, message: OutgoingMessage<'_>) {
let file = self.memory_file("inject");
let key = StreamKey::new(message.name()).unwrap_or_else(|err| {
panic!(
"the injected message to {:?} is not one a stream file takes: {err}",
message.name()
)
});
file.append(
key,
wire::encode(message.headers(), message.payload(), false),
);
}
fn published(&self, name: &str) -> Vec<RawMessage> {
self.memory_file("published").published(name)
}
}
#[cfg(feature = "testing")]
impl ConnectedFileBroker {
fn memory_file(&self, what: &str) -> &MemoryFile {
match &self.core.link {
Link::InProcess(file) => file,
Link::File(_) => panic!(
"TestableBroker::{what} reached a broker connected with `connect`; the harness \
drives the file `connect_in_process` produces"
),
}
}
}
pub use FilePublish as Publish;
pub mod prelude {
pub use ruststream::prelude::*;
pub use ruststream::{Positioned, Seeker};
pub use crate::file::{FileBroker, Publish};
pub use crate::{
FileBatchContext, FileContext, FilePosition, FileSeeker, FileStream, Position, SeekHandle,
};
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use futures::future::join_all;
use ruststream::{Broker as _, ConnectedBroker as _, OutgoingMessage, Publisher as _};
use sea_streamer_file::{FileErr, FileId, MessageSource, StreamMode};
use sea_streamer_types::{Buffer as _, Message as _, StreamKey};
use tokio::runtime::Runtime;
use tokio::time::timeout;
use super::FileBroker;
const STEP: Duration = Duration::from_secs(10);
const LIMIT: u64 = 8 * 1024;
fn body(n: usize) -> Vec<u8> {
let mut bytes = format!("{n:06}").into_bytes();
bytes.resize(300, b'.');
bytes
}
async fn read_back(path: &str, key: &StreamKey) -> Vec<usize> {
let mut source = MessageSource::new(FileId::new(path), StreamMode::Replay)
.await
.expect("the file opens for a replay");
let mut found = Vec::new();
loop {
match source.next().await {
Ok(msg) if msg.message.stream_key() == *key => {
let bytes = msg.message.message().as_bytes().to_vec();
let n = std::str::from_utf8(&bytes[..6]).expect("a number").parse();
found.push(n.expect("a number"));
}
Ok(_) => {}
Err(FileErr::NotEnoughBytes | FileErr::StreamEnded) => return found,
Err(e) => panic!("the file reads back: {e}"),
}
}
}
async fn publish_past_the_limit(
path: &str,
count: usize,
concurrent: usize,
brokers: usize,
) -> Vec<usize> {
let mut connected = Vec::new();
let mut publishers = Vec::new();
for _ in 0..brokers {
let broker = FileBroker::new(path).size_limit(LIMIT);
connected.push(broker.clone().connect().await.expect("connects"));
publishers.push(broker.publisher());
}
let mut succeeded = Vec::new();
for wave in (0..count).collect::<Vec<_>>().chunks(concurrent) {
let results = join_all(wave.iter().map(|&n| {
let publisher = &publishers[n % brokers];
async move {
let bytes = body(n);
let publish =
publisher.publish(OutgoingMessage::new("orders", &bytes[..]), None);
let result = timeout(STEP, publish).await.unwrap_or_else(|_| {
panic!("publish {n} never returned ({concurrent} at a time)")
});
(n, result)
}
}))
.await;
succeeded.extend(
results
.into_iter()
.filter(|(_, r)| r.is_ok())
.map(|(n, _)| n),
);
}
for connected in connected {
let ended = timeout(STEP, connected.shutdown()).await;
assert!(
ended.is_ok(),
"the broker never shut down ({concurrent} at a time)"
);
}
succeeded
}
#[test]
fn a_publish_at_the_size_limit_returns_and_one_reported_done_is_whole_in_the_file() {
let rt = Runtime::new().expect("a runtime");
rt.block_on(async {
for (concurrent, brokers) in [(1, 1), (8, 1), (8, 8)] {
let path = std::env::temp_dir()
.join(format!(
"ruststream-sea-limit-{}-{concurrent}-{brokers}.ss",
std::process::id()
))
.to_string_lossy()
.into_owned();
let _stale = std::fs::remove_file(&path);
let succeeded = publish_past_the_limit(&path, 64, concurrent, brokers).await;
let key = StreamKey::new("orders").expect("a stream key");
let held = timeout(STEP, read_back(&path, &key)).await;
let held = held.unwrap_or_else(|_| {
panic!("reading the file back never ended ({concurrent} at a time)")
});
std::fs::remove_file(&path).expect("the file is removed");
assert!(
succeeded.len() < 64,
"the limit refused nothing ({concurrent} at a time)"
);
for n in &succeeded {
assert!(
held.contains(n),
"publish {n} reported success but the file does not hold it whole \
({concurrent} at a time; reported {succeeded:?}, held {held:?})"
);
}
}
});
}
}