Recover poisoned caches

akd/crates/akd_storage/src/vrf_key_database.rs

@@ -102,8 +102,18 @@ impl VrfKeyDatabase {
 #[async_trait]
 impl VRFKeyStorage for VrfKeyDatabase {
     async fn retrieve(&self) -> Result<Vec<u8>, VrfError> {
-        if let Some(cached_key) = self.cached_vrf_key.read().expect("cache poisoned").as_ref() {
-            return Ok(cached_key.clone());
+        // Handle poisoned lock by forcing a cache reload
+        let cached = match self.cached_vrf_key.read() {
+            Ok(guard) => guard.as_ref().cloned(),
+            Err(_poisoned) => {
+                error!("VRF key cache lock was poisoned, forcing reload from database");
+                // Don't trust the cached value - force a reload
+                None
+            }
+        };
+
+        if let Some(cached_key) = cached {
+            return Ok(cached_key);
         }
 
         let result = match &self.get_vrf_key().await {
@@ -137,9 +147,17 @@ impl VRFKeyStorage for VrfKeyDatabase {
         };
 
         if result.is_ok() {
-            // update cached key
-            let mut write_guard = self.cached_vrf_key.write().expect("cache poisoned");
-            *write_guard = Some(result.as_ref().unwrap().clone());
+            // Update cached key, clearing poison if present
+            match self.cached_vrf_key.write() {
+                Ok(mut guard) => {
+                    *guard = Some(result.as_ref().unwrap().clone());
+                }
+                Err(poisoned) => {
+                    error!("VRF key cache lock was poisoned during write, recovering by caching fresh value");
+                    // We have a valid value from the database - use it to fix the poisoned cache
+                    *poisoned.into_inner() = Some(result.as_ref().unwrap().clone());
+                }
+            }
         }
 
         result

akd/crates/ms_database/src/pool.rs

@@ -52,8 +52,16 @@ impl ConnectionManager {
 
     /// Mark the pool as unhealthy. This is used to indicate that a connection should be replaced.
     pub async fn set_unhealthy(&self) {
-        let mut healthy = self.is_healthy.write().expect("poisoned is_healthy lock");
-        *healthy = false;
+        match self.is_healthy.write() {
+            Ok(mut guard) => {
+                *guard = false;
+            }
+            Err(poisoned) => {
+                // Even if poisoned, we can safely set the boolean value
+                tracing::error!("Health lock was poisoned, recovering and setting unhealthy");
+                *poisoned.into_inner() = false;
+            }
+        }
     }
 }
 
@@ -149,6 +157,13 @@ impl ManageConnection for ConnectionManager {
     }
 
     fn has_broken(&self, _conn: &mut Self::Connection) -> bool {
-        *self.is_healthy.read().expect("poisoned is_healthy lock")
+        match self.is_healthy.read() {
+            Ok(guard) => *guard,
+            Err(poisoned) => {
+                // If poisoned, assume unhealthy (safe default)
+                tracing::error!("Health lock was poisoned during read, assuming unhealthy");
+                *poisoned.into_inner()
+            }
+        }
     }
 }