From a4fe44e2721602b52c8ea78c0c138a1ba522f686 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Sat, 10 Aug 2024 06:14:22 -0700
Subject: [PATCH] redis/throttle/limit: surface connectivity issues better

I recently noted that an issue relating to failing redis auth
was not being surfaced in the diagnostic log.

It was being logged at debug level so wasn't visible by default.

This commit does two things to improve this:

* Surface that error at error level, but keeps "regular" over-limit
  events at debug level
* Explicitly pings the configured redis server on startup, so that
  connectivity can be validated at `--validate` time and not be
  an issue that is deferred until after the service is up and running.
---
 crates/kumo-server-common/src/lib.rs |  3 ++-
 crates/kumod/src/ready_queue.rs      | 12 +++++++++++-
 crates/mod-redis/src/lib.rs          | 15 +++++++++++++--
 crates/throttle/src/limit.rs         |  8 ++++++--
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/crates/kumo-server-common/src/lib.rs b/crates/kumo-server-common/src/lib.rs
index cd22e677..2c1154bb 100644
--- a/crates/kumo-server-common/src/lib.rs
+++ b/crates/kumo-server-common/src/lib.rs
@@ -201,9 +201,10 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
 
     kumo_mod.set(
         "configure_redis_throttles",
-        lua.create_function(move |lua, params: Value| {
+        lua.create_async_function(|lua, params: Value| async move {
             let key: RedisConnKey = from_lua_value(lua, params)?;
             let conn = key.open().map_err(any_err)?;
+            conn.ping().await.map_err(any_err)?;
             throttle::use_redis(conn).map_err(any_err)
         })?,
     )?;
diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs
index 75c8217d..0a536150 100644
--- a/crates/kumod/src/ready_queue.rs
+++ b/crates/kumod/src/ready_queue.rs
@@ -703,13 +703,23 @@ impl ReadyQueue {
                             self.connections.lock().push(handle);
                         }
                     }
-                    Err(err) => {
+                    Err(err @ throttle::Error::TooManyLeases(_)) => {
+                        // Over budget; we'll try again later
                         tracing::debug!(
                             "maintain {}: could not acquire connection lease: {err:#}",
                             self.name
                         );
                         break;
                     }
+                    Err(err) => {
+                        // Some kind of error trying to acquire the lease, could be
+                        // a redis/connectivity error, let's surface it
+                        tracing::error!(
+                            "maintain {}: could not acquire connection lease: {err:#}",
+                            self.name
+                        );
+                        break;
+                    }
                 }
             }
         }
diff --git a/crates/mod-redis/src/lib.rs b/crates/mod-redis/src/lib.rs
index 98eaf9a0..9d19cdf6 100644
--- a/crates/mod-redis/src/lib.rs
+++ b/crates/mod-redis/src/lib.rs
@@ -39,8 +39,7 @@ impl Manager for ClientManager {
         conn: &mut Self::Type,
         _metrics: &Metrics,
     ) -> RecycleResult<anyhow::Error> {
-        redis::cmd("PING")
-            .query_async(&mut *conn)
+        conn.ping()
             .await
             .map_err(|err| RecycleError::message(format!("{err:#}")))
     }
@@ -50,6 +49,12 @@ impl Manager for ClientManager {
 pub struct RedisConnection(Arc<RedisConnKey>);
 
 impl RedisConnection {
+    pub async fn ping(&self) -> anyhow::Result<()> {
+        let pool = self.0.get_pool()?;
+        let mut conn = pool.get().await.map_err(|err| anyhow::anyhow!("{err:#}"))?;
+        conn.ping().await
+    }
+
     pub async fn query(&self, cmd: Cmd) -> anyhow::Result<RedisValue> {
         let pool = self.0.get_pool()?;
         let mut conn = pool.get().await.map_err(|err| anyhow::anyhow!("{err:#}"))?;
@@ -257,6 +262,12 @@ pub enum ConnectionWrapper {
     Cluster(ClusterConnection),
 }
 
+impl ConnectionWrapper {
+    pub async fn ping(&mut self) -> anyhow::Result<()> {
+        Ok(redis::cmd("PING").query_async(self).await?)
+    }
+}
+
 impl ConnectionLike for ConnectionWrapper {
     // Required methods
     fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, RedisValue> {
diff --git a/crates/throttle/src/limit.rs b/crates/throttle/src/limit.rs
index 2ce00e3f..4780794d 100644
--- a/crates/throttle/src/limit.rs
+++ b/crates/throttle/src/limit.rs
@@ -1,5 +1,5 @@
 use crate::{Error, REDIS};
-use anyhow::anyhow;
+use anyhow::{anyhow, Context};
 use mod_redis::{RedisConnection, Script};
 use once_cell::sync::{Lazy, OnceCell};
 use std::collections::HashMap;
@@ -88,7 +88,11 @@ impl LimitSpec {
             .arg(self.limit)
             .arg(uuid_str);
 
-        match conn.invoke_script(script).await? {
+        match conn
+            .invoke_script(script)
+            .await
+            .context("error invoking redis lease acquisition script")?
+        {
             mod_redis::RedisValue::Okay => {}
             mod_redis::RedisValue::Int(next_expiration_interval) => {
                 return Err(Error::TooManyLeases(Duration::from_secs(