use super::pipeline::{MessageContext, MessageMiddleware};
use crate::common::error::Result;
use crate::common::platform::wall_clock_ms;
use crate::common::protocol::Frame;
use async_trait::async_trait;
use std::sync::Arc;
use tracing::{debug, info, warn};
pub struct LoggingMiddleware {
name: String,
log_level: LogLevel,
}
#[derive(Clone, Copy)]
pub enum LogLevel {
Debug,
Info,
Warn,
}
impl LoggingMiddleware {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
log_level: LogLevel::Info,
}
}
pub fn with_level(mut self, level: LogLevel) -> Self {
self.log_level = level;
self
}
}
#[async_trait]
impl MessageMiddleware for LoggingMiddleware {
async fn before(&self, ctx: &MessageContext) -> Result<Option<Frame>> {
match self.log_level {
LogLevel::Debug => {
debug!(
connection_id = ?ctx.connection_id,
message_id = %ctx.frame.message_id,
"Processing message"
);
}
LogLevel::Info => {
info!(
connection_id = ?ctx.connection_id,
message_id = %ctx.frame.message_id,
"Processing message"
);
}
LogLevel::Warn => {
warn!(
connection_id = ?ctx.connection_id,
message_id = %ctx.frame.message_id,
"Processing message"
);
}
}
Ok(None)
}
fn name(&self) -> &str {
&self.name
}
fn priority(&self) -> u32 {
10 }
}
pub struct MetricsMiddleware {
name: String,
}
impl MetricsMiddleware {
pub fn new(name: impl Into<String>) -> Self {
Self { name: name.into() }
}
}
#[async_trait]
impl MessageMiddleware for MetricsMiddleware {
async fn before(&self, ctx: &MessageContext) -> Result<Option<Frame>> {
ctx.set_metadata(
"start_time".to_string(),
wall_clock_ms().to_le_bytes().to_vec(),
)
.await;
Ok(None)
}
async fn after(&self, ctx: &MessageContext, response: Option<Frame>) -> Result<Option<Frame>> {
if let Some(start_bytes) = ctx.get_metadata("start_time").await {
let start_ms = u64::from_le_bytes(start_bytes.try_into().unwrap_or([0; 8]));
let duration_ms = wall_clock_ms().saturating_sub(start_ms);
debug!(
connection_id = ?ctx.connection_id,
message_id = %ctx.frame.message_id,
duration_ms,
"Message processed"
);
}
Ok(response)
}
fn name(&self) -> &str {
&self.name
}
fn priority(&self) -> u32 {
20
}
}
pub struct ValidationMiddleware {
name: String,
#[allow(clippy::type_complexity)]
validator: Arc<dyn Fn(&Frame) -> Result<()> + Send + Sync>,
}
impl ValidationMiddleware {
pub fn new<F>(name: impl Into<String>, validator: F) -> Self
where
F: Fn(&Frame) -> Result<()> + Send + Sync + 'static,
{
Self {
name: name.into(),
validator: Arc::new(validator),
}
}
}
#[async_trait]
impl MessageMiddleware for ValidationMiddleware {
async fn before(&self, ctx: &MessageContext) -> Result<Option<Frame>> {
(self.validator)(&ctx.frame)?;
Ok(None)
}
fn name(&self) -> &str {
&self.name
}
fn priority(&self) -> u32 {
5 }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::common::MessageParser;
use crate::common::platform::wall_clock_ms;
use crate::common::protocol::{FrameBuilder, Reliability, ping};
#[tokio::test]
async fn metrics_middleware_records_absolute_start_time_millis() {
let frame = FrameBuilder::new()
.with_command(crate::common::protocol::Command {
r#type: Some(
crate::common::protocol::flare::core::commands::command::Type::System(ping()),
),
})
.with_reliability(Reliability::BestEffort)
.build();
let ctx = MessageContext::new(frame, Some("conn-1".to_string()), MessageParser::json());
let metrics = MetricsMiddleware::new("metrics");
let before = wall_clock_ms();
metrics.before(&ctx).await.expect("before should succeed");
let after = wall_clock_ms();
let start_bytes = ctx
.get_metadata("start_time")
.await
.expect("metrics should store start_time");
let start_ms = u64::from_le_bytes(
start_bytes
.try_into()
.expect("start_time should be a u64 millis value"),
);
assert!(
(before..=after).contains(&start_ms),
"start_time should be an absolute wall-clock millis timestamp, got {start_ms}, expected within {before}..={after}"
);
}
}