use std::time::Duration;
use sofie_core::{Error, Result, SqlValue, Timestamp, Uuid};
use crate::{Context, Db, Outcome, Reach, Reply, Resolve, Step, query_values};
const KEPT_FOR: Duration = Duration::from_hours(24);
const ANSWERED_WITHIN: Duration = Duration::from_mins(1);
const PENDING: &str = "A request with this key is still being answered; try again shortly.";
const REUSED: &str = "This key was already used for a different request.";
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct IdempotencyKey(String);
impl IdempotencyKey {
pub fn generate() -> Result<Self> {
let id =
Uuid::now_v7().map_err(|error| Error::internal("drawing an idempotency key", error))?;
Ok(Self(id.to_string()))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
#[derive(Clone, Debug)]
pub struct Work {
key: String,
fingerprint: String,
}
impl Work {
#[must_use]
pub fn new(key: String, fingerprint: String) -> Self {
Self { key, fingerprint }
}
#[must_use]
pub fn key(&self) -> IdempotencyKey {
IdempotencyKey(self.key.clone())
}
}
pub trait KeyedWork: Context {
fn work(&self) -> Option<&Work>;
}
impl KeyedWork for crate::worker::JobContext {
fn work(&self) -> Option<&Work> {
Some(&self.work)
}
}
impl<C: KeyedWork> Resolve<C> for IdempotencyKey {
fn resolve(context: &mut C) -> impl Future<Output = Result<Self>> + Send {
let key = context.work().map(Work::key).ok_or_else(|| {
Error::internal_message(
"an IdempotencyKey is given only to a job or the handler of an idempotent route or a webhook",
)
});
std::future::ready(key)
}
}
pub struct Keyed<'a> {
db: &'a Db,
work: &'a Work,
}
impl<'a> Keyed<'a> {
#[must_use]
pub fn new(db: &'a Db, work: &'a Work) -> Self {
Self { db, work }
}
pub async fn replay(&self) -> Result<Option<Reply>> {
let dialect = self.db.driver().dialect();
let sql = format!(
"SELECT fingerprint, answered_at, reply, created_at FROM {} WHERE key = {} AND created_at > {}",
table(self.db),
dialect.placeholder(1),
dialect.placeholder(2)
);
let params = [SqlValue::Text(self.work.key.clone()), kept_since(self.db)?];
let rows = query_values(self.db.driver(), &sql, ¶ms).await?;
let Some(row) = rows.into_iter().next() else {
return Ok(None);
};
let lost_before = self.db.clock().now().checked_sub(ANSWERED_WITHIN)?;
stored_answer(row, self.work, lost_before.timestamp())
}
pub async fn commit(&self, outcome: &mut dyn Outcome, reach: Reach) -> Result<Option<Reply>> {
let claim = self.claim()?;
let Err(error) = self.db.commit_with(outcome, vec![claim], reach).await else {
return Ok(None);
};
match self.replay().await? {
Some(reply) => Ok(Some(reply)),
None => Err(error),
}
}
pub async fn record(&self, reply: &Reply) -> Result<()> {
let dialect = self.db.driver().dialect();
let sql = format!(
"UPDATE {} SET reply = {}, answered_at = {} WHERE key = {}",
table(self.db),
dialect.placeholder(1),
dialect.placeholder(2),
dialect.placeholder(3)
);
let reply = match reply {
Reply::Empty => SqlValue::Null,
answered => SqlValue::Json(answered.to_json()?),
};
let params = [
reply,
SqlValue::Timestamp(self.db.clock().now().timestamp()),
SqlValue::Text(self.work.key.clone()),
];
query_values(self.db.driver(), &sql, ¶ms).await?;
Ok(())
}
fn claim(&self) -> Result<Step> {
let dialect = self.db.driver().dialect();
let placeholders: Vec<String> = (1..=4).map(|index| dialect.placeholder(index)).collect();
let sql = format!(
"INSERT INTO {} AS stored (key, fingerprint, created_at) VALUES ({}, {}, {}) \
ON CONFLICT (key) DO UPDATE SET fingerprint = excluded.fingerprint, \
created_at = excluded.created_at, answered_at = NULL, reply = NULL \
WHERE stored.created_at <= {}",
table(self.db),
placeholders[0],
placeholders[1],
placeholders[2],
placeholders[3]
);
Ok(Step::Statement {
sql,
params: vec![
SqlValue::Text(self.work.key.clone()),
SqlValue::Text(self.work.fingerprint.clone()),
SqlValue::Timestamp(self.db.clock().now().timestamp()),
kept_since(self.db)?,
],
expect_rows: Some(crate::ExpectedRows::or_conflict(1)),
})
}
}
pub async fn sweep_expired_keys(db: &Db) -> Result<()> {
let sql = format!(
"DELETE FROM {} WHERE created_at <= {}",
table(db),
db.driver().dialect().placeholder(1)
);
query_values(db.driver(), &sql, &[kept_since(db)?]).await?;
Ok(())
}
fn kept_since(db: &Db) -> Result<SqlValue> {
let kept = db.clock().now().checked_sub(KEPT_FOR)?;
Ok(SqlValue::Timestamp(kept.timestamp()))
}
fn table(db: &Db) -> String {
db.driver().dialect().qualified("sofie", "idempotency")
}
fn stored_answer(row: Vec<SqlValue>, work: &Work, lost_before: Timestamp) -> Result<Option<Reply>> {
match <[SqlValue; 4]>::try_from(row) {
Ok([SqlValue::Text(fingerprint), ..]) if fingerprint != work.fingerprint => {
Err(Error::Rule(REUSED.to_owned()))
}
Ok([_, SqlValue::Null, _, SqlValue::Timestamp(claimed)]) if claimed > lost_before => {
Err(Error::Rule(PENDING.to_owned()))
}
Ok(
[
_,
SqlValue::Null | SqlValue::Timestamp(_),
SqlValue::Null,
_,
],
) => Ok(Some(Reply::Empty)),
Ok([_, SqlValue::Timestamp(_), SqlValue::Json(value), _]) => Ok(Some(Reply::Json(value))),
row => Err(Error::internal_message(format!(
"unexpected row in the stored answers: {:?}",
row.map(|row| row.map(|value| value.kind()))
))),
}
}
#[cfg(test)]
mod tests {
use sofie_core::Timestamp;
use super::*;
fn work() -> Work {
Work::new("key".to_owned(), "fingerprint".to_owned())
}
#[test]
fn a_generated_key_is_new_each_time() {
let key = IdempotencyKey::generate().unwrap();
assert_ne!(key, IdempotencyKey::generate().unwrap());
assert_eq!(key.as_str().len(), 36);
}
fn row(fingerprint: &str, answered: bool, reply: SqlValue) -> Vec<SqlValue> {
let answered_at = if answered {
SqlValue::Timestamp(Timestamp::UNIX_EPOCH)
} else {
SqlValue::Null
};
let claimed = SqlValue::Timestamp(Timestamp::UNIX_EPOCH);
vec![
SqlValue::Text(fingerprint.to_owned()),
answered_at,
reply,
claimed,
]
}
fn answer(row: Vec<SqlValue>) -> Result<Option<Reply>> {
stored_answer(row, &work(), Timestamp::MIN)
}
#[test]
fn an_answered_key_replays_its_reply() {
let json = serde_json::json!({ "id": 1 });
let empty = answer(row("fingerprint", true, SqlValue::Null)).unwrap();
let answered = answer(row("fingerprint", true, SqlValue::Json(json.clone()))).unwrap();
assert!(matches!(empty, Some(Reply::Empty)));
assert!(matches!(answered, Some(Reply::Json(value)) if value == json));
}
#[test]
fn a_pending_or_reused_key_is_refused() {
let pending = answer(row("fingerprint", false, SqlValue::Null));
let reused = answer(row("other", true, SqlValue::Null));
assert!(matches!(pending, Err(Error::Rule(message)) if message == PENDING));
assert!(matches!(reused, Err(Error::Rule(message)) if message == REUSED));
}
#[test]
fn a_claim_whose_answer_was_lost_replays_as_empty_instead_of_pending() {
let since_claimed = Timestamp::UNIX_EPOCH
.checked_add(Duration::from_secs(1))
.unwrap();
let lost = stored_answer(
row("fingerprint", false, SqlValue::Null),
&work(),
since_claimed,
);
assert!(matches!(lost, Ok(Some(Reply::Empty))), "{lost:?}");
}
}