use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use async_trait::async_trait;
use camel_api::{CamelError, Lifecycle, MetricsCollector};
const DRAIN_ZERO_SAMPLES_REQUIRED: u32 = 10;
#[derive(Default)]
struct BatchDepthProbeState {
zero_run: u32,
}
pub(crate) struct BatchDepthProbe {
expected: HashSet<String>,
queues: Mutex<HashMap<String, BatchDepthProbeState>>,
}
impl BatchDepthProbe {
pub(crate) fn new(expected: HashSet<String>) -> Self {
Self {
expected,
queues: Mutex::new(HashMap::new()),
}
}
pub(crate) fn reset(&self) {
for state in self
.queues
.lock()
.unwrap_or_else(|e| e.into_inner())
.values_mut()
{
state.zero_run = 0;
}
}
pub(crate) fn all_drained(&self) -> bool {
let queues = self.queues.lock().unwrap_or_else(|e| e.into_inner());
self.expected.iter().all(|label| {
queues
.get(label)
.is_some_and(|state| state.zero_run >= DRAIN_ZERO_SAMPLES_REQUIRED)
})
}
}
impl MetricsCollector for BatchDepthProbe {
fn record_exchange_duration(&self, _route_id: &str, _duration: Duration) {}
fn increment_errors(&self, _route_id: &str, _error_type: &str) {}
fn increment_exchanges(&self, _route_id: &str) {}
fn record_circuit_breaker_change(&self, _route_id: &str, _from: &str, _to: &str) {}
fn set_queue_depth(&self, queue: &str, depth: usize) {
if !self.expected.contains(queue) {
return;
}
let mut queues = self.queues.lock().unwrap_or_else(|e| e.into_inner());
let state = queues.entry(queue.to_string()).or_default();
if depth == 0 {
state.zero_run += 1;
} else {
state.zero_run = 0;
}
}
}
pub(crate) struct BatchProbeLifecycle(pub(crate) Arc<BatchDepthProbe>);
#[async_trait]
impl Lifecycle for BatchProbeLifecycle {
fn name(&self) -> &str {
"batch-depth-probe"
}
async fn start(&mut self) -> Result<(), CamelError> {
Ok(())
}
async fn stop(&mut self) -> Result<(), CamelError> {
Ok(())
}
fn as_metrics_collector(&self) -> Option<Arc<dyn MetricsCollector>> {
Some(Arc::clone(&self.0) as Arc<dyn MetricsCollector>)
}
}
pub(crate) async fn drain_until_empty(
probe: &BatchDepthProbe,
deadline: tokio::time::Instant,
) -> bool {
loop {
if probe.all_drained() {
return true;
}
let now = tokio::time::Instant::now();
let nap = deadline
.saturating_duration_since(now)
.min(Duration::from_millis(100));
if nap.is_zero() {
return false;
}
tokio::time::sleep(nap).await;
}
}