use super::store;
use crate::{
ServerFrame,
relay::{RelayRequest, RelayResult},
};
use anyhow::{Context, Result, bail};
use rusqlite::{Connection, OptionalExtension, params};
use serde_json::{Value, json};
use std::sync::{Arc, Mutex, MutexGuard};
use uuid::Uuid;
#[derive(Clone)]
pub struct Queue {
connection: Arc<Mutex<Connection>>,
}
impl Queue {
pub fn open(state: &store::Store) -> Result<Self> {
Ok(Self {
connection: Arc::new(Mutex::new(open_database(state)?)),
})
}
fn database(&self) -> Result<MutexGuard<'_, Connection>> {
self.connection
.lock()
.map_err(|_| anyhow::anyhow!("database operation lock poisoned"))
}
}
fn open_database(state: &store::Store) -> Result<Connection> {
let path = state.directory().join("relay.sqlite3");
store::secure_open(&path)?;
let conn = Connection::open(path)?;
conn.busy_timeout(std::time::Duration::from_secs(5))?;
conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL; PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS cursors(session TEXT NOT NULL, actor TEXT NOT NULL, sequence INTEGER NOT NULL, PRIMARY KEY(session,actor));
CREATE TABLE IF NOT EXISTS generations(session TEXT PRIMARY KEY,generation INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS inbox(session TEXT NOT NULL, delivery_id TEXT NOT NULL, actor TEXT NOT NULL, sequence INTEGER NOT NULL, frame TEXT NOT NULL, delivered INTEGER NOT NULL DEFAULT 0, attempts INTEGER NOT NULL DEFAULT 0, next_attempt INTEGER NOT NULL DEFAULT 0, PRIMARY KEY(session,delivery_id), UNIQUE(session,actor,sequence));
CREATE TABLE IF NOT EXISTS outbox(request_id TEXT PRIMARY KEY,session TEXT NOT NULL,request TEXT NOT NULL,state TEXT NOT NULL DEFAULT 'pending',result TEXT,error TEXT,attempts INTEGER NOT NULL DEFAULT 0,next_attempt INTEGER NOT NULL DEFAULT 0,created INTEGER NOT NULL,expected_generation INTEGER);
CREATE INDEX IF NOT EXISTS inbox_pending ON inbox(delivered,next_attempt);
CREATE INDEX IF NOT EXISTS outbox_pending ON outbox(state,next_attempt);")?;
let generation_column: i64 = conn.query_row(
"SELECT COUNT(*) FROM pragma_table_info('outbox') WHERE name='expected_generation'",
[],
|r| r.get(0),
)?;
if generation_column == 0 {
conn.execute(
"ALTER TABLE outbox ADD COLUMN expected_generation INTEGER",
[],
)?;
}
Ok(conn)
}
impl Queue {
pub fn cursor(&self, session: &str, actor: &str) -> Result<i64> {
Ok(self
.database()?
.query_row(
"SELECT sequence FROM cursors WHERE session=?1 AND actor=?2",
params![session, actor],
|r| r.get(0),
)
.optional()?
.unwrap_or(0))
}
}
impl Queue {
pub fn generation(&self, session: &str) -> Result<Option<i64>> {
Ok(self
.database()?
.query_row(
"SELECT generation FROM generations WHERE session=?1",
params![session],
|r| r.get(0),
)
.optional()?)
}
}
pub fn stream_session(session: &str, generation: Option<i64>) -> String {
generation.map_or_else(|| session.to_owned(), |n| format!("{session}#{n}"))
}
impl Queue {
pub fn adopt_generation(&self, session: &str, next: Option<i64>) -> Result<()> {
let Some(next) = next else { return Ok(()) };
let mut conn = self.database()?;
let tx = conn.transaction()?;
let previous: Option<i64> = tx
.query_row(
"SELECT generation FROM generations WHERE session=?1",
params![session],
|r| r.get(0),
)
.optional()?;
if previous.is_some() && previous != Some(next) {
let stream = stream_session(session, previous);
tx.execute(
"UPDATE inbox SET delivered=-1 WHERE session=?1 AND delivered=0",
params![stream],
)?;
tx.execute("UPDATE outbox SET state='failed',error=?2 WHERE session=?1 AND state='pending'",params![session,json!({"code":"testing_generation_changed","message":"test environment changed; inspect then explicitly resubmit this request"}).to_string()])?;
}
tx.execute("INSERT INTO generations(session,generation) VALUES(?1,?2) ON CONFLICT(session) DO UPDATE SET generation=excluded.generation",params![session,next])?;
tx.commit()?;
Ok(())
}
}
impl Queue {
pub fn receive(&self, session: &str, frame: &ServerFrame, raw: &str) -> Result<i64> {
let (id, actor, sequence) = frame
.delivery_position()
.context("frame has no delivery sequence")?;
let mut conn = self.database()?;
let tx = conn.transaction()?;
let old: Option<(String, i64, String)> = tx
.query_row(
"SELECT actor,sequence,frame FROM inbox WHERE session=?1 AND delivery_id=?2",
params![session, id.to_string()],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
)
.optional()?;
if let Some((old_actor, old_sequence, old_frame)) = old {
if old_actor != actor
|| old_sequence != sequence
|| !same_delivery_identity(&serde_json::from_str(&old_frame)?, frame)
{
bail!("server reused a delivery ID with a different immutable identity")
}
} else {
tx.execute(
"INSERT INTO inbox(session,delivery_id,actor,sequence,frame) VALUES(?1,?2,?3,?4,?5)",
params![session, id.to_string(), actor, sequence, raw],
)?;
}
let mut through: i64 = tx
.query_row(
"SELECT sequence FROM cursors WHERE session=?1 AND actor=?2",
params![session, actor],
|r| r.get(0),
)
.optional()?
.unwrap_or(0);
loop {
let next: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM inbox WHERE session=?1 AND actor=?2 AND sequence=?3)",
params![session, actor, through + 1],
|r| r.get(0),
)?;
if !next {
break;
}
through += 1;
}
tx.execute("INSERT INTO cursors(session,actor,sequence) VALUES(?1,?2,?3) ON CONFLICT(session,actor) DO UPDATE SET sequence=excluded.sequence",params![session,actor,through])?;
tx.commit()?;
Ok(through)
}
}
fn same_delivery_identity(old: &ServerFrame, new: &ServerFrame) -> bool {
match (old, new) {
(ServerFrame::Message { message: old, .. }, ServerFrame::Message { message: new, .. }) => {
old.id == new.id
&& old.conversation_id == new.conversation_id
&& old.sender == new.sender
&& old.sequence == new.sequence
&& old.created_at == new.created_at
}
(
ServerFrame::Receipt {
message_id: old, ..
},
ServerFrame::Receipt {
message_id: new, ..
},
) => old == new,
_ => false,
}
}
impl Queue {
pub fn enqueue(&self, request: &RelayRequest, original: &Value) -> Result<()> {
let conn = self.database()?;
let raw = serde_json::to_string(original)?;
let previous: Option<String> = conn
.query_row(
"SELECT request FROM outbox WHERE request_id=?1",
params![request.request_id.to_string()],
|r| r.get(0),
)
.optional()?;
if let Some(previous) = previous {
if serde_json::from_str::<Value>(&previous)? != *original {
bail!("request_id was already used for a different request")
}
return Ok(());
}
let session = store::session_key(&request.profile, request.testing_environment_id);
let known: Option<i64> = conn
.query_row(
"SELECT generation FROM generations WHERE session=?1",
params![session],
|r| r.get(0),
)
.optional()?;
conn.execute(
"INSERT INTO outbox(request_id,session,request,created,expected_generation) VALUES(?1,?2,?3,?4,?5)",
params![
request.request_id.to_string(),
store::session_key(&request.profile, request.testing_environment_id),
raw,
(store::now() as i64),
request.testing_generation.or(known)
],
)?;
Ok(())
}
}
impl Queue {
pub fn result(&self, id: Uuid) -> Result<Option<RelayResult>> {
let row: Option<(String, Option<String>, Option<String>, String)> = self
.database()?
.query_row(
"SELECT state,result,error,request FROM outbox WHERE request_id=?1",
params![id.to_string()],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
)
.optional()?;
row.map(|(state, result, error, request)| {
Ok(RelayResult {
request_id: id,
state,
result: result.map(|s| serde_json::from_str(&s)).transpose()?,
error: error.map(|s| serde_json::from_str(&s)).transpose()?,
request: serde_json::from_str(&request)?,
})
})
.transpose()
}
}
impl Queue {
pub fn pending(&self) -> Result<Vec<RelayRequest>> {
let conn = self.database()?;
conn.execute("UPDATE outbox SET expected_generation=(SELECT generation FROM generations WHERE generations.session=outbox.session) WHERE state='pending' AND expected_generation IS NULL AND EXISTS (SELECT 1 FROM generations WHERE generations.session=outbox.session)", [])?;
let mut stmt=conn.prepare("SELECT request,expected_generation FROM outbox o WHERE state='pending' AND next_attempt<=?1 AND NOT EXISTS (SELECT 1 FROM outbox previous WHERE previous.session=o.session AND previous.state='pending' AND previous.rowid<o.rowid) ORDER BY rowid LIMIT 50")?;
let rows = stmt
.query_map(params![(store::now() as i64)], |r| {
Ok((r.get::<_, String>(0)?, r.get::<_, Option<i64>>(1)?))
})?
.collect::<std::result::Result<Vec<_>, _>>()?;
rows.iter()
.map(|(s, generation)| {
let mut request: RelayRequest = serde_json::from_str(s)?;
request.testing_generation = generation.or(request.testing_generation);
Ok(request)
})
.collect()
}
}
impl Queue {
pub fn finish(&self, id: Uuid, result: Option<&Value>, error: Option<&Value>) -> Result<()> {
self.database()?.execute(
"UPDATE outbox SET state=?2,result=?3,error=?4 WHERE request_id=?1",
params![
id.to_string(),
if error.is_some() {
"failed"
} else {
"completed"
},
result.map(serde_json::to_string).transpose()?,
error.map(serde_json::to_string).transpose()?
],
)?;
Ok(())
}
}
impl Queue {
pub fn retry(&self, id: Uuid, error: &Value) -> Result<()> {
self.database()?.execute("UPDATE outbox SET attempts=attempts+1,next_attempt=?2+MIN(60,1 << MIN(attempts,6)),error=?3 WHERE request_id=?1",params![id.to_string(),(store::now() as i64),serde_json::to_string(error)?])?;
Ok(())
}
}
pub struct WebhookWork {
pub session: String,
pub delivery_id: String,
pub frame: Value,
}
impl Queue {
pub fn webhook_work(&self) -> Result<Vec<WebhookWork>> {
let conn = self.database()?;
let mut stmt=conn.prepare("SELECT session,delivery_id,frame FROM inbox i WHERE delivered=0 AND next_attempt<=?1 AND NOT EXISTS (SELECT 1 FROM inbox prior WHERE prior.session=i.session AND prior.actor=i.actor AND prior.delivered=0 AND prior.sequence<i.sequence) ORDER BY rowid LIMIT 50")?;
let rows = stmt
.query_map(params![(store::now() as i64)], |r| {
Ok((
r.get::<_, String>(0)?,
r.get::<_, String>(1)?,
r.get::<_, String>(2)?,
))
})?
.collect::<std::result::Result<Vec<_>, _>>()?;
rows.into_iter()
.map(|(session, delivery_id, frame)| {
Ok(WebhookWork {
session,
delivery_id,
frame: serde_json::from_str(&frame)?,
})
})
.collect()
}
}
impl Queue {
pub fn webhook_is_pending(&self, session: &str, id: &str) -> Result<bool> {
Ok(self.database()?.query_row(
"SELECT EXISTS(SELECT 1 FROM inbox WHERE session=?1 AND delivery_id=?2 AND delivered=0)",
params![session, id],
|row| row.get(0),
)?)
}
pub fn webhook_done(
&self,
session: &str,
id: &str,
success: bool,
receipt: Option<&RelayRequest>,
) -> Result<()> {
let mut conn = self.database()?;
let tx = conn.transaction()?;
let changed = tx.execute(
"UPDATE inbox SET delivered=?3,attempts=attempts+1,next_attempt=?4+MIN(300,1 << MIN(attempts,9)) WHERE session=?1 AND delivery_id=?2 AND delivered=0",
params![session, id, success, store::now() as i64],
)?;
if changed == 0 {
tx.commit()?;
return Ok(());
}
if success && let Some(request) = receipt {
tx.execute(
"INSERT OR IGNORE INTO outbox(request_id,session,request,created,expected_generation) VALUES(?1,?2,?3,?4,?5)",
params![
request.request_id.to_string(),
store::session_key(&request.profile, request.testing_environment_id),
serde_json::to_string(request)?,
store::now() as i64,
request.testing_generation
],
)?;
}
tx.commit()?;
Ok(())
}
}
impl Queue {
pub fn stats(&self) -> Result<Value> {
let conn = self.database()?;
let pending: i64 = conn.query_row(
"SELECT count(*) FROM outbox WHERE state='pending'",
[],
|r| r.get(0),
)?;
let callbacks: i64 =
conn.query_row("SELECT count(*) FROM inbox WHERE delivered=0", [], |r| {
r.get(0)
})?;
let failed: i64 = conn.query_row(
"SELECT count(*) FROM outbox WHERE state='failed'",
[],
|r| r.get(0),
)?;
Ok(
json!({"pending_requests":pending,"pending_webhooks":callbacks,"failed_requests":failed}),
)
}
}
impl Queue {
pub fn expire_transient_commands(&self) -> Result<()> {
self.database()?.execute("UPDATE outbox SET state='failed',error=?1 WHERE state='pending' AND json_extract(request,'$.request.operation')='set_presence'",params![json!({"code":"transient_expired","message":"presence change expired when the daemon stopped; send a current activity"}).to_string()])?;
Ok(())
}
}