mod pending_connect;
pub use pending_connect::PendingPushConnect;
use crate::remote::mailbox::{
Mailbox, MailboxItem, MailboxLimits, MailboxStats, MailboxWords, ReaderEnd, offer_or_end,
};
const PUSH_WORDS: MailboxWords = MailboxWords {
receive: "push receive failed",
item: "server push",
reader: "push reader",
};
use alloc::format;
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec;
use alloc::vec::Vec;
use core::time::Duration;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::sync::Mutex;
use std::sync::mpsc::Sender;
use std::thread::JoinHandle;
use std::time::Instant;
use liminal::protocol::{
CausalContext, Frame, MessageEnvelope, PUBLISH_ACK_POSITION_VERSION, ProtocolError,
ProtocolVersion, SchemaId, WorkerRegisterOutcome, WorkerRegistration, decode, encode,
encoded_len,
};
use super::flush::{
FLUSH_BUDGET, FlushLedger, FlushMode, FlushOutcome, PublishRejection, PublishVerdict,
};
use crate::SdkError;
use crate::remote::SETUP_TIMEOUT;
const CLIENT_MIN_VERSION: ProtocolVersion = ProtocolVersion::new(1, 0);
const CLIENT_MAX_VERSION: ProtocolVersion = PUBLISH_ACK_POSITION_VERSION;
const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
const READ_CHUNK_BYTES: usize = 4096;
const DROP_DRAIN_BUDGET: Duration = Duration::from_secs(5);
const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
const APPLICATION_STREAM_ID: u32 = 1;
pub const OBSERVABILITY_CHANNEL: &str = "aion.observability.v1";
#[derive(Clone, PartialEq, Eq)]
pub struct PushedFrame {
correlation_id: u64,
payload: Vec<u8>,
}
impl core::fmt::Debug for PushedFrame {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
formatter
.debug_struct("PushedFrame")
.field("correlation_id", &self.correlation_id)
.field("payload_bytes", &self.payload.len())
.finish()
}
}
impl MailboxItem for PushedFrame {
fn payload_len(&self) -> usize {
self.payload.len()
}
}
impl PushedFrame {
#[must_use]
pub const fn correlation_id(&self) -> u64 {
self.correlation_id
}
#[must_use]
pub fn payload(&self) -> &[u8] {
&self.payload
}
#[must_use]
pub fn into_payload(self) -> Vec<u8> {
self.payload
}
}
#[derive(Debug)]
pub struct PushClient {
writer: Arc<Mutex<TcpStream>>,
mailbox: Arc<Mailbox<PushedFrame>>,
reader: Option<JoinHandle<()>>,
ledger: Arc<FlushLedger>,
}
impl PushClient {
#[must_use]
pub const fn with_setup_deadline(address: &str, deadline: Duration) -> PendingPushConnect<'_> {
PendingPushConnect::new(address, deadline)
}
pub fn connect(address: &str) -> Result<Self, SdkError> {
Self::connect_with_auth(address, &[])
}
pub fn connect_with_auth(address: &str, auth_token: &[u8]) -> Result<Self, SdkError> {
Self::connect_configured(address, auth_token, None, SETUP_TIMEOUT)
}
pub fn connect_with_registration(
address: &str,
registration: WorkerRegistration,
) -> Result<Self, SdkError> {
Self::connect_with_registration_and_auth(address, registration, &[])
}
pub fn connect_with_registration_and_auth(
address: &str,
registration: WorkerRegistration,
auth_token: &[u8],
) -> Result<Self, SdkError> {
Self::connect_configured(address, auth_token, Some(registration), SETUP_TIMEOUT)
}
fn connect_configured(
address: &str,
auth_token: &[u8],
registration: Option<WorkerRegistration>,
setup_deadline: Duration,
) -> Result<Self, SdkError> {
let mut stream = connect_socket(address, setup_deadline)?;
handshake(&mut stream, auth_token, setup_deadline)?;
if let Some(registration) = registration {
register(&mut stream, registration, setup_deadline)?;
}
Self::start_reader(stream)
}
fn start_reader(stream: TcpStream) -> Result<Self, SdkError> {
Self::start_reader_with(stream, MailboxLimits::PUSH)
}
fn start_reader_with(stream: TcpStream, limits: MailboxLimits) -> Result<Self, SdkError> {
stream
.set_read_timeout(None)
.map_err(|source| SdkError::Connection {
description: format!("failed to clear the push read deadline: {source}"),
})?;
let read_stream = stream.try_clone().map_err(|source| SdkError::Protocol {
description: format!("failed to clone push socket for reader thread: {source}"),
})?;
let mailbox = Arc::new(Mailbox::new(limits, PUSH_WORDS));
let reader_mailbox = Arc::clone(&mailbox);
let (ledger, verdicts) = FlushLedger::new();
let ledger = Arc::new(ledger);
let reader_ledger = Arc::clone(&ledger);
let reader = std::thread::Builder::new()
.name("liminal-push-reader".to_string())
.spawn(move || {
let end = run_reader(read_stream, &reader_mailbox, &verdicts, &reader_ledger);
reader_mailbox.close(end);
})
.map_err(|source| SdkError::Protocol {
description: format!("failed to start push reader thread: {source}"),
})?;
Ok(Self {
writer: Arc::new(Mutex::new(stream)),
mailbox,
reader: Some(reader),
ledger,
})
}
fn await_reader_exit(&self, budget: Duration) -> bool {
self.mailbox
.wait_until(budget, |_, end| end.is_some())
.unwrap_or(false)
}
fn shutdown_socket(&self, how: Shutdown) {
if let Ok(stream) = self.writer.lock() {
let _ = stream.shutdown(how);
}
}
pub fn recv_timeout(&self, timeout: Duration) -> Result<PushedFrame, SdkError> {
self.mailbox.take_push_timeout(timeout)
}
pub fn recv(&self) -> Result<PushReceiveOutcome, SdkError> {
receive(&self.mailbox, true)
}
pub fn receiver_handle(&self) -> Result<PushReceiver, SdkError> {
self.mailbox.claim_receiver()?;
Ok(PushReceiver {
mailbox: Arc::clone(&self.mailbox),
})
}
pub fn mailbox_stats(&self) -> Result<MailboxStats, SdkError> {
self.mailbox.stats()
}
pub fn reply(&self, correlation_id: u64, payload: Vec<u8>) -> Result<(), SdkError> {
let frame = Frame::new_push_reply(APPLICATION_STREAM_ID, correlation_id, payload)
.map_err(|error| protocol_error(&error))?;
let mut writer = self.writer.lock().map_err(|error| SdkError::Connection {
description: format!("push writer lock poisoned: {error}"),
})?;
write_frame(&mut writer, &frame)
}
pub fn refuse(&self, correlation_id: u64, code: &str, words: &str) -> Result<(), SdkError> {
self.writer_handle().refuse(correlation_id, code, words)
}
#[must_use]
pub fn writer_handle(&self) -> PushWriter {
PushWriter {
writer: Arc::clone(&self.writer),
ledger: Arc::clone(&self.ledger),
}
}
pub fn publish(&self, channel: &str, payload: Vec<u8>) -> Result<(), SdkError> {
self.writer_handle().publish(channel, payload)
}
pub fn flush(&self) -> Result<FlushOutcome, SdkError> {
let (failures, unresolved) = self.ledger.drain(FLUSH_BUDGET)?;
Ok(FlushOutcome::new(
failures,
unresolved,
FlushMode::VerdictOnly,
))
}
pub fn close(self) -> Result<FlushOutcome, SdkError> {
self.mailbox.abandon()?;
let (failures, unresolved) = self.ledger.drain(FLUSH_BUDGET)?;
let mode = if Arc::strong_count(&self.writer) == 1 {
FlushMode::FlushedAndHalfClosed
} else {
FlushMode::VerdictOnly
};
drop(self);
Ok(FlushOutcome::new(failures, unresolved, mode))
}
}
#[derive(Debug)]
pub enum PushReceiveOutcome {
Frame(PushedFrame),
Closed {
cause: PushCloseCause,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PushCloseCause {
OwnerClose,
PeerClose,
ReaderFailure {
error: SdkError,
},
}
#[derive(Debug)]
pub struct PushReceiver {
mailbox: Arc<Mailbox<PushedFrame>>,
}
impl PushReceiver {
pub fn recv(&self) -> Result<PushReceiveOutcome, SdkError> {
receive(&self.mailbox, false)
}
}
fn receive(mailbox: &Mailbox<PushedFrame>, direct: bool) -> Result<PushReceiveOutcome, SdkError> {
use crate::remote::mailbox::BlockingTake;
Ok(match mailbox.take_blocking(direct)? {
BlockingTake::Item(frame) => PushReceiveOutcome::Frame(frame),
BlockingTake::Abandoned => PushReceiveOutcome::Closed {
cause: PushCloseCause::OwnerClose,
},
BlockingTake::Ended(ReaderEnd::PeerClosed(_)) => PushReceiveOutcome::Closed {
cause: PushCloseCause::PeerClose,
},
BlockingTake::Ended(ReaderEnd::Failed(error)) => PushReceiveOutcome::Closed {
cause: PushCloseCause::ReaderFailure { error },
},
BlockingTake::Ended(ReaderEnd::Stopped) => PushReceiveOutcome::Closed {
cause: PushCloseCause::ReaderFailure {
error: SdkError::Connection {
description: "push reader stopped without a peer close".to_string(),
},
},
},
})
}
#[derive(Clone, Debug)]
pub struct PushWriter {
writer: Arc<Mutex<TcpStream>>,
ledger: Arc<FlushLedger>,
}
impl PushWriter {
pub fn refuse(&self, correlation_id: u64, code: &str, words: &str) -> Result<(), SdkError> {
let refusal = liminal::protocol::WorkerRefusal::new(code, words)
.map_err(|error| refusal_error(&error))?;
let frame = Frame::new_push_refusal(APPLICATION_STREAM_ID, correlation_id, &refusal)
.map_err(|error| refusal_error(&error))?;
let mut writer = self.writer.lock().map_err(|_| SdkError::Connection {
description: "push writer lock poisoned".to_string(),
})?;
write_frame(&mut writer, &frame)
}
pub fn publish(&self, channel: &str, payload: Vec<u8>) -> Result<(), SdkError> {
let envelope = MessageEnvelope::new(
SchemaId::new([0_u8; SchemaId::WIRE_LEN]),
CausalContext::independent(),
payload,
);
let frame = Frame::new_publish(APPLICATION_STREAM_ID, channel, envelope)
.map_err(|error| protocol_error(&error))?;
let mut writer = self.writer.lock().map_err(|error| SdkError::Connection {
description: format!("push writer lock poisoned: {error}"),
})?;
write_frame(&mut writer, &frame)?;
if channel != OBSERVABILITY_CHANNEL {
self.ledger.record_written();
}
Ok(())
}
pub fn reply(&self, correlation_id: u64, payload: Vec<u8>) -> Result<(), SdkError> {
let frame = Frame::new_push_reply(APPLICATION_STREAM_ID, correlation_id, payload)
.map_err(|error| protocol_error(&error))?;
let mut writer = self.writer.lock().map_err(|error| SdkError::Connection {
description: format!("push writer lock poisoned: {error}"),
})?;
write_frame(&mut writer, &frame)
}
}
impl Drop for PushClient {
fn drop(&mut self) {
let sole_owner = Arc::strong_count(&self.writer) == 1;
if self.mailbox.abandon().is_err() {
self.shutdown_socket(Shutdown::Both);
}
if sole_owner {
self.shutdown_socket(Shutdown::Write);
if !self.await_reader_exit(DROP_DRAIN_BUDGET) {
self.shutdown_socket(Shutdown::Both);
}
} else {
self.shutdown_socket(Shutdown::Read);
}
if let Some(reader) = self.reader.take() {
reader.join().ok();
}
}
}
fn connect_socket(address: &str, setup_deadline: Duration) -> Result<TcpStream, SdkError> {
let stream = TcpStream::connect(address).map_err(|source| SdkError::Connection {
description: format!("failed to connect push client to {address}: {source}"),
})?;
stream
.set_nodelay(true)
.map_err(|source| SdkError::Connection {
description: format!("failed to disable Nagle for {address}: {source}"),
})?;
stream
.set_read_timeout(Some(setup_deadline))
.map_err(|source| SdkError::Connection {
description: format!("failed to set the push setup deadline for {address}: {source}"),
})?;
stream
.set_write_timeout(Some(WRITE_TIMEOUT))
.map_err(|source| SdkError::Connection {
description: format!("failed to set push write timeout for {address}: {source}"),
})?;
Ok(stream)
}
pub(super) fn register(
stream: &mut TcpStream,
registration: WorkerRegistration,
setup_deadline: Duration,
) -> Result<(), SdkError> {
let frame = Frame::WorkerRegister {
flags: 0,
registration,
};
write_frame(stream, &frame)?;
let mut buffer = Vec::new();
match read_one_frame(stream, &mut buffer, setup_deadline)? {
Frame::WorkerRegisterAck {
outcome: WorkerRegisterOutcome::Accepted,
..
} => Ok(()),
Frame::WorkerRegisterAck {
outcome: WorkerRegisterOutcome::Rejected { reason },
..
} => Err(SdkError::Protocol {
description: format!("server rejected worker registration: {reason}"),
}),
other => Err(SdkError::Protocol {
description: format!(
"expected WorkerRegisterAck during registration, received {:?}",
other.frame_type()
),
}),
}
}
fn handshake(
stream: &mut TcpStream,
auth_token: &[u8],
setup_deadline: Duration,
) -> Result<(), SdkError> {
let connect = Frame::Connect {
flags: 0,
min_version: CLIENT_MIN_VERSION,
max_version: CLIENT_MAX_VERSION,
auth_token: auth_token.to_vec(),
};
write_frame(stream, &connect)?;
let mut buffer = Vec::new();
match read_one_frame(stream, &mut buffer, setup_deadline)? {
Frame::ConnectAck { .. } => Ok(()),
Frame::ConnectError {
reason_code,
message,
..
} => Err(SdkError::Connection {
description: format!(
"server rejected push connection (reason {reason_code}): {}",
message.unwrap_or_else(|| "no detail".to_string())
),
}),
other => Err(SdkError::Protocol {
description: format!(
"expected ConnectAck during push handshake, received {:?}",
other.frame_type()
),
}),
}
}
fn run_reader(
mut stream: TcpStream,
mailbox: &Mailbox<PushedFrame>,
verdicts: &Sender<PublishVerdict>,
ledger: &FlushLedger,
) -> ReaderEnd {
let mut buffer = Vec::new();
loop {
match next_frame(&mut stream, &mut buffer) {
Ok(Some(Frame::Push {
correlation_id,
payload,
..
})) => {
let pushed = PushedFrame {
correlation_id,
payload,
};
if let Some(end) = offer_or_end(mailbox, pushed) {
return end;
}
}
Ok(Some(Frame::PublishAck { .. })) => {
if verdicts.send(PublishVerdict::Accepted).is_err() {
return ReaderEnd::Stopped;
}
ledger.record_arrival();
}
Ok(Some(Frame::PublishError {
reason_code,
message,
..
})) => {
let rejection = PublishRejection::new(reason_code, message);
if verdicts.send(PublishVerdict::Rejected(rejection)).is_err() {
return ReaderEnd::Stopped;
}
ledger.record_arrival();
}
Ok(None) => return ReaderEnd::PeerClosed(peer_closed()),
Ok(Some(_)) => {
if let Err(error) = mailbox.record_ignored_frame() {
return ReaderEnd::Failed(error);
}
}
Err(error) => return ReaderEnd::Failed(error),
}
}
}
fn next_frame(stream: &mut TcpStream, buffer: &mut Vec<u8>) -> Result<Option<Frame>, SdkError> {
loop {
match decode(buffer) {
Ok((frame, consumed)) => {
buffer.drain(..consumed);
return Ok(Some(frame));
}
Err(
ProtocolError::IncompleteHeader { .. } | ProtocolError::TruncatedPayload { .. },
) => match fill_buffer(stream, buffer)? {
FillOutcome::Read => {}
FillOutcome::Closed if buffer.is_empty() => return Ok(None),
FillOutcome::Closed => return Err(peer_closed()),
FillOutcome::TimedOut => {
return Err(SdkError::Connection {
description: "the push reader's steady-state socket reported a read \
deadline it should not carry"
.to_string(),
});
}
},
Err(error) => return Err(protocol_error(&error)),
}
}
}
fn read_one_frame(
stream: &mut TcpStream,
buffer: &mut Vec<u8>,
setup_deadline: Duration,
) -> Result<Frame, SdkError> {
let deadline = Instant::now() + setup_deadline;
loop {
match decode(buffer) {
Ok((frame, consumed)) => {
buffer.drain(..consumed);
return Ok(frame);
}
Err(
ProtocolError::IncompleteHeader { .. } | ProtocolError::TruncatedPayload { .. },
) => match fill_buffer(stream, buffer)? {
FillOutcome::Read => {}
FillOutcome::Closed => return Err(peer_closed()),
FillOutcome::TimedOut => {
if Instant::now() >= deadline {
return Err(SdkError::Connection {
description:
"push connection timed out waiting for a control-frame reply"
.to_string(),
});
}
}
},
Err(error) => return Err(protocol_error(&error)),
}
}
}
fn fill_buffer(stream: &mut TcpStream, buffer: &mut Vec<u8>) -> Result<FillOutcome, SdkError> {
if buffer.len() > MAX_FRAME_BYTES {
return Err(SdkError::Protocol {
description: format!(
"push frame exceeded {MAX_FRAME_BYTES} bytes without a complete frame"
),
});
}
let mut chunk = [0_u8; READ_CHUNK_BYTES];
match stream.read(&mut chunk) {
Ok(0) => Ok(FillOutcome::Closed),
Ok(read) => {
let Some(received) = chunk.get(..read) else {
return Err(SdkError::Protocol {
description: "push socket read reported more bytes than the buffer holds"
.to_string(),
});
};
buffer.extend_from_slice(received);
Ok(FillOutcome::Read)
}
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
) =>
{
Ok(FillOutcome::TimedOut)
}
Err(error) => Err(SdkError::Connection {
description: format!("failed to read from push connection: {error}"),
}),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FillOutcome {
Closed,
Read,
TimedOut,
}
fn peer_closed() -> SdkError {
SdkError::Connection {
description: "server closed the push connection".to_string(),
}
}
fn write_frame(stream: &mut TcpStream, frame: &Frame) -> Result<(), SdkError> {
let len = encoded_len(frame).map_err(|error| protocol_error(&error))?;
let mut bytes = vec![0_u8; len];
let written = encode(frame, &mut bytes).map_err(|error| protocol_error(&error))?;
let encoded = bytes.get(..written).ok_or_else(|| SdkError::Protocol {
description: "push wire encoder reported an invalid byte count".to_string(),
})?;
stream
.write_all(encoded)
.map_err(|source| SdkError::Connection {
description: format!("failed to write push frame: {source}"),
})?;
stream.flush().map_err(|source| SdkError::Connection {
description: format!("failed to flush push frame: {source}"),
})
}
fn protocol_error(error: &ProtocolError) -> SdkError {
SdkError::Protocol {
description: format!("push wire codec error: {error}"),
}
}
fn refusal_error(error: &ProtocolError) -> SdkError {
SdkError::Protocol {
description: error
.message()
.unwrap_or("worker_refusal_protocol_fault")
.to_string(),
}
}
#[cfg(test)]
#[path = "push_refusal_tests.rs"]
mod refusal_tests;
#[cfg(test)]
#[path = "push_client_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "push_client_mailbox_tests.rs"]
mod mailbox_tests;
#[cfg(test)]
#[path = "push_client_receive_tests.rs"]
mod receive_tests;