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);
8const 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 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, ¶ms).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 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, ¶ms).await?;
129 Ok(())
130 }
131
132 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
160pub 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}