use std::io::{self, Read, Write};
use weir_core::{
Durability, Envelope, HEADER_LEN, Header, MAX_ACK_BATCH_PAYLOAD_LEN,
MAX_BATCH_RECORDS_HARD_CAP, MAX_PAYLOAD_HARD_CAP, MAX_TRACKED_ACK_PAYLOAD_LEN, MessageType,
NackReason, RecordCoordinate, decode_ack_batch, encode_batch_body,
};
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct BatchOutcome {
pub accepted: Vec<bool>,
}
impl BatchOutcome {
#[must_use]
pub fn all_accepted(&self) -> bool {
self.accepted.iter().all(|&ok| ok)
}
#[must_use]
pub fn rejected_indices(&self) -> Vec<usize> {
self.accepted
.iter()
.enumerate()
.filter_map(|(i, &ok)| (!ok).then_some(i))
.collect()
}
}
const MAX_RESPONSE_PAYLOAD_LEN: usize = 2;
fn max_response_payload_len(message_type: MessageType) -> usize {
match message_type {
MessageType::AckTracked => MAX_TRACKED_ACK_PAYLOAD_LEN,
MessageType::AckBatch => MAX_ACK_BATCH_PAYLOAD_LEN,
_ => MAX_RESPONSE_PAYLOAD_LEN,
}
}
#[cfg(unix)]
pub type DefaultTransport = std::os::unix::net::UnixStream;
#[cfg(not(unix))]
pub type DefaultTransport = std::net::TcpStream;
#[derive(Debug)]
#[non_exhaustive]
pub enum ClientError {
Io(io::Error),
Nack(NackReason),
VersionMismatch {
daemon_version: u8,
},
UnknownNack(u8),
PayloadTooLarge {
len: usize,
limit: usize,
},
Protocol(String),
NoDefaultDurability,
EmptyPayload,
}
impl ClientError {
#[must_use]
pub fn is_recoverable(&self) -> bool {
match self {
Self::PayloadTooLarge { .. } | Self::NoDefaultDurability | Self::EmptyPayload => true,
Self::Nack(NackReason::InternalError) => true,
Self::Nack(_) => false,
Self::VersionMismatch { .. } | Self::UnknownNack(_) => false,
Self::Io(_) | Self::Protocol(_) => false,
#[allow(unreachable_patterns)]
_ => false,
}
}
}
impl std::fmt::Display for ClientError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "socket I/O error: {e}"),
Self::Nack(r) => write!(f, "server nack: {r:?}"),
Self::VersionMismatch { daemon_version } => write!(
f,
"wire version mismatch: daemon speaks v{daemon_version}, this client speaks v{}",
weir_core::WIRE_VERSION
),
Self::UnknownNack(b) => write!(f, "server nack with unknown reason {b:#04x}"),
Self::PayloadTooLarge { len, limit } => write!(
f,
"payload too large: {len} bytes exceeds the {limit}-byte protocol hard cap"
),
Self::Protocol(msg) => write!(f, "protocol error: {msg}"),
Self::NoDefaultDurability => {
write!(f, "push_default called but no default durability was set")
}
Self::EmptyPayload => write!(
f,
"payload is empty; rejected locally (an empty payload is the WAB end-of-records sentinel)"
),
}
}
}
impl std::error::Error for ClientError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => Some(e),
_ => None,
}
}
}
impl From<io::Error> for ClientError {
fn from(e: io::Error) -> Self {
Self::Io(e)
}
}
pub struct WeirClient<S = DefaultTransport> {
pub(crate) stream: S,
pub(crate) default_durability: Option<Durability>,
pub(crate) poisoned: bool,
pub(crate) closed_after_nack: bool,
}
impl<S> std::fmt::Debug for WeirClient<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WeirClient")
.field("default_durability", &self.default_durability)
.field("poisoned", &self.poisoned)
.field("closed_after_nack", &self.closed_after_nack)
.finish_non_exhaustive()
}
}
impl<S: Read + Write> WeirClient<S> {
pub fn set_default_durability(&mut self, durability: Durability) {
self.default_durability = Some(durability);
}
#[must_use = "a dropped push result hides whether the record was acked; an \
unhandled Nack also closes the connection"]
pub fn push(
&mut self,
payload: impl AsRef<[u8]>,
durability: Durability,
) -> Result<(), ClientError> {
let resp = self.send_record(MessageType::Push, payload, durability)?;
match resp.header().message_type() {
MessageType::Ack => Ok(()),
MessageType::Nack => Err(self.surface_nack(resp.payload())),
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"expected Ack or Nack, got {other:?}"
)))
}
}
}
#[must_use = "a dropped push result hides whether the record was acked; an \
unhandled Nack also closes the connection"]
pub fn push_tracked(
&mut self,
payload: impl AsRef<[u8]>,
durability: Durability,
) -> Result<RecordCoordinate, ClientError> {
let resp = self.send_record(MessageType::PushTracked, payload, durability)?;
match resp.header().message_type() {
MessageType::AckTracked => {
RecordCoordinate::decode(resp.payload()).map_err(|e| {
self.poisoned = true;
ClientError::Protocol(format!("malformed record coordinate: {e}"))
})
}
MessageType::Nack => Err(self.surface_nack(resp.payload())),
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"expected AckTracked or Nack, got {other:?}"
)))
}
}
}
fn send_record(
&mut self,
message_type: MessageType,
payload: impl AsRef<[u8]>,
durability: Durability,
) -> Result<Envelope, ClientError> {
self.ensure_usable()?;
let bytes = payload.as_ref();
if bytes.is_empty() {
return Err(ClientError::EmptyPayload);
}
if bytes.len() > MAX_PAYLOAD_HARD_CAP {
return Err(ClientError::PayloadTooLarge {
len: bytes.len(),
limit: MAX_PAYLOAD_HARD_CAP,
});
}
let payload = weir_core::Payload::copy_from_slice(bytes);
let header = Header::new(message_type, durability, 0);
let frame = Envelope::new(header, payload).encode();
self.send_frame(&frame)
}
fn send_frame(&mut self, frame: &[u8]) -> Result<Envelope, ClientError> {
if let Err(write_err) = self.stream.write_all(frame) {
if let Ok(resp) = self.read_response()
&& resp.header().message_type() == MessageType::Nack
{
self.closed_after_nack = true;
return Err(nack_error(resp.payload()));
}
self.poisoned = true;
return Err(write_err.into());
}
self.read_response()
}
fn note_nack(&mut self, err: ClientError) -> ClientError {
if !matches!(err, ClientError::Nack(NackReason::InternalError)) {
self.closed_after_nack = true;
}
err
}
fn surface_nack(&mut self, payload: &[u8]) -> ClientError {
match nack_error(payload) {
ClientError::Protocol(msg) => {
self.poisoned = true;
ClientError::Protocol(msg)
}
other => self.note_nack(other),
}
}
pub fn push_default(&mut self, payload: impl AsRef<[u8]>) -> Result<(), ClientError> {
let d = self
.default_durability
.ok_or(ClientError::NoDefaultDurability)?;
self.push(payload, d)
}
#[must_use = "a batch can partially succeed; inspect `accepted` or call \
`all_accepted()` — discarding it hides which records need retrying"]
pub fn push_batch(
&mut self,
records: &[impl AsRef<[u8]>],
durability: Durability,
) -> Result<BatchOutcome, ClientError> {
self.ensure_usable()?;
if records.is_empty() {
return Err(ClientError::EmptyPayload);
}
if records.len() > MAX_BATCH_RECORDS_HARD_CAP {
return Err(ClientError::PayloadTooLarge {
len: records.len(),
limit: MAX_BATCH_RECORDS_HARD_CAP,
});
}
let slices: Vec<&[u8]> = records.iter().map(AsRef::as_ref).collect();
for s in &slices {
if s.is_empty() {
return Err(ClientError::EmptyPayload);
}
if s.len() > MAX_PAYLOAD_HARD_CAP {
return Err(ClientError::PayloadTooLarge {
len: s.len(),
limit: MAX_PAYLOAD_HARD_CAP,
});
}
}
let body = encode_batch_body(&slices);
let header = Header::new(MessageType::PushBatch, durability, 0);
let frame = Envelope::new(header, body).encode();
let resp = self.send_frame(&frame)?;
match resp.header().message_type() {
MessageType::AckBatch => {
let accepted = decode_ack_batch(resp.payload(), slices.len()).map_err(|e| {
self.poisoned = true;
ClientError::Protocol(format!("malformed AckBatch payload: {e}"))
})?;
Ok(BatchOutcome { accepted })
}
MessageType::Ack => {
self.poisoned = true;
Err(ClientError::Protocol(
"daemon answered a PushBatch with a bare Ack; the records may be \
durable but their individual outcomes are unknown"
.to_string(),
))
}
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"unexpected response to a PushBatch: {other:?}"
)))
}
}
}
pub fn push_tracked_default(
&mut self,
payload: impl AsRef<[u8]>,
) -> Result<RecordCoordinate, ClientError> {
let d = self
.default_durability
.ok_or(ClientError::NoDefaultDurability)?;
self.push_tracked(payload, d)
}
pub fn health_check(&mut self) -> Result<(), ClientError> {
self.ensure_usable()?;
let header = Header::new(MessageType::HealthCheck, Durability::Durable, 0);
let frame = Envelope::new(header, vec![]).encode();
if let Err(e) = self.stream.write_all(&frame) {
self.poisoned = true;
return Err(e.into());
}
let resp = self.read_response()?;
match resp.header().message_type() {
MessageType::HealthCheckResponse => Ok(()),
MessageType::Nack => Err(self.surface_nack(resp.payload())),
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"expected HealthCheckResponse, got {other:?}"
)))
}
}
}
#[must_use]
pub fn is_poisoned(&self) -> bool {
self.poisoned || self.closed_after_nack
}
pub(crate) fn ensure_usable(&self) -> Result<(), ClientError> {
if self.closed_after_nack {
return Err(ClientError::Protocol(
"connection closed by the daemon after a Nack; \
drop this client and reconnect before sending again"
.into(),
));
}
if self.poisoned {
return Err(ClientError::Protocol(
"client connection poisoned by a prior response-read error/timeout; \
drop this client and reconnect before sending again"
.into(),
));
}
Ok(())
}
fn read_response(&mut self) -> Result<Envelope, ClientError> {
let result = self.read_response_inner();
if result.is_err() {
self.poisoned = true;
}
result
}
fn read_response_inner(&mut self) -> Result<Envelope, ClientError> {
let mut header_buf = [0u8; HEADER_LEN];
self.stream.read_exact(&mut header_buf)?;
let header =
Header::decode(&header_buf).map_err(|e| ClientError::Protocol(e.to_string()))?;
let payload_len = header.payload_len() as usize;
let cap = max_response_payload_len(header.message_type());
if payload_len > cap {
return Err(ClientError::Protocol(format!(
"response payload_len {payload_len} exceeds the {cap}-byte cap for \
{:?}; refusing to allocate",
header.message_type()
)));
}
let mut payload_buf = vec![0u8; payload_len];
if payload_len > 0 {
self.stream.read_exact(&mut payload_buf)?;
}
let mut crc_buf = [0u8; 4];
self.stream.read_exact(&mut crc_buf)?;
let expected = u32::from_le_bytes(crc_buf);
let computed = crc32fast::hash(&payload_buf);
if expected != computed {
return Err(ClientError::Protocol(format!(
"response payload CRC mismatch: expected {expected:#010x}, computed {computed:#010x}"
)));
}
Ok(Envelope::new(header, payload_buf))
}
#[cfg_attr(not(feature = "tls"), allow(dead_code))]
pub(crate) fn from_parts(stream: S, default_durability: Option<Durability>) -> Self {
Self {
stream,
default_durability,
poisoned: false,
closed_after_nack: false,
}
}
}
pub(crate) fn nack_error(payload: &[u8]) -> ClientError {
match payload.first().copied() {
Some(b) => match NackReason::try_from(b) {
Ok(NackReason::VersionMismatch) => match payload.get(1).copied() {
Some(daemon_version) => ClientError::VersionMismatch { daemon_version },
None => ClientError::Nack(NackReason::VersionMismatch),
},
Ok(r) => ClientError::Nack(r),
Err(_) => ClientError::UnknownNack(b),
},
None => ClientError::Protocol("Nack frame had empty payload".into()),
}
}