use super::error::{TransportError, TransportResult};
use super::filter::FilteredDlqEntry;
use super::types::{Message, SendResult};
use super::work_batch::{Record, WorkBatch};
use std::fmt::{Debug, Display};
use std::future::Future;
pub trait CommitToken: Clone + Send + Sync + Debug + Display + 'static {
fn as_str(&self) -> String {
format!("{self}")
}
}
#[derive(Debug)]
pub struct RecvBatch<T: CommitToken> {
pub messages: Vec<Message<T>>,
pub dlq_entries: Vec<FilteredDlqEntry>,
pub filtered_tokens: Vec<T>,
}
impl<T: CommitToken> RecvBatch<T> {
#[must_use]
pub fn empty() -> Self {
Self {
messages: Vec::new(),
dlq_entries: Vec::new(),
filtered_tokens: Vec::new(),
}
}
#[must_use]
pub fn from_messages(messages: Vec<Message<T>>) -> Self {
Self {
messages,
dlq_entries: Vec::new(),
filtered_tokens: Vec::new(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.messages.is_empty() && self.filtered_tokens.is_empty()
}
#[must_use]
pub fn len(&self) -> usize {
self.messages.len()
}
}
pub trait TransportBase: Send + Sync {
fn close(&self) -> impl Future<Output = TransportResult<()>> + Send;
fn is_healthy(&self) -> bool;
fn name(&self) -> &'static str;
fn healthcheck(&self) -> impl Future<Output = TransportResult<()>> + Send {
async move {
if self.is_healthy() {
Ok(())
} else {
Err(TransportError::Connection(format!(
"{} transport failed boot healthcheck (not healthy)",
self.name()
)))
}
}
}
}
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
pub struct HealthcheckConfig {
#[serde(default = "default_healthcheck_enabled")]
pub enabled: bool,
#[serde(default = "default_healthcheck_timeout_ms")]
pub timeout_ms: u64,
}
fn default_healthcheck_enabled() -> bool {
true
}
fn default_healthcheck_timeout_ms() -> u64 {
5000
}
impl Default for HealthcheckConfig {
fn default() -> Self {
Self {
enabled: default_healthcheck_enabled(),
timeout_ms: default_healthcheck_timeout_ms(),
}
}
}
pub async fn boot_healthcheck<T: TransportBase>(
transport: &T,
cfg: HealthcheckConfig,
) -> TransportResult<()> {
if !cfg.enabled {
return Ok(());
}
let probe = transport.healthcheck();
match tokio::time::timeout(std::time::Duration::from_millis(cfg.timeout_ms), probe).await {
Ok(result) => result,
Err(_elapsed) => Err(TransportError::Connection(format!(
"{} transport boot healthcheck timed out after {}ms",
transport.name(),
cfg.timeout_ms
))),
}
}
pub trait TransportSender: TransportBase {
fn send(&self, key: &str, payload: bytes::Bytes) -> impl Future<Output = SendResult> + Send;
fn send_batch(&self, records: &[Record]) -> impl Future<Output = SendResult> + Send {
async move {
for record in records {
let key = record.key.as_deref().unwrap_or("");
match self.send(key, record.payload.clone()).await {
SendResult::Ok | SendResult::FilteredDlq => {}
other @ (SendResult::Backpressured | SendResult::Fatal(_)) => return other,
}
}
SendResult::Ok
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RecvLimits {
pub max_records: usize,
pub max_bytes: u64,
}
pub trait TransportReceiver: TransportBase {
type Token: CommitToken;
fn recv(
&self,
max: usize,
) -> impl Future<Output = TransportResult<WorkBatch<Self::Token>>> + Send;
fn recv_limited(
&self,
limits: RecvLimits,
) -> impl Future<Output = TransportResult<WorkBatch<Self::Token>>> + Send {
self.recv(limits.max_records)
}
fn commit(&self, tokens: &[Self::Token]) -> impl Future<Output = TransportResult<()>> + Send;
}
pub trait Transport: TransportSender + TransportReceiver {}
impl<T: TransportSender + TransportReceiver> Transport for T {}
pub trait FromCascade: Default + serde::Serialize + serde::de::DeserializeOwned + 'static {
#[must_use]
fn from_cascade_key(key: &str) -> Self {
#[cfg(feature = "config")]
{
if let Some(cfg) = crate::config::try_get()
&& let Ok(value) = cfg.unmarshal_key_registered::<Self>(key)
{
return value;
}
}
#[cfg(not(feature = "config"))]
let _ = key;
Self::default()
}
}
#[cfg(test)]
mod healthcheck_tests {
use super::*;
use std::time::Duration;
struct DefaultProbe {
healthy: bool,
}
impl TransportBase for DefaultProbe {
async fn close(&self) -> TransportResult<()> {
Ok(())
}
fn is_healthy(&self) -> bool {
self.healthy
}
fn name(&self) -> &'static str {
"default-probe"
}
}
struct ActiveProbe {
ok: bool,
hang: bool,
}
impl TransportBase for ActiveProbe {
async fn close(&self) -> TransportResult<()> {
Ok(())
}
fn is_healthy(&self) -> bool {
true
}
fn name(&self) -> &'static str {
"active-probe"
}
async fn healthcheck(&self) -> TransportResult<()> {
if self.hang {
tokio::time::sleep(Duration::from_secs(60)).await;
}
if self.ok {
Ok(())
} else {
Err(TransportError::Connection("active probe rejected".into()))
}
}
}
#[tokio::test]
async fn default_healthcheck_delegates_to_is_healthy() {
assert!(
boot_healthcheck(
&DefaultProbe { healthy: true },
HealthcheckConfig::default()
)
.await
.is_ok()
);
assert!(
boot_healthcheck(
&DefaultProbe { healthy: false },
HealthcheckConfig::default()
)
.await
.is_err(),
"default probe must fail-fast when not healthy"
);
}
#[tokio::test]
async fn disabled_skips_probe_even_if_unhealthy() {
let t = ActiveProbe {
ok: false,
hang: true,
};
let cfg = HealthcheckConfig {
enabled: false,
timeout_ms: 10,
};
assert!(
boot_healthcheck(&t, cfg).await.is_ok(),
"disabled must skip"
);
}
#[tokio::test]
async fn active_probe_failure_fails_fast() {
let t = ActiveProbe {
ok: false,
hang: false,
};
assert!(
boot_healthcheck(&t, HealthcheckConfig::default())
.await
.is_err()
);
}
#[tokio::test(start_paused = true)]
async fn hanging_probe_times_out() {
let t = ActiveProbe {
ok: true,
hang: true,
};
let cfg = HealthcheckConfig {
enabled: true,
timeout_ms: 50,
};
assert!(
boot_healthcheck(&t, cfg).await.is_err(),
"a hanging probe must hit the boot timeout"
);
}
}