diff --git a/crates/kumod/src/delivery_metrics.rs b/crates/kumod/src/delivery_metrics.rs index 3885dc77..301360ba 100644 --- a/crates/kumod/src/delivery_metrics.rs +++ b/crates/kumod/src/delivery_metrics.rs @@ -26,7 +26,7 @@ impl DeliveryMetrics { self.connection_total.inc(); self.global_connection_total.inc(); MetricsWrappedConnection { - client, + client: Some(client), metrics: self.clone(), armed: true, } @@ -62,7 +62,7 @@ impl DeliveryMetrics { /// and decrements them when dropped #[derive(Debug)] pub struct MetricsWrappedConnection { - client: T, + client: Option, metrics: DeliveryMetrics, armed: bool, } @@ -72,11 +72,20 @@ impl MetricsWrappedConnection { pub fn map_connection(mut self, client: O) -> MetricsWrappedConnection { self.armed = false; MetricsWrappedConnection { - client, + client: Some(client), metrics: self.metrics.clone(), armed: true, } } + + pub fn take(mut self) -> T { + if self.armed { + self.metrics.connection_gauge.dec(); + self.metrics.global_connection_gauge.dec(); + self.armed = false; + } + self.client.take().expect("to take only once") + } } impl Drop for MetricsWrappedConnection { @@ -91,12 +100,12 @@ impl Drop for MetricsWrappedConnection { impl std::ops::Deref for MetricsWrappedConnection { type Target = T; fn deref(&self) -> &T { - &self.client + self.client.as_ref().expect("to be valid") } } impl std::ops::DerefMut for MetricsWrappedConnection { fn deref_mut(&mut self) -> &mut T { - &mut self.client + self.client.as_mut().expect("to be valid") } } diff --git a/crates/kumod/src/lua_deliver.rs b/crates/kumod/src/lua_deliver.rs index f298f13d..278fea8e 100644 --- a/crates/kumod/src/lua_deliver.rs +++ b/crates/kumod/src/lua_deliver.rs @@ -1,3 +1,4 @@ +use crate::delivery_metrics::MetricsWrappedConnection; use crate::logging::{log_disposition, LogDisposition}; use crate::ready_queue::{Dispatcher, QueueDispatcher}; use crate::runtime::{rt_spawn, spawn}; @@ -25,7 +26,7 @@ pub struct LuaDeliveryProtocol { pub struct LuaQueueDispatcher { lua_config: LuaConfig, proto_config: LuaDeliveryProtocol, - connection: Option, + connection: Option>, peer_address: ResolvedAddress, } @@ -49,7 +50,7 @@ impl LuaQueueDispatcher { impl QueueDispatcher for LuaQueueDispatcher { async fn close_connection(&mut self, _dispatcher: &mut Dispatcher) -> anyhow::Result { if let Some(connection) = self.connection.take() { - self.lua_config.remove_registry_value(connection)?; + self.lua_config.remove_registry_value(connection.take())?; Ok(true) } else { Ok(false) @@ -57,6 +58,7 @@ impl QueueDispatcher for LuaQueueDispatcher { } async fn attempt_connection(&mut self, dispatcher: &mut Dispatcher) -> anyhow::Result<()> { + let connection_wrapper = dispatcher.metrics.wrap_connection(()); let components = QueueNameComponents::parse(&dispatcher.queue_name); let connection = self .lua_config @@ -66,7 +68,8 @@ impl QueueDispatcher for LuaQueueDispatcher { ) .await?; - self.connection.replace(connection); + self.connection + .replace(connection_wrapper.map_connection(connection)); Ok(()) } diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index dff0d0ab..d9bb61f4 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/crates/kumod/src/ready_queue.rs @@ -487,8 +487,6 @@ impl Dispatcher { } if let Err(err) = queue_dispatcher.attempt_connection(&mut dispatcher).await { connection_failures.push(format!("{err:#}")); - dispatcher.metrics.connection_gauge.dec(); - dispatcher.metrics.global_connection_gauge.dec(); if !queue_dispatcher .have_more_connection_candidates(&mut dispatcher) .await