use std::collections::VecDeque;
use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::DBOS;
use crate::checkpoint::{Built, PendingStep, StepPlacement, StepWriter, revive};
use crate::connection::Connection;
use crate::context::Ctx;
use crate::error::{DurableError, EngineOnly, Error, Result};
use crate::serialization::SerializationFormat;
use crate::serialization::{PORTABLE_JSON, PayloadEncoding, decode_tagged, encode};
use crate::sysdb::types::{
AwaitedStream, EncodedValue, Outcome, StepTiming, Timestamp, WrittenBy, step_names,
};
use crate::sysdb::{STREAM_CLOSED, is_stream_closed};
pub const DEFAULT_STREAM_POLLING_INTERVAL: Duration = Duration::from_secs(1);
pub const MIN_STREAM_POLLING_INTERVAL: Duration = Duration::from_millis(1);
const STREAM_PAGE: i32 = 100;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct ReadStreamOptions {
pub polling_interval: Option<Duration>,
pub timeout: Option<Duration>,
}
pub fn write_stream<'a, T, E>(key: &'a str, value: &T) -> PendingStep<'a, (), E>
where
T: Serialize,
E: DurableError + 'a,
{
write_stream_with(key, value, WriteStreamOptions::default())
}
pub fn write_stream_with<'a, T, E>(
key: &'a str,
value: &T,
options: WriteStreamOptions,
) -> PendingStep<'a, (), E>
where
T: Serialize,
E: DurableError + 'a,
{
pending_write(key, StreamWrite::Value(value, options.serialization))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct WriteStreamOptions {
pub serialization: SerializationFormat,
}
pub fn close_stream<'a, E>(key: &'a str) -> PendingStep<'a, (), E>
where
E: DurableError + 'a,
{
pending_write::<(), E>(key, StreamWrite::Close)
}
enum StreamWrite<'v, T> {
Value(&'v T, SerializationFormat),
Close,
}
impl<T> StreamWrite<'_, T> {
fn name(&self) -> &'static str {
match self {
StreamWrite::Value(..) => step_names::WRITE_STREAM,
StreamWrite::Close => step_names::CLOSE_STREAM,
}
}
fn operation(&self) -> &'static str {
match self {
StreamWrite::Value(..) => "write_stream",
StreamWrite::Close => "close_stream",
}
}
}
fn pending_write<'a, T, E>(key: &'a str, write: StreamWrite<'_, T>) -> PendingStep<'a, (), E>
where
T: Serialize,
E: DurableError + 'a,
{
let name = write.name();
let built = Ctx::current()
.ok_or_else(|| Error::NotInWorkflow {
operation: write.operation().into(),
})
.and_then(|ctx| {
let encoded = match write {
StreamWrite::Value(value, serialization) => {
let encoding = PayloadEncoding::resolve(serialization, ctx.encoding());
Some((encode(value, &encoding, "stream value")?, encoding))
}
StreamWrite::Close => None,
};
Ok((
(Arc::clone(ctx.executor().connection()), encoded),
StepPlacement::at(ctx),
))
});
PendingStep::placed(name, built, move |(conn, encoded), placement| async move {
let (workflow_id, step_id, owner_xid, written_by) = match &placement {
StepPlacement::Recorded { ctx, step_id } => (
ctx.workflow_id(),
*step_id,
ctx.owner_xid(),
WrittenBy::Workflow,
),
StepPlacement::InsideStep { ctx } => (
ctx.workflow_id(),
ctx.step_id()
.expect("a context inside a step knows its step id"),
ctx.owner_xid(),
WrittenBy::Step,
),
StepPlacement::Outside | StepPlacement::ClientConnection => {
unreachable!("a stream write is placed from the ambient workflow's own context")
}
};
let sysdb = conn.sysdb();
match &encoded {
Some((value, encoding)) => {
sysdb
.write_stream(
workflow_id,
step_id,
key,
value,
Some(encoding.name()),
written_by,
owner_xid,
)
.await
}
None => {
sysdb
.close_stream(workflow_id, step_id, key, written_by, owner_xid)
.await
}
}
.map_err(Error::from_execution)
.map_err(Error::lift)
})
}
pub fn read_stream<T, E>(
workflow_id: &str,
key: &str,
options: ReadStreamOptions,
) -> StreamReader<T, E>
where
T: DeserializeOwned,
E: DurableError,
{
StreamReader::new(
StepPlacement::ambient_connection("read_stream"),
workflow_id,
key,
0,
options,
"read_stream",
)
}
pub fn read_stream_from<T, E>(
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
) -> StreamReader<T, E>
where
T: DeserializeOwned,
E: DurableError,
{
StreamReader::new(
StepPlacement::ambient_connection("read_stream_from"),
workflow_id,
key,
offset,
options,
"read_stream_from",
)
}
pub fn read_stream_value<'a, T, E>(
workflow_id: &'a str,
key: &'a str,
offset: i32,
options: ReadStreamOptions,
) -> PendingStep<'a, T, E>
where
T: DeserializeOwned + 'a,
E: DurableError + 'a,
{
read_value(
StepPlacement::ambient_connection("read_stream_value"),
workflow_id,
key,
offset,
options,
)
}
impl DBOS {
pub fn read_stream<T: DeserializeOwned>(
&self,
workflow_id: &str,
key: &str,
options: ReadStreamOptions,
) -> StreamReader<T> {
StreamReader::new(
self.executor("read_stream")
.map(|executor| Arc::clone(executor.connection())),
workflow_id,
key,
0,
options,
"read_stream",
)
}
pub fn read_stream_from<T: DeserializeOwned>(
&self,
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
) -> StreamReader<T> {
StreamReader::new(
self.executor("read_stream_from")
.map(|executor| Arc::clone(executor.connection())),
workflow_id,
key,
offset,
options,
"read_stream_from",
)
}
pub fn read_stream_value<'a, T: DeserializeOwned + 'a>(
&self,
workflow_id: &'a str,
key: &'a str,
offset: i32,
options: ReadStreamOptions,
) -> PendingStep<'a, T> {
read_value(
self.executor("read_stream_value")
.map(|executor| Arc::clone(executor.connection())),
workflow_id,
key,
offset,
options,
)
}
}
impl crate::Client {
pub fn read_stream<T: DeserializeOwned>(
&self,
workflow_id: &str,
key: &str,
options: ReadStreamOptions,
) -> StreamReader<T> {
StreamReader::new(
Ok(Arc::clone(self.connection())),
workflow_id,
key,
0,
options,
"read_stream",
)
}
pub fn read_stream_from<T: DeserializeOwned>(
&self,
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
) -> StreamReader<T> {
StreamReader::new(
Ok(Arc::clone(self.connection())),
workflow_id,
key,
offset,
options,
"read_stream_from",
)
}
pub async fn read_stream_value<T: DeserializeOwned>(
&self,
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
) -> Result<T> {
read_value(
Ok(Arc::clone(self.connection())),
workflow_id,
key,
offset,
options,
)
.await
}
}
#[must_use = "a stream reader reads nothing until `next` is awaited"]
pub struct StreamReader<T, E = EngineOnly> {
source: std::result::Result<Arc<Connection>, Option<Error>>,
operation: &'static str,
read: ReadOf,
offset: i32,
buffer: VecDeque<EncodedValue>,
ended: bool,
replaying: bool,
marker: PhantomData<fn() -> (T, E)>,
}
impl<T, E> std::fmt::Debug for StreamReader<T, E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StreamReader")
.field("workflow_id", &self.read.workflow_id)
.field("key", &self.read.key)
.field("offset", &self.offset)
.field("ended", &self.ended)
.finish_non_exhaustive()
}
}
impl<T, E> StreamReader<T, E>
where
T: DeserializeOwned,
E: DurableError,
{
fn new(
source: Result<Arc<Connection>>,
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
operation: &'static str,
) -> Self {
let read = ReadOf::new(workflow_id, key, options);
let source = source.and_then(|conn| {
read.validate(offset, operation)?;
Ok(conn)
});
Self {
source: source.map_err(Some),
operation,
read,
offset,
buffer: VecDeque::new(),
ended: false,
replaying: true,
marker: PhantomData,
}
}
#[allow(clippy::should_implement_trait)]
pub fn next(&mut self) -> PendingStep<'_, Option<T>, E> {
let Some(built) = self.place() else {
return PendingStep::settled(step_names::READ_STREAM, || Ok(None));
};
PendingStep::placed(
step_names::READ_STREAM,
built,
move |conn, placement| async move {
let read = self.read_next(&conn, &placement).await;
if read.is_err() {
self.ended = true;
}
read.map_err(Error::lift)
},
)
}
#[must_use]
pub fn offset(&self) -> i32 {
self.offset
}
fn place(&mut self) -> Option<Built<Arc<Connection>>> {
if self.ended {
return None;
}
let built = match &mut self.source {
Ok(conn) => place(Arc::clone(conn), self.operation),
Err(refused) => Err(refused
.take()
.expect("a refusal is reported once, and the reader then ends")),
};
self.ended = built.is_err();
Some(built)
}
async fn read_next(
&mut self,
conn: &Connection,
placement: &StepPlacement,
) -> Result<Option<T>> {
let name = step_names::READ_STREAM;
let started_at = Timestamp::now();
if self.replaying {
match check(conn, placement, name).await? {
Some(Replayed::Value(output, tag)) => {
self.offset += 1;
return decode_tagged(
Some(&output),
tag.as_deref(),
conn.serializer(),
"stream value",
)
.map(Some);
}
Some(Replayed::Ended) => {
self.ended = true;
return Ok(None);
}
None => self.replaying = placement.step_position().is_none(),
}
}
if self.buffer.is_empty() {
match self
.read
.await_values(conn, placement, name, self.offset, STREAM_PAGE, started_at)
.await?
{
AwaitedStream::Values(values) => {
let mut values = values.into_iter();
let first = values.next().expect("a page of values is never empty");
let closed = is_close(&first);
self.buffer.push_back(first);
if !closed {
self.buffer
.extend(values.take_while(|later| !is_close(later)));
}
}
AwaitedStream::Ended => {
record_end(conn, placement, name, started_at).await?;
self.ended = true;
return Ok(None);
}
AwaitedStream::TimedOut => {
return Err(self
.read
.time_out(conn, placement, name, started_at)
.await?);
}
}
}
let front = self.buffer.front().expect("a value is buffered by now");
if is_close(front) {
record_end(conn, placement, name, started_at).await?;
self.buffer.clear();
self.ended = true;
return Ok(None);
}
record_value(conn, placement, name, front, started_at).await?;
let entry = self.buffer.pop_front().expect("the value just recorded");
self.offset += 1;
decode_tagged(
Some(&entry.value),
entry.serialization.as_deref(),
conn.serializer(),
"stream value",
)
.map(Some)
}
}
fn read_value<'a, T, E>(
source: Result<Arc<Connection>>,
workflow_id: &str,
key: &str,
offset: i32,
options: ReadStreamOptions,
) -> PendingStep<'a, T, E>
where
T: DeserializeOwned + 'a,
E: DurableError + 'a,
{
let name = step_names::READ_STREAM_VALUE;
let read = ReadOf::new(workflow_id, key, options);
let built = source
.and_then(|conn| {
read.validate(offset, "read_stream_value")?;
Ok(conn)
})
.and_then(|conn| place(conn, "read_stream_value"));
PendingStep::placed(name, built, move |conn, placement| async move {
async {
let started_at = Timestamp::now();
let ended = || Error::StreamTimeout {
workflow_id: read.workflow_id.clone(),
key: read.key.clone(),
timeout: None,
};
match check(&conn, &placement, name).await? {
Some(Replayed::Value(output, tag)) => {
return decode_tagged(
Some(&output),
tag.as_deref(),
conn.serializer(),
"stream value",
);
}
Some(Replayed::Ended) => return Err(ended()),
None => {}
}
let entry = match read
.await_values(&conn, &placement, name, offset, 1, started_at)
.await?
{
AwaitedStream::Values(values) => values
.into_iter()
.next()
.expect("a page of values is never empty"),
AwaitedStream::Ended => {
record_end(&conn, &placement, name, started_at).await?;
return Err(ended());
}
AwaitedStream::TimedOut => {
return Err(read.time_out(&conn, &placement, name, started_at).await?);
}
};
if is_close(&entry) {
record_end(&conn, &placement, name, started_at).await?;
return Err(ended());
}
record_value(&conn, &placement, name, &entry, started_at).await?;
decode_tagged(
Some(&entry.value),
entry.serialization.as_deref(),
conn.serializer(),
"stream value",
)
}
.await
.map_err(Error::lift)
})
}
#[derive(Debug, Clone)]
struct ReadOf {
workflow_id: String,
key: String,
polling_interval: Duration,
timeout: Option<Duration>,
}
impl ReadOf {
fn new(workflow_id: &str, key: &str, options: ReadStreamOptions) -> Self {
Self {
workflow_id: workflow_id.to_owned(),
key: key.to_owned(),
polling_interval: options
.polling_interval
.unwrap_or(DEFAULT_STREAM_POLLING_INTERVAL),
timeout: options.timeout,
}
}
fn validate(&self, offset: i32, operation: &'static str) -> Result<()> {
if offset < 0 {
return Err(Error::InvalidArgument {
operation: operation.into(),
detail: format!("the offset must not be negative, got {offset}"),
});
}
if self.polling_interval < MIN_STREAM_POLLING_INTERVAL {
return Err(Error::InvalidArgument {
operation: operation.into(),
detail: format!(
"the polling interval must be at least {}ms, got {:?}",
MIN_STREAM_POLLING_INTERVAL.as_millis(),
self.polling_interval
),
});
}
Ok(())
}
async fn await_values(
&self,
conn: &Connection,
placement: &StepPlacement,
name: &str,
offset: i32,
page: i32,
started_at: Timestamp,
) -> Result<AwaitedStream> {
let deadline = self.timeout.map(|timeout| {
started_at
.checked_add(timeout)
.unwrap_or(Timestamp::from_epoch_ms(i64::MAX))
});
let wait = conn.sysdb().await_stream_values(
&self.workflow_id,
&self.key,
offset,
page,
deadline,
self.polling_interval,
);
let awaited = match placement {
StepPlacement::Recorded { ctx, .. } => tokio::select! {
awaited = wait => awaited,
() = crate::step::observe_cancellation(ctx) => {
return Err(Error::WorkflowCancelled {
workflow_id: ctx.workflow_id().to_owned(),
});
}
},
_ => wait.await,
};
match awaited {
Err(refused) if refused.should_record() => {
let refused = Error::SystemDatabase(refused);
let encoding = conn.serializer().encoding();
let encoded = encode(&refused, &encoding, "stream read refusal")?;
record(
conn,
placement,
name,
Outcome::Error(&encoded),
Some(encoding.name()),
started_at,
)
.await?;
Err(refused)
}
awaited => awaited.map_err(Error::from_execution),
}
}
async fn time_out(
&self,
conn: &Connection,
placement: &StepPlacement,
name: &str,
started_at: Timestamp,
) -> Result<Error> {
let timed_out = Error::StreamTimeout {
workflow_id: self.workflow_id.clone(),
key: self.key.clone(),
timeout: self.timeout,
};
let encoding = conn.serializer().encoding();
let encoded = encode(&timed_out, &encoding, "stream timeout")?;
record(
conn,
placement,
name,
Outcome::Error(&encoded),
Some(encoding.name()),
started_at,
)
.await?;
Ok(timed_out)
}
}
fn place(conn: Arc<Connection>, operation: &'static str) -> Built<Arc<Connection>> {
let placement = StepPlacement::of(&conn, operation)?;
Ok((conn, placement))
}
enum Replayed {
Value(String, Option<String>),
Ended,
}
async fn check(
conn: &Connection,
placement: &StepPlacement,
name: &str,
) -> Result<Option<Replayed>> {
let Some(step) = placement.check(conn, name).await? else {
return Ok(None);
};
tracing::debug!(step_id = ?placement.step_id(), name, "replaying a stream read");
if let Some(error) = step.error {
return Err(revive(
&error,
step.serialization.as_deref(),
conn.serializer(),
name,
));
}
let output = step.output.unwrap_or_default();
Ok(Some(
if is_stream_closed(&output, step.serialization.as_deref()) {
Replayed::Ended
} else {
Replayed::Value(output, step.serialization)
},
))
}
async fn record_value(
conn: &Connection,
placement: &StepPlacement,
name: &str,
entry: &EncodedValue,
started_at: Timestamp,
) -> Result<()> {
record(
conn,
placement,
name,
Outcome::Output(Some(&entry.value)),
entry.serialization.as_deref(),
started_at,
)
.await
}
async fn record_end(
conn: &Connection,
placement: &StepPlacement,
name: &str,
started_at: Timestamp,
) -> Result<()> {
record(
conn,
placement,
name,
Outcome::Output(Some(STREAM_CLOSED)),
Some(PORTABLE_JSON),
started_at,
)
.await
}
fn is_close(entry: &EncodedValue) -> bool {
is_stream_closed(&entry.value, entry.serialization.as_deref())
}
async fn record(
conn: &Connection,
placement: &StepPlacement,
name: &str,
outcome: Outcome<'_>,
serialization: Option<&str>,
started_at: Timestamp,
) -> Result<()> {
let Some(StepWriter {
workflow_id,
step_id,
owner_xid,
}) = placement.step_writer()
else {
return Ok(());
};
conn.sysdb()
.record_step(
workflow_id,
step_id,
name,
outcome,
serialization,
Some(StepTiming {
started_at,
completed_at: Timestamp::now(),
}),
owner_xid,
)
.await
.map_err(Error::from_execution)
}