use alloc::format;
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::time::Duration;
use std::net::{Shutdown, TcpStream};
use std::thread::JoinHandle;
use liminal::protocol::{
DeliveredId, Frame, MessageId, ProtocolError, SchemaId, WorkerRegistration,
};
use crate::SdkError;
use crate::remote::SETUP_TIMEOUT;
mod subscription_reader;
mod subscription_setup;
use subscription_reader::run_reader;
mod subscription_signals;
mod subscription_wire;
use crate::remote::mailbox::{Mailbox, MailboxItem, MailboxLimits, MailboxStats, MailboxWords};
use subscription_setup::{connect_socket, handshake, subscribe};
use subscription_wire::write_frame;
const SUBSCRIPTION_WORDS: MailboxWords = MailboxWords {
receive: "subscription receive failed",
item: "delivery",
reader: "subscription reader",
};
const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
const READ_CHUNK_BYTES: usize = 4096;
const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024;
const SUBSCRIPTION_STREAM_ID: u32 = 1;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DeliveredMessage {
delivery_seq: u64,
schema_id: SchemaId,
payload: Vec<u8>,
message: Option<DeliveredId>,
}
impl MailboxItem for DeliveredMessage {
fn payload_len(&self) -> usize {
self.payload.len()
}
}
impl DeliveredMessage {
#[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<DeliveredId> {
self.message
}
#[must_use]
pub fn reply_parent(&self) -> Option<MessageId> {
self.message.map(DeliveredId::reply_parent)
}
}
#[derive(Debug)]
pub struct SubscriptionStream {
writer: TcpStream,
subscription_id: u64,
mailbox: Arc<Mailbox<DeliveredMessage>>,
reader: Option<JoinHandle<()>>,
}
impl SubscriptionStream {
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> {
Self::open_configured(address, channel, accepted_schemas, auth_token, None, limits)
}
pub fn open_registered(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
registration: WorkerRegistration,
) -> Result<Self, SdkError> {
Self::open_configured(
address,
channel,
accepted_schemas,
&[],
Some(registration),
MailboxLimits::SUBSCRIPTION,
)
}
fn open_configured(
address: &str,
channel: &str,
accepted_schemas: Vec<SchemaId>,
auth_token: &[u8],
registration: Option<WorkerRegistration>,
limits: MailboxLimits,
) -> Result<Self, SdkError> {
let mut stream = connect_socket(address)?;
let mut buffer = Vec::new();
handshake(&mut stream, &mut buffer, auth_token)?;
if let Some(registration) = registration {
super::push_client::register(&mut stream, registration, SETUP_TIMEOUT)?;
}
let subscription_id = subscribe(&mut stream, &mut buffer, channel, accepted_schemas)?;
stream
.set_read_timeout(None)
.map_err(|source| SdkError::Connection {
description: format!("failed to clear the subscription read deadline: {source}"),
})?;
let read_stream = stream.try_clone().map_err(|source| SdkError::Protocol {
description: format!("failed to clone subscription socket for reader thread: {source}"),
})?;
Self::start(stream, read_stream, buffer, subscription_id, limits)
}
fn start(
stream: TcpStream,
read_stream: TcpStream,
buffer: Vec<u8>,
subscription_id: u64,
limits: MailboxLimits,
) -> Result<Self, SdkError> {
let mailbox = Arc::new(Mailbox::new(limits, SUBSCRIPTION_WORDS));
let reader_mailbox = Arc::clone(&mailbox);
let reader = std::thread::Builder::new()
.name("liminal-subscription-reader".to_string())
.spawn(move || {
let end = run_reader(read_stream, buffer, &reader_mailbox);
reader_mailbox.close(end);
})
.map_err(|source| SdkError::Protocol {
description: format!("failed to start subscription reader thread: {source}"),
})?;
Ok(Self {
writer: stream,
subscription_id,
mailbox,
reader: Some(reader),
})
}
pub fn recv_timeout(&self, timeout: Duration) -> Result<DeliveredMessage, 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
}
}
impl Drop for SubscriptionStream {
fn drop(&mut self) {
let unsubscribe = Frame::Unsubscribe {
flags: 0,
stream_id: SUBSCRIPTION_STREAM_ID,
subscription_id: self.subscription_id,
};
let _ = write_frame(&mut self.writer, &unsubscribe);
let _ = write_frame(&mut self.writer, &Frame::Disconnect { flags: 0 });
let _ = self.writer.shutdown(Shutdown::Both);
self.mailbox.release();
if let Some(reader) = self.reader.take() {
reader.join().ok();
}
}
}
#[cfg(test)]
use subscription_setup::CLIENT_MAX_VERSION;
#[cfg(test)]
use alloc::vec;
#[cfg(test)]
use liminal::protocol::{decode, encode, encoded_len};
#[cfg(test)]
use std::io::{Read, Write};
#[cfg(test)]
use subscription_wire::protocol_error;
#[cfg(test)]
#[path = "subscription_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "subscription_mailbox_tests.rs"]
mod mailbox_tests;