use std::time::{Duration, Instant};
use anyhow::{anyhow, Context, Result};
use thiserror::Error;
use crate::registry::{db, now_epoch_s, pid_alive, Registry};
const POLL_INTERVAL: Duration = Duration::from_millis(100);
const RELEASE_SQL: &str = "DELETE FROM bidi_locks WHERE browser_name = ?1 AND holder_pid = ?2";
fn release(
conn: &rusqlite::Connection,
browser_name: &str,
holder_pid: u32,
) -> rusqlite::Result<usize> {
conn.execute(RELEASE_SQL, rusqlite::params![browser_name, holder_pid])
}
#[derive(Debug, Error)]
#[error("Firefox BiDi for {browser_name} is held by PID {holder_pid} (waited {waited_ms}ms)")]
pub struct BidiLockBusy {
pub browser_name: String,
pub holder_pid: u32,
pub waited_ms: u64,
}
#[derive(Debug)]
pub struct BidiLockGuard {
browser_name: String,
holder_pid: u32,
db_path: std::path::PathBuf,
}
impl BidiLockGuard {
pub fn browser_name(&self) -> &str {
&self.browser_name
}
pub fn holder_pid(&self) -> u32 {
self.holder_pid
}
}
impl Drop for BidiLockGuard {
fn drop(&mut self) {
if let Ok(conn) = rusqlite::Connection::open(&self.db_path) {
let _ = conn.busy_timeout(std::time::Duration::from_secs(5));
let _ = release(&conn, &self.browser_name, self.holder_pid);
}
}
}
impl Registry {
pub fn bidi_lock_acquire(
&self,
browser_name: &str,
timeout: Duration,
) -> Result<BidiLockGuard> {
let my_pid = std::process::id();
let start = Instant::now();
loop {
let now = now_epoch_s();
let attempt = self.conn.execute(
"INSERT INTO bidi_locks (browser_name, holder_pid, acquired_at_epoch_s) \
VALUES (?1, ?2, ?3)",
rusqlite::params![browser_name, my_pid, now],
);
match attempt {
Ok(_) => {
return Ok(BidiLockGuard {
browser_name: browser_name.to_string(),
holder_pid: my_pid,
db_path: self.db_path.clone(),
});
}
Err(rusqlite::Error::SqliteFailure(e, _))
if e.code == rusqlite::ErrorCode::ConstraintViolation =>
{
if let Some(existing) = self.bidi_lock_holder(browser_name)? {
if !pid_alive(existing.holder_pid) {
let _ = release(&self.conn, browser_name, existing.holder_pid);
continue;
}
if start.elapsed() >= timeout {
return Err(BidiLockBusy {
browser_name: browser_name.to_string(),
holder_pid: existing.holder_pid,
waited_ms: start.elapsed().as_millis() as u64,
}
.into());
}
std::thread::sleep(POLL_INTERVAL);
} else {
continue;
}
}
Err(e) => {
return Err(anyhow!(e))
.with_context(|| format!("acquire bidi_lock for {browser_name}"));
}
}
}
}
pub fn bidi_lock_holder(&self, browser_name: &str) -> Result<Option<BidiLockRow>> {
db::query_optional(
&self.conn,
"SELECT browser_name, holder_pid, acquired_at_epoch_s \
FROM bidi_locks WHERE browser_name = ?1",
[browser_name],
|r| {
Ok(BidiLockRow {
browser_name: r.get(0)?,
holder_pid: r.get::<_, i64>(1)? as u32,
acquired_at_epoch_s: r.get(2)?,
})
},
)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BidiLockRow {
pub browser_name: String,
pub holder_pid: u32,
pub acquired_at_epoch_s: i64,
}
#[cfg(test)]
mod tests {
use super::*;
fn fresh_at(path: &std::path::Path) -> Registry {
Registry::open_at(path).unwrap()
}
#[test]
fn acquire_when_unlocked_succeeds_immediately() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let reg = fresh_at(tmp.path());
let guard = reg
.bidi_lock_acquire("brave", Duration::from_secs(1))
.unwrap();
assert_eq!(guard.holder_pid(), std::process::id());
let holder = reg.bidi_lock_holder("brave").unwrap().unwrap();
assert_eq!(holder.holder_pid, std::process::id());
}
#[test]
fn drop_releases_the_lock() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let reg = fresh_at(tmp.path());
{
let _guard = reg.bidi_lock_acquire("b", Duration::from_secs(1)).unwrap();
assert!(reg.bidi_lock_holder("b").unwrap().is_some());
}
assert!(reg.bidi_lock_holder("b").unwrap().is_none());
}
#[test]
fn contention_against_self_pid_times_out_with_typed_error() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let reg = fresh_at(tmp.path());
let _g = reg.bidi_lock_acquire("b", Duration::from_secs(1)).unwrap();
let start = Instant::now();
let err = reg
.bidi_lock_acquire("b", Duration::from_millis(250))
.expect_err("must time out");
let elapsed = start.elapsed();
assert!(elapsed >= Duration::from_millis(200));
assert!(elapsed < Duration::from_secs(1));
let typed = err
.downcast_ref::<BidiLockBusy>()
.expect("typed BidiLockBusy");
assert_eq!(typed.holder_pid, std::process::id());
}
#[test]
fn blocks_until_granted_on_release() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let path = tmp.path().to_path_buf();
let reg = fresh_at(&path);
let guard = reg.bidi_lock_acquire("b", Duration::from_secs(1)).unwrap();
const HOLD: Duration = Duration::from_millis(300);
let holder = std::thread::spawn(move || {
std::thread::sleep(HOLD);
drop(guard); });
let reg2 = fresh_at(&path);
let start = Instant::now();
let g2 = reg2
.bidi_lock_acquire("b", Duration::from_secs(5))
.expect("must eventually acquire after release");
let elapsed = start.elapsed();
holder.join().unwrap();
assert_eq!(g2.holder_pid(), std::process::id());
assert!(
elapsed >= HOLD,
"acquired before release: elapsed {elapsed:?} < hold {HOLD:?}"
);
let holder_row = reg2.bidi_lock_holder("b").unwrap().unwrap();
assert_eq!(holder_row.holder_pid, std::process::id());
}
#[test]
fn stale_pid_holder_is_evicted_on_acquire() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let reg = fresh_at(tmp.path());
let dead_pid: u32 = 1; let bogus_pid: u32 = 999_999_999;
reg.conn
.execute(
"INSERT INTO bidi_locks (browser_name, holder_pid, acquired_at_epoch_s) \
VALUES (?1, ?2, ?3)",
rusqlite::params!["b", bogus_pid, crate::registry::now_epoch_s()],
)
.unwrap();
let _ = dead_pid;
let guard = reg.bidi_lock_acquire("b", Duration::from_secs(2)).unwrap();
assert_eq!(guard.holder_pid(), std::process::id());
}
}