taquba_workflow/
effects.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::{Arc, Mutex};
3
4use serde::{Deserialize, Serialize};
5use taquba::EnqueueRequest;
6
7use crate::error::{Error, Result};
8use crate::keys::RESERVED_KV_PREFIX;
9
10/// Application KV effects staged during a step, applied in the same
11/// transaction as the settlement that commits the step's outcome.
12///
13/// Obtained through [`Step::effects`](crate::Step::effects). Writes and
14/// deletes staged here are applied atomically with the acknowledgement
15/// of the [`StepOutcome`](crate::StepOutcome) the runner returns
16/// (`Continue`, `Succeed`, `Fail` and `Cancel`): either the settlement
17/// and every staged effect commit together or none of them do. Delivery
18/// is at-least-once, so a retried step stages its effects again and
19/// every staged value must be correct when applied more than once. No
20/// effects are applied when the runner returns a
21/// [`StepError`](crate::StepError) or when an external
22/// [`WorkflowRuntime::cancel`](crate::WorkflowRuntime::cancel)
23/// overrides the runner's outcome.
24///
25/// Each operation is validated when it is staged: keys must not start
26/// with the reserved `workflow/` prefix
27/// ([`RESERVED_KV_PREFIX`](crate::RESERVED_KV_PREFIX)), values are
28/// capped at [`taquba::MAX_KV_VALUE_SIZE`] each (an effects set has no
29/// aggregate cap) and a key cannot be staged for both a write and a
30/// delete within one step. The handle is sealed once `run_step`
31/// returns; staging through a clone retained past that point returns
32/// [`Error::EffectsSealed`].
33///
34/// The handle is cheap to clone and clones share one accumulator. Use
35/// [`EffectsHandle::detached`] when constructing a
36/// [`Step`](crate::Step) in tests.
37#[derive(Debug, Clone)]
38pub struct EffectsHandle {
39    inner: Arc<Mutex<EffectsState>>,
40}
41
42#[derive(Debug, Default)]
43struct EffectsState {
44    staged: StagedEffects,
45    /// Writes under the reserved namespace applied only when the step's
46    /// failure terminates the run; see [`EffectsHandle::put_reserved_on_failure`].
47    on_failure: HashMap<Vec<u8>, Vec<u8>>,
48    sealed: bool,
49}
50
51impl EffectsState {
52    fn put(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
53        self.check_key(&key)?;
54        if value.len() > taquba::MAX_KV_VALUE_SIZE {
55            return Err(Error::Queue(taquba::Error::KvValueTooLarge {
56                size: value.len(),
57                max: taquba::MAX_KV_VALUE_SIZE,
58            }));
59        }
60        if self.staged.deletes.contains(&key) {
61            return Err(Error::ConflictingKvEffect(display_key(&key)));
62        }
63        self.staged.writes.insert(key, value);
64        Ok(())
65    }
66
67    /// Stage a write under a key of the reserved namespace, for the
68    /// crate's own records; validated like [`put`](Self::put) otherwise.
69    fn put_reserved(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
70        self.check_reserved(&value)?;
71        self.staged.writes.insert(key, value);
72        Ok(())
73    }
74
75    /// Stage a reserved-namespace write into the failure set.
76    fn put_reserved_on_failure(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
77        self.check_reserved(&value)?;
78        self.on_failure.insert(key, value);
79        Ok(())
80    }
81
82    fn check_reserved(&self, value: &[u8]) -> Result<()> {
83        if self.sealed {
84            return Err(Error::EffectsSealed);
85        }
86        if value.len() > taquba::MAX_KV_VALUE_SIZE {
87            return Err(Error::Queue(taquba::Error::KvValueTooLarge {
88                size: value.len(),
89                max: taquba::MAX_KV_VALUE_SIZE,
90            }));
91        }
92        Ok(())
93    }
94
95    /// Seal the state and move out both staged sets.
96    fn seal_and_take(&mut self) -> SealedEffects {
97        self.sealed = true;
98        SealedEffects {
99            outcome: std::mem::take(&mut self.staged),
100            on_failure: std::mem::take(&mut self.on_failure),
101        }
102    }
103
104    fn delete(&mut self, key: Vec<u8>) -> Result<()> {
105        self.check_key(&key)?;
106        if self.staged.writes.contains_key(&key) {
107            return Err(Error::ConflictingKvEffect(display_key(&key)));
108        }
109        self.staged.deletes.insert(key);
110        Ok(())
111    }
112
113    fn check_key(&self, key: &[u8]) -> Result<()> {
114        if self.sealed {
115            return Err(Error::EffectsSealed);
116        }
117        if key.starts_with(RESERVED_KV_PREFIX.as_bytes()) {
118            return Err(Error::ReservedKvKey(display_key(key)));
119        }
120        Ok(())
121    }
122}
123
124/// Writes and deletes accumulated by an [`EffectsHandle`]. Stored in
125/// the step-output replay record so a replayed delivery applies the
126/// same effects.
127#[derive(Debug, Clone, Default, Serialize, Deserialize)]
128pub(crate) struct StagedEffects {
129    pub(crate) writes: HashMap<Vec<u8>, Vec<u8>>,
130    pub(crate) deletes: HashSet<Vec<u8>>,
131}
132
133/// Everything an [`EffectsHandle`] holds when its delivery returns: the
134/// effects of the returned outcome and the writes reserved for a
135/// terminating failure.
136#[derive(Debug, Default)]
137pub(crate) struct SealedEffects {
138    pub(crate) outcome: StagedEffects,
139    pub(crate) on_failure: HashMap<Vec<u8>, Vec<u8>>,
140}
141
142impl EffectsHandle {
143    /// Build a handle bound to no delivery, for constructing a
144    /// [`Step`](crate::Step) in tests. A detached handle accepts and
145    /// validates effects like a delivery-bound one, is never sealed and
146    /// its staged effects are never applied.
147    pub fn detached() -> Self {
148        Self::for_delivery()
149    }
150
151    pub(crate) fn for_delivery() -> Self {
152        Self {
153            inner: Arc::new(Mutex::new(EffectsState::default())),
154        }
155    }
156
157    /// Stage a write of `value` under `key` in the caller KV namespace.
158    ///
159    /// # Errors
160    ///
161    /// [`Error::ReservedKvKey`] when `key` starts with the reserved
162    /// `workflow/` prefix, [`Error::Queue`] with
163    /// [`taquba::Error::KvValueTooLarge`] when `value` exceeds
164    /// [`taquba::MAX_KV_VALUE_SIZE`], [`Error::ConflictingKvEffect`]
165    /// when `key` is already staged for a delete and
166    /// [`Error::EffectsSealed`] when the step has returned.
167    pub fn put(&self, key: impl Into<Vec<u8>>, value: impl Into<Vec<u8>>) -> Result<()> {
168        self.inner.lock().unwrap().put(key.into(), value.into())
169    }
170
171    /// Stage a delete of `key` from the caller KV namespace.
172    ///
173    /// # Errors
174    ///
175    /// [`Error::ReservedKvKey`] when `key` starts with the reserved
176    /// `workflow/` prefix, [`Error::ConflictingKvEffect`] when `key` is
177    /// already staged for a write and [`Error::EffectsSealed`] when the
178    /// step has returned.
179    pub fn delete(&self, key: impl Into<Vec<u8>>) -> Result<()> {
180        self.inner.lock().unwrap().delete(key.into())
181    }
182
183    /// Stage a write under a key of the reserved `workflow/` namespace,
184    /// for the crate's own records, applied with the returned outcome
185    /// like [`put`](Self::put).
186    pub(crate) fn put_reserved(&self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
187        self.inner.lock().unwrap().put_reserved(key, value)
188    }
189
190    /// Stage a write under a key of the reserved `workflow/` namespace
191    /// that is applied only when the step returns a
192    /// [`StepError`](crate::StepError) that terminates the run: a
193    /// permanent error, or a transient one on the step's last attempt.
194    /// A retried attempt applies nothing, and the writes of an acking
195    /// outcome are the ones staged through [`put`](Self::put).
196    pub(crate) fn put_reserved_on_failure(&self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
197        self.inner
198            .lock()
199            .unwrap()
200            .put_reserved_on_failure(key, value)
201    }
202
203    /// Seal the handle and move out everything staged. An effect staged
204    /// after the seal could not join the settlement, so later staging
205    /// attempts return [`Error::EffectsSealed`].
206    pub(crate) fn seal_and_take(&self) -> SealedEffects {
207        self.inner.lock().unwrap().seal_and_take()
208    }
209}
210
211/// Effects staged by a [`TerminalHook`](crate::TerminalHook) during a
212/// notification delivery, applied in the same transaction as the
213/// notification job's acknowledgement.
214///
215/// Passed to
216/// [`TerminalHook::on_termination`](crate::TerminalHook::on_termination).
217/// Beyond the KV writes and deletes of [`EffectsHandle`] (validated by
218/// the same rules), a hook stages follow-up enqueues, so work driven by
219/// a run's termination is created atomically with the notification
220/// being acknowledged. Effects are applied only when the hook returns
221/// `Ok`; a retried hook stages its effects again.
222///
223/// The handle is sealed once the hook returns; staging through a clone
224/// retained past that point returns [`Error::EffectsSealed`]. Use
225/// [`TerminalEffects::detached`] when invoking a hook directly in
226/// tests.
227#[derive(Debug, Clone)]
228pub struct TerminalEffects {
229    inner: Arc<Mutex<TerminalState>>,
230}
231
232#[derive(Debug, Default)]
233struct TerminalState {
234    kv: EffectsState,
235    enqueues: Vec<EnqueueRequest>,
236}
237
238impl TerminalEffects {
239    /// Build a handle bound to no delivery, for invoking a
240    /// [`TerminalHook`](crate::TerminalHook) directly in tests. A
241    /// detached handle accepts and validates effects like a
242    /// delivery-bound one, is never sealed and its staged effects are
243    /// never applied.
244    pub fn detached() -> Self {
245        Self::for_delivery()
246    }
247
248    pub(crate) fn for_delivery() -> Self {
249        Self {
250            inner: Arc::new(Mutex::new(TerminalState::default())),
251        }
252    }
253
254    /// Stage a follow-up enqueue, committed with the notification's
255    /// acknowledgement.
256    ///
257    /// # Errors
258    ///
259    /// [`Error::EffectsSealed`] when the hook has returned.
260    pub fn enqueue(&self, request: EnqueueRequest) -> Result<()> {
261        let mut state = self.inner.lock().unwrap();
262        if state.kv.sealed {
263            return Err(Error::EffectsSealed);
264        }
265        state.enqueues.push(request);
266        Ok(())
267    }
268
269    /// Stage a write of `value` under `key` in the caller KV namespace.
270    ///
271    /// # Errors
272    ///
273    /// As [`EffectsHandle::put`].
274    pub fn put(&self, key: impl Into<Vec<u8>>, value: impl Into<Vec<u8>>) -> Result<()> {
275        self.inner.lock().unwrap().kv.put(key.into(), value.into())
276    }
277
278    /// Stage a delete of `key` from the caller KV namespace.
279    ///
280    /// # Errors
281    ///
282    /// As [`EffectsHandle::delete`].
283    pub fn delete(&self, key: impl Into<Vec<u8>>) -> Result<()> {
284        self.inner.lock().unwrap().kv.delete(key.into())
285    }
286
287    /// Seal the handle and move out everything staged.
288    pub(crate) fn seal_and_take(&self) -> (StagedEffects, Vec<EnqueueRequest>) {
289        let mut state = self.inner.lock().unwrap();
290        let sealed = state.kv.seal_and_take();
291        (sealed.outcome, std::mem::take(&mut state.enqueues))
292    }
293}
294
295fn display_key(key: &[u8]) -> String {
296    String::from_utf8_lossy(key).into_owned()
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn staging_validates_keys_values_and_conflicts() {
305        let handle = EffectsHandle::detached();
306        assert!(matches!(
307            handle.put("workflow/x", "v"),
308            Err(Error::ReservedKvKey(_))
309        ));
310        assert!(matches!(
311            handle.delete("workflow/x"),
312            Err(Error::ReservedKvKey(_))
313        ));
314        let oversized = vec![0u8; taquba::MAX_KV_VALUE_SIZE + 1];
315        assert!(matches!(
316            handle.put("k", oversized),
317            Err(Error::Queue(taquba::Error::KvValueTooLarge { .. }))
318        ));
319        handle.put("a", "v").unwrap();
320        assert!(matches!(
321            handle.delete("a"),
322            Err(Error::ConflictingKvEffect(_))
323        ));
324        handle.delete("b").unwrap();
325        assert!(matches!(
326            handle.put("b", "v"),
327            Err(Error::ConflictingKvEffect(_))
328        ));
329    }
330
331    #[test]
332    fn a_sealed_handle_rejects_staging() {
333        let handle = EffectsHandle::for_delivery();
334        let clone = handle.clone();
335        handle.put("a", "v").unwrap();
336        let sealed = handle.seal_and_take();
337        assert_eq!(sealed.outcome.writes.len(), 1);
338        assert!(matches!(clone.put("b", "v"), Err(Error::EffectsSealed)));
339        assert!(matches!(
340            clone.put_reserved_on_failure(b"workflow/x".to_vec(), b"v".to_vec()),
341            Err(Error::EffectsSealed)
342        ));
343        assert!(matches!(clone.delete("b"), Err(Error::EffectsSealed)));
344    }
345
346    #[test]
347    fn the_terminal_handle_applies_the_staging_and_seal_rules() {
348        let handle = TerminalEffects::for_delivery();
349        assert!(matches!(
350            handle.put("workflow/x", "v"),
351            Err(Error::ReservedKvKey(_))
352        ));
353        assert!(matches!(
354            handle.delete("workflow/x"),
355            Err(Error::ReservedKvKey(_))
356        ));
357        handle.put("a", "v").unwrap();
358        assert!(matches!(
359            handle.delete("a"),
360            Err(Error::ConflictingKvEffect(_))
361        ));
362        handle.delete("b").unwrap();
363        handle
364            .enqueue(taquba::EnqueueRequest {
365                queue: "side".to_string(),
366                payload: Vec::new(),
367                options: Default::default(),
368            })
369            .unwrap();
370        let (staged, enqueues) = handle.seal_and_take();
371        assert_eq!(staged.writes.len(), 1);
372        assert_eq!(staged.deletes.len(), 1);
373        assert_eq!(enqueues.len(), 1);
374        assert!(matches!(handle.put("c", "v"), Err(Error::EffectsSealed)));
375        assert!(matches!(handle.delete("c"), Err(Error::EffectsSealed)));
376        assert!(matches!(
377            handle.enqueue(taquba::EnqueueRequest {
378                queue: "side".to_string(),
379                payload: Vec::new(),
380                options: Default::default(),
381            }),
382            Err(Error::EffectsSealed)
383        ));
384    }
385
386    #[test]
387    fn clones_share_one_accumulator() {
388        let handle = EffectsHandle::for_delivery();
389        let clone = handle.clone();
390        clone.put("a", "v").unwrap();
391        clone.delete("b").unwrap();
392        let staged = handle.seal_and_take().outcome;
393        assert_eq!(staged.writes.get(b"a".as_slice()), Some(&b"v".to_vec()));
394        assert!(staged.deletes.contains(b"b".as_slice()));
395    }
396}
taquba_workflow/effects.rs

taquba_workflow/
effects.rs
