diff --git a/akd/crates/akd_storage/src/vrf_key_database.rs b/akd/crates/akd_storage/src/vrf_key_database.rs index 97cd2d7e6f..4c98503c19 100644 --- a/akd/crates/akd_storage/src/vrf_key_database.rs +++ b/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, 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 diff --git a/akd/crates/ms_database/src/pool.rs b/akd/crates/ms_database/src/pool.rs index 8daaa0e12a..17b0b8f13d 100644 --- a/akd/crates/ms_database/src/pool.rs +++ b/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() + } + } } }