use alloc::format;
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::time::Duration;
use std::net::TcpStream;
use std::thread::JoinHandle;
use liminal::protocol::{Frame, SchemaId, decode};
use liminal_protocol::outcome::ReconnectState;
use spin::Mutex;
use crate::SdkError;
use crate::remote::SETUP_TIMEOUT;
use crate::remote::mailbox::{
Mailbox, MailboxItem, MailboxLimits, MailboxStats, MailboxWords, ReaderEnd, offer_or_end,
};
const WEBSOCKET_SUBSCRIPTION_WORDS: MailboxWords = MailboxWords {
receive: "websocket subscription receive failed",
item: "delivery",
reader: "subscription reader",
};
use super::binding::{AttemptFateOutcome, OpenRequestDecision, WebSocketAuthorityBinding};
use super::connection_error;
use super::core::{
DriverOutput, FrameCorrelation, ResponseExpectation, SocketCommand, SocketEvent,
WebSocketFrameDriver,
};
use super::liminal_ws_message_bound;
use super::std_socket::{SocketRead, WsSocket};
const CLIENT_MIN_VERSION: liminal::protocol::ProtocolVersion =
liminal::protocol::ProtocolVersion::new(1, 0);
const CLIENT_MAX_VERSION: liminal::protocol::ProtocolVersion =
liminal::protocol::DELIVER_MESSAGE_ID_VERSION;
const SUBSCRIPTION_STREAM_ID: u32 = 1;
const SUBSCRIBE_MAX_IN_FLIGHT: u32 = 1024;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WebSocketDeliveredMessage {
delivery_seq: u64,
schema_id: SchemaId,
payload: Vec<u8>,
message: Option<liminal::protocol::DeliveredId>,
}
impl WebSocketDeliveredMessage {
#[must_use]
pub const fn delivery_seq(&self) -> u64 {
self.delivery_seq
}
#[must_use]
pub const fn schema_id(&self) -> SchemaId {
self.schema_id
}
#[must_use]
pub fn payload(&self) -> &[u8] {
&self.payload
}
#[must_use]
pub fn into_payload(self) -> Vec<u8> {
self.payload
}
#[must_use]
pub const fn message(&self) -> Option<liminal::protocol::DeliveredId> {
self.message
}
#[must_use]
pub fn reply_parent(&self) -> Option<liminal::protocol::MessageId> {
self.message
.map(liminal::protocol::DeliveredId::reply_parent)
}
}
impl MailboxItem for WebSocketDeliveredMessage {
fn payload_len(&self) -> usize {
self.payload.len()
}
}
#[derive(Debug)]
pub struct WebSocketSubscriptionStream {
pub(super) shutdown: TcpStream,
subscription_id: u64,
pub(super) mailbox: Arc<Mailbox<WebSocketDeliveredMessage>>,
binding: Arc<Mutex<WebSocketAuthorityBinding>>,
reader: Option<JoinHandle<()>>,
}
impl WebSocketSubscriptionStream {
pub fn open(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
) -> Result<Self, SdkError> {
Self::open_with_auth(address, channel, accepted_schemas, &[])
}
pub fn open_with_auth(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
auth_token: &[u8],
) -> Result<Self, SdkError> {
Self::open_with_limits(
address,
channel,
accepted_schemas,
auth_token,
MailboxLimits::SUBSCRIPTION,
)
}
fn open_with_limits(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
auth_token: &[u8],
limits: MailboxLimits,
) -> Result<Self, SdkError> {
let message_bound = liminal_ws_message_bound()?;
let mut binding = WebSocketAuthorityBinding::new();
match binding.request_open() {
OpenRequestDecision::Authorized { .. } => {}
OpenRequestDecision::Refused(refusal) => {
return Err(connection_error(&format!(
"client authority refused the subscription open: {refusal:?}"
)));
}
}
match Self::open_link(
address,
channel,
accepted_schemas,
message_bound,
auth_token,
) {
Ok((socket, driver, subscription_id, pending)) => {
match binding.connection_established() {
AttemptFateOutcome::Recorded { .. } => {}
AttemptFateOutcome::Refused(refusal) => {
return Err(SdkError::Protocol {
description: format!(
"client authority refused the Connected fate for the \
subscription open: {refusal:?}"
),
});
}
}
Self::start(socket, driver, binding, subscription_id, pending, limits)
}
Err(error) => match binding.open_failed() {
AttemptFateOutcome::Recorded { .. } => Err(error),
AttemptFateOutcome::Refused(refusal) => Err(SdkError::Protocol {
description: format!(
"subscription open failed ({error}) and the client authority \
refused the Failed fate: {refusal:?}"
),
}),
},
}
}
pub fn recv_timeout(&self, timeout: Duration) -> Result<WebSocketDeliveredMessage, SdkError> {
self.mailbox.take(timeout)
}
pub fn stats(&self) -> Result<MailboxStats, SdkError> {
self.mailbox.stats()
}
#[must_use]
pub const fn subscription_id(&self) -> u64 {
self.subscription_id
}
#[must_use]
pub fn reconnect_state(&self) -> ReconnectState {
self.binding.lock().reconnect_state()
}
fn open_link(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
message_bound: usize,
auth_token: &[u8],
) -> Result<
(
WsSocket,
WebSocketFrameDriver,
u64,
Vec<WebSocketDeliveredMessage>,
),
SdkError,
> {
let mut driver = WebSocketFrameDriver::new();
let command = driver
.command_open()
.map_err(|refusal| SdkError::Protocol {
description: format!("subscription driver refused its first open: {refusal:?}"),
})?;
if command != SocketCommand::Open {
return Err(SdkError::Protocol {
description: "subscription driver emitted a non-open first command".to_string(),
});
}
let mut socket = WsSocket::connect(address, message_bound)?;
socket.set_read_timeout(Some(SETUP_TIMEOUT))?;
let step = driver.handle_event(SocketEvent::Opened);
if step.output != DriverOutput::Opened {
return Err(SdkError::Protocol {
description: format!("subscription driver refused the opened socket: {step:?}"),
});
}
let mut pending = Vec::new();
let connect = Frame::Connect {
flags: 0,
min_version: CLIENT_MIN_VERSION,
max_version: CLIENT_MAX_VERSION,
auth_token: auth_token.to_vec(),
};
match setup_exchange(&mut socket, &mut driver, &connect, &mut pending)? {
Frame::ConnectAck { .. } => {}
Frame::ConnectError {
reason_code,
message,
..
} => {
return Err(connection_error(&format!(
"server rejected subscription connection (reason {reason_code}): {}",
message.unwrap_or_else(|| "no detail".to_string())
)));
}
other => {
return Err(unexpected_setup_frame("ConnectAck", &other));
}
}
let subscribe = Frame::Subscribe {
flags: 0,
stream_id: SUBSCRIPTION_STREAM_ID,
channel: channel.to_string(),
accepted_schemas,
max_in_flight: SUBSCRIBE_MAX_IN_FLIGHT,
};
let subscription_id =
match setup_exchange(&mut socket, &mut driver, &subscribe, &mut pending)? {
Frame::SubscribeAck {
subscription_id, ..
} => subscription_id,
Frame::SubscribeError {
reason_code,
message,
..
} => {
return Err(SdkError::Protocol {
description: format!(
"server rejected subscribe (reason {reason_code}): {}",
message.unwrap_or_else(|| "no detail".to_string())
),
});
}
other => {
return Err(unexpected_setup_frame("SubscribeAck", &other));
}
};
Ok((socket, driver, subscription_id, pending))
}
fn start(
socket: WsSocket,
driver: WebSocketFrameDriver,
binding: WebSocketAuthorityBinding,
subscription_id: u64,
pending: Vec<WebSocketDeliveredMessage>,
limits: MailboxLimits,
) -> Result<Self, SdkError> {
socket.set_read_timeout(None)?;
let shutdown = socket.try_clone_stream()?;
let binding = Arc::new(Mutex::new(binding));
let reader_binding = Arc::clone(&binding);
let mailbox = Arc::new(Mailbox::new(limits, WEBSOCKET_SUBSCRIPTION_WORDS));
let reader_mailbox = Arc::clone(&mailbox);
let reader = std::thread::Builder::new()
.name("liminal-ws-subscription-reader".to_string())
.spawn(move || {
let end = run_reader(socket, driver, &reader_binding, &reader_mailbox, pending);
reader_mailbox.close(end);
})
.map_err(|source| SdkError::Protocol {
description: format!(
"failed to start websocket subscription reader thread: {source}"
),
})?;
Ok(Self {
shutdown,
subscription_id,
mailbox,
binding,
reader: Some(reader),
})
}
}
impl Drop for WebSocketSubscriptionStream {
fn drop(&mut self) {
self.shutdown.shutdown(std::net::Shutdown::Both).ok();
self.mailbox.release();
if let Some(reader) = self.reader.take() {
reader.join().ok();
}
}
}
fn setup_exchange(
socket: &mut WsSocket,
driver: &mut WebSocketFrameDriver,
request: &Frame,
pending: &mut Vec<WebSocketDeliveredMessage>,
) -> Result<Frame, SdkError> {
let bytes = super::encode_frame(request)?;
let command = driver
.command_send(bytes, ResponseExpectation::Correlated)
.map_err(|refusal| SdkError::Protocol {
description: format!("subscription driver refused the setup send: {refusal:?}"),
})?;
let SocketCommand::SendBinary(payload) = command else {
return Err(SdkError::Protocol {
description: "subscription driver emitted a non-send command for a send".to_string(),
});
};
if let Err(failure) = socket.send_binary(payload) {
let step = driver.handle_event(SocketEvent::Failed(failure));
if step.command == Some(SocketCommand::Close) {
socket.execute_close();
}
return Err(connection_error(&format!(
"failed to send subscription setup frame: {}",
socket
.last_failure_detail()
.unwrap_or("websocket send failed")
)));
}
loop {
let event = match socket.read_event() {
SocketRead::TimedOut => {
return Err(connection_error(
"subscription connection timed out waiting for a control-frame reply",
));
}
SocketRead::Event(event) => event,
};
let step = driver.handle_event(event);
if step.command == Some(SocketCommand::Close) {
socket.execute_close();
}
match step.output {
DriverOutput::Frame { bytes, correlation } => {
let frame = decode_message(&bytes)?;
match correlation {
FrameCorrelation::UnsolicitedDelivery => {
if let Some(message) = delivered_message(frame) {
pending.push(message);
}
}
FrameCorrelation::CorrelatedResponse | FrameCorrelation::UnsolicitedFrame => {
return Ok(frame);
}
}
}
DriverOutput::Terminal(terminal) => {
return Err(connection_error(&format!(
"subscription connection terminated during setup: {terminal:?}"
)));
}
DriverOutput::Opened
| DriverOutput::PostTerminalIgnored(_)
| DriverOutput::Refused(_) => {
return Err(SdkError::Protocol {
description: format!(
"subscription driver produced an unexpected setup output: {:?}",
step.output
),
});
}
}
}
}
fn run_reader(
mut socket: WsSocket,
mut driver: WebSocketFrameDriver,
binding: &Mutex<WebSocketAuthorityBinding>,
mailbox: &Mailbox<WebSocketDeliveredMessage>,
pending: Vec<WebSocketDeliveredMessage>,
) -> ReaderEnd {
let mut first_error: Option<ReaderEnd> = None;
for message in pending {
if let Some(end) = offer_or_end(mailbox, message) {
close_link(&mut socket, &mut driver);
first_error = Some(end);
break;
}
}
loop {
let event = match socket.read_event() {
SocketRead::TimedOut => continue,
SocketRead::Event(event) => event,
};
let step = driver.handle_event(event);
if step.command == Some(SocketCommand::Close) {
socket.execute_close();
}
match step.output {
DriverOutput::Frame { bytes, correlation } => {
if first_error.is_some() {
continue;
}
first_error = handle_frame(&bytes, correlation, mailbox);
if first_error.is_some() {
close_link(&mut socket, &mut driver);
}
}
DriverOutput::Terminal(terminal) => {
let _outcome = binding.lock().established_terminal(&terminal);
return first_error.unwrap_or(ReaderEnd::Stopped);
}
DriverOutput::PostTerminalIgnored(_) => {
return first_error.unwrap_or(ReaderEnd::Stopped);
}
DriverOutput::Opened | DriverOutput::Refused(_) => {}
}
}
}
fn handle_frame(
bytes: &[u8],
correlation: FrameCorrelation,
mailbox: &Mailbox<WebSocketDeliveredMessage>,
) -> Option<ReaderEnd> {
let frame = match decode_message(bytes) {
Ok(frame) => frame,
Err(error) => return Some(ReaderEnd::Failed(error)),
};
match (correlation, frame) {
(
FrameCorrelation::UnsolicitedDelivery,
Frame::Deliver {
delivery_seq,
envelope,
message,
..
},
) => offer_or_end(
mailbox,
WebSocketDeliveredMessage {
delivery_seq,
schema_id: envelope.schema_id,
payload: envelope.payload,
message,
},
),
(_, Frame::Disconnect { .. }) => Some(ReaderEnd::Stopped),
(
_,
Frame::SubscribeError {
reason_code,
message,
..
},
) => Some(ReaderEnd::Failed(subscription_ended(reason_code, message))),
_ => mailbox.record_ignored_frame().err().map(ReaderEnd::Failed),
}
}
fn close_link(socket: &mut WsSocket, driver: &mut WebSocketFrameDriver) {
if driver.command_close().is_ok() {
socket.execute_close();
}
}
fn decode_message(bytes: &[u8]) -> Result<Frame, SdkError> {
match decode(bytes) {
Ok((frame, consumed)) if consumed == bytes.len() => Ok(frame),
Ok((_, consumed)) => Err(SdkError::Protocol {
description: format!(
"subscription decode consumed {consumed} of {} message bytes",
bytes.len()
),
}),
Err(error) => Err(SdkError::Protocol {
description: format!("subscription wire codec error: {error}"),
}),
}
}
fn delivered_message(frame: Frame) -> Option<WebSocketDeliveredMessage> {
match frame {
Frame::Deliver {
delivery_seq,
envelope,
message,
..
} => Some(WebSocketDeliveredMessage {
delivery_seq,
schema_id: envelope.schema_id,
payload: envelope.payload,
message,
}),
_ => None,
}
}
fn subscription_ended(reason_code: u16, message: Option<alloc::string::String>) -> SdkError {
SdkError::Protocol {
description: format!(
"server ended the subscription (reason {reason_code}): {}",
message.unwrap_or_else(|| "no detail".to_string())
),
}
}
fn unexpected_setup_frame(expected: &str, actual: &Frame) -> SdkError {
SdkError::Protocol {
description: format!(
"expected {expected} during subscription setup, received {:?}",
actual.frame_type()
),
}
}
#[cfg(test)]
#[path = "subscription_decode_tests.rs"]
mod decode_tests;
#[cfg(test)]
mod tests {
use super::SETUP_TIMEOUT;
use core::time::Duration;
#[test]
fn websocket_subscription_source_has_no_retired_reader_poll_family() {
const SOURCE: &str = include_str!("subscription.rs");
let production = SOURCE.split("#[cfg(test)]").next().unwrap_or(SOURCE);
for forbidden in [
"READER_POLL_TIMEOUT",
"AtomicBool",
"stop.load",
"stop.store",
"re-check the stop flag",
"poll the stop flag",
] {
assert!(
!production.contains(forbidden),
"retired websocket-subscription-reader poll-family source \
`{forbidden}` reappeared"
);
}
}
#[test]
fn the_named_setup_deadline_is_the_ratified_five_seconds() {
assert_eq!(SETUP_TIMEOUT, Duration::from_secs(5));
}
const DELIVERED: liminal::protocol::DeliveredId = liminal::protocol::DeliveredId {
boot: 1_791_328_512_345,
id: 42,
};
fn deliver(message: Option<liminal::protocol::DeliveredId>) -> liminal::protocol::Frame {
liminal::protocol::Frame::Deliver {
flags: 0,
stream_id: super::SUBSCRIPTION_STREAM_ID,
delivery_seq: 1,
envelope: liminal::protocol::MessageEnvelope::new(
liminal::protocol::SchemaId::new([7; liminal::protocol::SchemaId::WIRE_LEN]),
liminal::protocol::CausalContext::independent(),
alloc::vec![1],
),
message,
}
}
#[test]
fn the_subscription_offers_the_delivered_id_version() {
assert_eq!(
super::CLIENT_MAX_VERSION,
liminal::protocol::DELIVER_MESSAGE_ID_VERSION
);
}
#[test]
fn a_delivered_message_keeps_its_id_and_names_its_reply_parent() {
let with_id = super::delivered_message(deliver(Some(DELIVERED)));
assert_eq!(
with_id
.as_ref()
.and_then(super::WebSocketDeliveredMessage::message),
Some(DELIVERED)
);
assert_eq!(
with_id.and_then(|message| message.reply_parent()),
Some(liminal::protocol::MessageId::from("1791328512345:42"))
);
let without = super::delivered_message(deliver(None));
assert_eq!(
without
.as_ref()
.and_then(super::WebSocketDeliveredMessage::message),
None
);
assert_eq!(without.and_then(|message| message.reply_parent()), None);
}
}