use std::sync::OnceLock;
use tokio::sync::{Mutex, MutexGuard};
static HARN_STATE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
const LEAKY_STATE_ENV_VARS: &[&str] = &[
harn_vm::runtime_paths::HARN_STATE_DIR_ENV,
harn_vm::runtime_paths::HARN_RUN_DIR_ENV,
harn_vm::runtime_paths::HARN_WORKTREE_DIR_ENV,
harn_vm::event_log::HARN_EVENT_LOG_BACKEND_ENV,
harn_vm::event_log::HARN_EVENT_LOG_DIR_ENV,
harn_vm::event_log::HARN_EVENT_LOG_SQLITE_PATH_ENV,
harn_vm::event_log::HARN_EVENT_LOG_QUEUE_DEPTH_ENV,
"HARN_MCP_OAUTH_AUTHORIZATION_SERVERS",
"HARN_MCP_OAUTH_INTROSPECTION_URL",
"HARN_MCP_OAUTH_RESOURCE",
"HARN_MCP_OAUTH_AUDIENCE",
"HARN_MCP_OAUTH_SCOPES",
];
fn clear_leaky_state_env() {
for name in LEAKY_STATE_ENV_VARS {
std::env::remove_var(name);
}
}
fn state_mutex() -> &'static Mutex<()> {
HARN_STATE_LOCK.get_or_init(|| Mutex::new(()))
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Holder {
Task(tokio::task::Id),
Thread(std::thread::ThreadId),
}
static HOLDER: std::sync::Mutex<Option<Holder>> = std::sync::Mutex::new(None);
fn current_holder() -> Holder {
match tokio::task::try_id() {
Some(id) => Holder::Task(id),
None => Holder::Thread(std::thread::current().id()),
}
}
fn reject_reentrant_acquire() {
let holder = current_holder();
let already_held = *HOLDER.lock().expect("harn-state holder") == Some(holder);
assert!(
!already_held,
"this test already holds the harn-state lock; acquiring it again would deadlock. \
There is exactly one lock over the process environment — take it once, at the top \
of the test, and pass the guard down if an inner helper needs it."
);
}
pub struct HarnStateGuard {
_inner: MutexGuard<'static, ()>,
}
impl Drop for HarnStateGuard {
fn drop(&mut self) {
*HOLDER.lock().expect("harn-state holder") = None;
}
}
fn finish_acquire(inner: MutexGuard<'static, ()>) -> HarnStateGuard {
*HOLDER.lock().expect("harn-state holder") = Some(current_holder());
clear_leaky_state_env();
HarnStateGuard { _inner: inner }
}
pub fn lock_harn_state() -> HarnStateGuard {
reject_reentrant_acquire();
finish_acquire(state_mutex().blocking_lock())
}
pub async fn lock_harn_state_async() -> HarnStateGuard {
reject_reentrant_acquire();
finish_acquire(state_mutex().lock().await)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
#[should_panic(expected = "already holds the harn-state lock")]
async fn a_second_acquire_from_the_same_test_panics_instead_of_hanging() {
let _first = lock_harn_state_async().await;
let _second = lock_harn_state_async().await;
}
#[tokio::test]
async fn releasing_the_lock_lets_the_same_test_take_it_again() {
drop(lock_harn_state_async().await);
let _second = lock_harn_state_async().await;
}
}