From f40a7ac4e289dadc07cc46346ff109d2bd5c3cf5 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Mon, 19 Jun 2023 08:48:38 -0700
Subject: [PATCH] fix: negative connection counters

This was really due to some stray decrements left over from an earlier
refactor; the newer code tracks the counts using RAII so the manual
decrements should not be needed any more.

To make the counters consistent when using lua connections, introduce
the same connection wrapping container over there, plus some glue
to make it possible to move the inner connection out of the wrapper
to destroy it when the connection is closed.

closes: #66
---
 crates/kumod/src/delivery_metrics.rs | 19 ++++++++++++++-----
 crates/kumod/src/lua_deliver.rs      |  9 ++++++---
 crates/kumod/src/ready_queue.rs      |  2 --
 3 files changed, 20 insertions(+), 10 deletions(-)

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<T> {
-    client: T,
+    client: Option<T>,
     metrics: DeliveryMetrics,
     armed: bool,
 }
@@ -72,11 +72,20 @@ impl<T> MetricsWrappedConnection<T> {
     pub fn map_connection<O>(mut self, client: O) -> MetricsWrappedConnection<O> {
         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<T> Drop for MetricsWrappedConnection<T> {
@@ -91,12 +100,12 @@ impl<T> Drop for MetricsWrappedConnection<T> {
 impl<T> std::ops::Deref for MetricsWrappedConnection<T> {
     type Target = T;
     fn deref(&self) -> &T {
-        &self.client
+        self.client.as_ref().expect("to be valid")
     }
 }
 
 impl<T> std::ops::DerefMut for MetricsWrappedConnection<T> {
     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<RegistryKey>,
+    connection: Option<MetricsWrappedConnection<RegistryKey>>,
     peer_address: ResolvedAddress,
 }
 
@@ -49,7 +50,7 @@ impl LuaQueueDispatcher {
 impl QueueDispatcher for LuaQueueDispatcher {
     async fn close_connection(&mut self, _dispatcher: &mut Dispatcher) -> anyhow::Result<bool> {
         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