use std::time::Duration;
use secrecy::{ExposeSecret, SecretString};
use sqlx::{Connection, Executor, PgConnection, PgPool, Row};
use tokio::sync::watch;
use crate::error::{KernelError, Result};
pub const WRITER_LOCK_KEY: i64 = 0x0067_776b_0077_7231;
pub const LOCK_PROBE_INTERVAL: Duration = Duration::from_secs(5);
#[derive(Debug)]
pub struct WriterLock {
cancelled: watch::Receiver<bool>,
probe: tokio::task::JoinHandle<()>,
}
impl WriterLock {
pub async fn acquire(database_url: &SecretString) -> Result<Self> {
Self::acquire_with_interval(database_url, LOCK_PROBE_INTERVAL).await
}
pub async fn acquire_with_interval(
database_url: &SecretString,
probe_interval: Duration,
) -> Result<Self> {
let mut conn = PgConnection::connect(database_url.expose_secret()).await?;
let acquired: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
.bind(WRITER_LOCK_KEY)
.fetch_one(&mut conn)
.await?;
if !acquired {
return Err(KernelError::Writer(
"another kernel already holds the writer lock on this database; \
exactly one process may write per store"
.to_owned(),
));
}
let (tx, cancelled) = watch::channel(false);
let probe = tokio::spawn(async move {
loop {
tokio::time::sleep(probe_interval).await;
if conn.execute("SELECT 1").await.is_err() {
let _ = tx.send(true);
return;
}
if tx.is_closed() {
return;
}
}
});
Ok(Self { cancelled, probe })
}
pub fn is_cancelled(&self) -> bool {
*self.cancelled.borrow()
}
pub async fn cancelled(&self) {
let mut rx = self.cancelled.clone();
while !*rx.borrow_and_update() {
if rx.changed().await.is_err() {
return; }
}
}
}
impl Drop for WriterLock {
fn drop(&mut self) {
self.probe.abort();
}
}
pub async fn claim_epoch(pool: &PgPool) -> Result<i64> {
let row = sqlx::query(
"UPDATE gwk_internal.writer SET epoch = epoch + 1 WHERE id = 1 RETURNING epoch",
)
.fetch_optional(pool)
.await?
.ok_or_else(|| {
KernelError::Schema(
"gwk_internal.writer has no singleton row: this database was not initialized by \
`gw admin init`"
.to_owned(),
)
})?;
Ok(row.try_get("epoch")?)
}