#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
use std::fs;
use std::future::{Future, ready};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[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::{OnceCell, RwLock};
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: RwLock<Option<FileProducer>>,
pub(crate) runtime: Handle,
}
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>,
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,
cell: Arc::new(OnceCell::new()),
}
}
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);
}
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 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)?;
}
Ok::<_, SeaFileError>(Arc::new(Core {
link: Link::File(Disk {
streamer,
producer: RwLock::new(Some(producer)),
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 Some(producer) = disk.producer.write().await.take() else {
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);
let held = match &core.link {
Link::File(disk) => disk.producer.read().await,
#[cfg(feature = "testing")]
Link::InProcess(file) => {
file.append(key, payload);
return Ok(());
}
};
let producer = held.as_ref().ok_or(SeaFileError::NotConnected)?;
let appended = append(producer, &key, &payload).await;
drop(held);
appended.map_err(|e| SeaFileError::Publish {
stream: msg.name().to_owned(),
source: box_err(e),
})
}
}
async fn append(
producer: &FileProducer,
key: &StreamKey,
payload: &[u8],
) -> Result<(), StreamErr<FileErr>> {
producer.send_to(key, payload)?.await?;
let mut flusher = producer.clone();
flusher.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,
};
}