Skip to main content

sofie_db/
idempotency.rs

1use std::time::Duration;
2
3use sofie_core::{Error, Result, SqlValue, Timestamp, Uuid};
4
5use crate::{Context, Db, Outcome, Reach, Reply, Resolve, Step, query_values};
6
7const KEPT_FOR: Duration = Duration::from_hours(24);
8/// The claim commits with the work, so a claim left unanswered this long is work whose answer was
9/// lost after it committed; it replays as empty rather than refusing every retry for a day.
10const ANSWERED_WITHIN: Duration = Duration::from_mins(1);
11const PENDING: &str = "A request with this key is still being answered; try again shortly.";
12const REUSED: &str = "This key was already used for a different request.";
13
14#[derive(Clone, Debug, PartialEq, Eq)]
15pub struct IdempotencyKey(String);
16
17impl IdempotencyKey {
18    /// For a vendor call outside keyed work: build it once per call and send it on every retry.
19    pub fn generate() -> Result<Self> {
20        let id =
21            Uuid::now_v7().map_err(|error| Error::internal("drawing an idempotency key", error))?;
22        Ok(Self(id.to_string()))
23    }
24
25    #[must_use]
26    pub fn as_str(&self) -> &str {
27        &self.0
28    }
29}
30
31#[derive(Clone, Debug)]
32pub struct Work {
33    key: String,
34    fingerprint: String,
35}
36
37impl Work {
38    #[must_use]
39    pub fn new(key: String, fingerprint: String) -> Self {
40        Self { key, fingerprint }
41    }
42
43    #[must_use]
44    pub fn key(&self) -> IdempotencyKey {
45        IdempotencyKey(self.key.clone())
46    }
47}
48
49pub trait KeyedWork: Context {
50    fn work(&self) -> Option<&Work>;
51}
52
53impl KeyedWork for crate::worker::JobContext {
54    fn work(&self) -> Option<&Work> {
55        Some(&self.work)
56    }
57}
58
59impl<C: KeyedWork> Resolve<C> for IdempotencyKey {
60    fn resolve(context: &mut C) -> impl Future<Output = Result<Self>> + Send {
61        let key = context.work().map(Work::key).ok_or_else(|| {
62            Error::internal_message(
63                "an IdempotencyKey is given only to a job or the handler of an idempotent route or a webhook",
64            )
65        });
66        std::future::ready(key)
67    }
68}
69
70pub struct Keyed<'a> {
71    db: &'a Db,
72    work: &'a Work,
73}
74
75impl<'a> Keyed<'a> {
76    #[must_use]
77    pub fn new(db: &'a Db, work: &'a Work) -> Self {
78        Self { db, work }
79    }
80
81    pub async fn replay(&self) -> Result<Option<Reply>> {
82        let dialect = self.db.driver().dialect();
83        let sql = format!(
84            "SELECT fingerprint, answered_at, reply, created_at FROM {} WHERE key = {} AND created_at > {}",
85            table(self.db),
86            dialect.placeholder(1),
87            dialect.placeholder(2)
88        );
89        let params = [SqlValue::Text(self.work.key.clone()), kept_since(self.db)?];
90        let rows = query_values(self.db.driver(), &sql, &params).await?;
91        let Some(row) = rows.into_iter().next() else {
92            return Ok(None);
93        };
94        let lost_before = self.db.clock().now().checked_sub(ANSWERED_WITHIN)?;
95        stored_answer(row, self.work, lost_before.timestamp())
96    }
97
98    /// When another request claimed the key first, its answer is returned and nothing written.
99    pub async fn commit(&self, outcome: &mut dyn Outcome, reach: Reach) -> Result<Option<Reply>> {
100        let claim = self.claim()?;
101        let Err(error) = self.db.commit_with(outcome, vec![claim], reach).await else {
102            return Ok(None);
103        };
104        match self.replay().await? {
105            Some(reply) => Ok(Some(reply)),
106            None => Err(error),
107        }
108    }
109
110    pub async fn record(&self, reply: &Reply) -> Result<()> {
111        let dialect = self.db.driver().dialect();
112        let sql = format!(
113            "UPDATE {} SET reply = {}, answered_at = {} WHERE key = {}",
114            table(self.db),
115            dialect.placeholder(1),
116            dialect.placeholder(2),
117            dialect.placeholder(3)
118        );
119        let reply = match reply {
120            Reply::Empty => SqlValue::Null,
121            answered => SqlValue::Json(answered.to_json()?),
122        };
123        let params = [
124            reply,
125            SqlValue::Timestamp(self.db.clock().now().timestamp()),
126            SqlValue::Text(self.work.key.clone()),
127        ];
128        query_values(self.db.driver(), &sql, &params).await?;
129        Ok(())
130    }
131
132    /// A live claim affects no row, which fails the commit as `Conflict`.
133    fn claim(&self) -> Result<Step> {
134        let dialect = self.db.driver().dialect();
135        let placeholders: Vec<String> = (1..=4).map(|index| dialect.placeholder(index)).collect();
136        let sql = format!(
137            "INSERT INTO {} AS stored (key, fingerprint, created_at) VALUES ({}, {}, {}) \
138             ON CONFLICT (key) DO UPDATE SET fingerprint = excluded.fingerprint, \
139             created_at = excluded.created_at, answered_at = NULL, reply = NULL \
140             WHERE stored.created_at <= {}",
141            table(self.db),
142            placeholders[0],
143            placeholders[1],
144            placeholders[2],
145            placeholders[3]
146        );
147        Ok(Step::Statement {
148            sql,
149            params: vec![
150                SqlValue::Text(self.work.key.clone()),
151                SqlValue::Text(self.work.fingerprint.clone()),
152                SqlValue::Timestamp(self.db.clock().now().timestamp()),
153                kept_since(self.db)?,
154            ],
155            expect_rows: Some(crate::ExpectedRows::or_conflict(1)),
156        })
157    }
158}
159
160/// Run on a schedule, not by a request: a client must not choose when its request pays for it.
161pub async fn sweep_expired_keys(db: &Db) -> Result<()> {
162    let sql = format!(
163        "DELETE FROM {} WHERE created_at <= {}",
164        table(db),
165        db.driver().dialect().placeholder(1)
166    );
167    query_values(db.driver(), &sql, &[kept_since(db)?]).await?;
168    Ok(())
169}
170
171fn kept_since(db: &Db) -> Result<SqlValue> {
172    let kept = db.clock().now().checked_sub(KEPT_FOR)?;
173    Ok(SqlValue::Timestamp(kept.timestamp()))
174}
175
176fn table(db: &Db) -> String {
177    db.driver().dialect().qualified("sofie", "idempotency")
178}
179
180fn stored_answer(row: Vec<SqlValue>, work: &Work, lost_before: Timestamp) -> Result<Option<Reply>> {
181    match <[SqlValue; 4]>::try_from(row) {
182        Ok([SqlValue::Text(fingerprint), ..]) if fingerprint != work.fingerprint => {
183            Err(Error::Rule(REUSED.to_owned()))
184        }
185        Ok([_, SqlValue::Null, _, SqlValue::Timestamp(claimed)]) if claimed > lost_before => {
186            Err(Error::Rule(PENDING.to_owned()))
187        }
188        Ok(
189            [
190                _,
191                SqlValue::Null | SqlValue::Timestamp(_),
192                SqlValue::Null,
193                _,
194            ],
195        ) => Ok(Some(Reply::Empty)),
196        Ok([_, SqlValue::Timestamp(_), SqlValue::Json(value), _]) => Ok(Some(Reply::Json(value))),
197        row => Err(Error::internal_message(format!(
198            "unexpected row in the stored answers: {:?}",
199            row.map(|row| row.map(|value| value.kind()))
200        ))),
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use sofie_core::Timestamp;
207
208    use super::*;
209
210    fn work() -> Work {
211        Work::new("key".to_owned(), "fingerprint".to_owned())
212    }
213
214    #[test]
215    fn a_generated_key_is_new_each_time() {
216        let key = IdempotencyKey::generate().unwrap();
217
218        assert_ne!(key, IdempotencyKey::generate().unwrap());
219        assert_eq!(key.as_str().len(), 36);
220    }
221
222    fn row(fingerprint: &str, answered: bool, reply: SqlValue) -> Vec<SqlValue> {
223        let answered_at = if answered {
224            SqlValue::Timestamp(Timestamp::UNIX_EPOCH)
225        } else {
226            SqlValue::Null
227        };
228        let claimed = SqlValue::Timestamp(Timestamp::UNIX_EPOCH);
229        vec![
230            SqlValue::Text(fingerprint.to_owned()),
231            answered_at,
232            reply,
233            claimed,
234        ]
235    }
236
237    fn answer(row: Vec<SqlValue>) -> Result<Option<Reply>> {
238        stored_answer(row, &work(), Timestamp::MIN)
239    }
240
241    #[test]
242    fn an_answered_key_replays_its_reply() {
243        let json = serde_json::json!({ "id": 1 });
244
245        let empty = answer(row("fingerprint", true, SqlValue::Null)).unwrap();
246        let answered = answer(row("fingerprint", true, SqlValue::Json(json.clone()))).unwrap();
247
248        assert!(matches!(empty, Some(Reply::Empty)));
249        assert!(matches!(answered, Some(Reply::Json(value)) if value == json));
250    }
251
252    #[test]
253    fn a_pending_or_reused_key_is_refused() {
254        let pending = answer(row("fingerprint", false, SqlValue::Null));
255        let reused = answer(row("other", true, SqlValue::Null));
256
257        assert!(matches!(pending, Err(Error::Rule(message)) if message == PENDING));
258        assert!(matches!(reused, Err(Error::Rule(message)) if message == REUSED));
259    }
260
261    #[test]
262    fn a_claim_whose_answer_was_lost_replays_as_empty_instead_of_pending() {
263        let since_claimed = Timestamp::UNIX_EPOCH
264            .checked_add(Duration::from_secs(1))
265            .unwrap();
266
267        let lost = stored_answer(
268            row("fingerprint", false, SqlValue::Null),
269            &work(),
270            since_claimed,
271        );
272
273        assert!(matches!(lost, Ok(Some(Reply::Empty))), "{lost:?}");
274    }
275}