imagegen_bridge_runtime/
idempotency.rs

1//! Scoped, bounded in-memory idempotency coordination.
2
3use std::{
4    collections::HashMap,
5    io,
6    sync::Arc,
7    time::{Duration, Instant},
8};
9
10use imagegen_bridge_core::{BridgeError, ErrorCode, ImageRequest, ImageResponse};
11use sha2::{Digest, Sha256};
12use tokio::sync::{Mutex, watch};
13use tokio_util::sync::CancellationToken;
14
15const MAX_SCOPE_BYTES: usize = 256;
16
17/// Bounds for replay records and abandoned in-flight operations.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct IdempotencyConfig {
20    /// Maximum keys retained at once.
21    pub max_entries: usize,
22    /// Maximum serialized response bytes retained for replay.
23    pub max_completed_bytes: usize,
24    /// Completed response replay lifetime.
25    pub completed_ttl: Duration,
26    /// Safety lifetime for a leader that disappears without completing.
27    pub in_flight_ttl: Duration,
28    /// Retention for a provider operation whose outcome is unknown.
29    pub unknown_ttl: Duration,
30}
31
32impl Default for IdempotencyConfig {
33    fn default() -> Self {
34        Self {
35            max_entries: 10_000,
36            max_completed_bytes: 256 * 1024 * 1024,
37            completed_ttl: Duration::from_secs(24 * 60 * 60),
38            in_flight_ttl: Duration::from_secs(31 * 60),
39            unknown_ttl: Duration::from_secs(24 * 60 * 60),
40        }
41    }
42}
43
44pub(crate) struct IdempotencyCoordinator {
45    inner: Mutex<IdempotencyState>,
46    config: IdempotencyConfig,
47}
48
49impl IdempotencyCoordinator {
50    pub(crate) fn new(config: IdempotencyConfig) -> Result<Self, BridgeError> {
51        if config.max_entries == 0
52            || config.max_completed_bytes == 0
53            || config.completed_ttl.is_zero()
54            || config.in_flight_ttl.is_zero()
55            || config.unknown_ttl.is_zero()
56        {
57            return Err(BridgeError::new(
58                ErrorCode::Configuration,
59                "idempotency limits and retention periods must be greater than zero",
60            ));
61        }
62        Ok(Self {
63            inner: Mutex::new(IdempotencyState::default()),
64            config,
65        })
66    }
67
68    pub(crate) async fn begin(
69        &self,
70        scope: &str,
71        key: &str,
72        request: &ImageRequest,
73    ) -> Result<IdempotencyAction, BridgeError> {
74        validate_scope(scope)?;
75        let fingerprint = request_fingerprint(request)?;
76        let record_key = RecordKey {
77            scope: scope.to_owned(),
78            key: key.to_owned(),
79        };
80        let now = Instant::now();
81        let mut state = self.inner.lock().await;
82        state.cleanup(now, self.config.in_flight_ttl);
83        state.access_clock = state.access_clock.saturating_add(1);
84        let access = state.access_clock;
85        if let Some(entry) = state.entries.get_mut(&record_key) {
86            if entry.fingerprint != fingerprint {
87                return Err(BridgeError::new(
88                    ErrorCode::IdempotencyConflict,
89                    "idempotency key was already used for a different request",
90                ));
91            }
92            entry.last_access = access;
93            let receiver = entry.sender.subscribe();
94            let outcome = receiver.borrow().clone();
95            return match outcome {
96                EntryOutcome::Completed(response) => {
97                    Ok(IdempotencyAction::Cached(Box::new((*response).clone())))
98                }
99                EntryOutcome::Failed(error) | EntryOutcome::Unknown(error) => Err(error),
100                EntryOutcome::Pending => Ok(IdempotencyAction::Wait(receiver)),
101            };
102        }
103        state.make_capacity(self.config.max_entries)?;
104        let (sender, _receiver) = watch::channel(EntryOutcome::Pending);
105        state.entries.insert(
106            record_key.clone(),
107            Entry {
108                fingerprint,
109                created_at: now,
110                expires_at: None,
111                retained_bytes: 0,
112                uncertain: false,
113                last_access: access,
114                sender,
115            },
116        );
117        Ok(IdempotencyAction::Leader(IdempotencyToken {
118            record_key,
119            fingerprint,
120        }))
121    }
122
123    pub(crate) async fn complete(&self, token: IdempotencyToken, response: ImageResponse) {
124        let retained_bytes = response_weight(&response).unwrap_or(usize::MAX);
125        let mut state = self.inner.lock().await;
126        state.complete(&token, response, retained_bytes, &self.config);
127    }
128
129    pub(crate) async fn fail(&self, token: IdempotencyToken, error: BridgeError) {
130        let mut state = self.inner.lock().await;
131        if error
132            .details
133            .get("outcome")
134            .and_then(serde_json::Value::as_str)
135            == Some("unknown")
136        {
137            if let Some(entry) = state.entries.get_mut(&token.record_key)
138                && entry.fingerprint == token.fingerprint
139            {
140                entry.expires_at = Some(Instant::now() + self.config.unknown_ttl);
141                entry.uncertain = true;
142                entry.sender.send_replace(EntryOutcome::Unknown(
143                    error.retryable(false).with_detail("outcome", "unknown"),
144                ));
145            }
146        } else if state
147            .entries
148            .get(&token.record_key)
149            .is_some_and(|entry| entry.fingerprint == token.fingerprint)
150            && let Some(entry) = state.remove(&token.record_key)
151        {
152            entry.sender.send_replace(EntryOutcome::Failed(error));
153        }
154    }
155}
156
157pub(crate) enum IdempotencyAction {
158    Leader(IdempotencyToken),
159    Cached(Box<ImageResponse>),
160    Wait(watch::Receiver<EntryOutcome>),
161}
162
163impl IdempotencyAction {
164    pub(crate) async fn wait(
165        mut receiver: watch::Receiver<EntryOutcome>,
166        deadline: Instant,
167        cancellation: &CancellationToken,
168    ) -> Result<ImageResponse, BridgeError> {
169        loop {
170            let outcome = receiver.borrow().clone();
171            match outcome {
172                EntryOutcome::Completed(response) => return Ok((*response).clone()),
173                EntryOutcome::Failed(error) | EntryOutcome::Unknown(error) => return Err(error),
174                EntryOutcome::Pending => {}
175            }
176            if deadline <= Instant::now() {
177                return Err(BridgeError::new(
178                    ErrorCode::Timeout,
179                    "request deadline elapsed while waiting for idempotent result",
180                ));
181            }
182            tokio::select! {
183                changed = receiver.changed() => {
184                    if changed.is_err() {
185                        return Err(BridgeError::new(
186                            ErrorCode::Internal,
187                            "idempotency leader ended without a result",
188                        ));
189                    }
190                }
191                () = cancellation.cancelled() => {
192                    return Err(BridgeError::new(
193                        ErrorCode::Cancelled,
194                        "request was cancelled while waiting for idempotent result",
195                    ));
196                }
197                () = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {
198                    return Err(BridgeError::new(
199                        ErrorCode::Timeout,
200                        "request deadline elapsed while waiting for idempotent result",
201                    ));
202                }
203            }
204        }
205    }
206}
207
208#[derive(Clone)]
209pub(crate) enum EntryOutcome {
210    Pending,
211    Completed(Arc<ImageResponse>),
212    Failed(BridgeError),
213    Unknown(BridgeError),
214}
215
216pub(crate) struct IdempotencyToken {
217    record_key: RecordKey,
218    fingerprint: [u8; 32],
219}
220
221#[derive(Debug, Clone, PartialEq, Eq, Hash)]
222struct RecordKey {
223    scope: String,
224    key: String,
225}
226
227struct Entry {
228    fingerprint: [u8; 32],
229    created_at: Instant,
230    expires_at: Option<Instant>,
231    retained_bytes: usize,
232    uncertain: bool,
233    last_access: u64,
234    sender: watch::Sender<EntryOutcome>,
235}
236
237#[derive(Default)]
238struct IdempotencyState {
239    entries: HashMap<RecordKey, Entry>,
240    completed_bytes: usize,
241    access_clock: u64,
242}
243
244impl IdempotencyState {
245    fn cleanup(&mut self, now: Instant, in_flight_ttl: Duration) {
246        let stale: Vec<_> = self
247            .entries
248            .iter()
249            .filter_map(|(key, entry)| {
250                let expired = entry.expires_at.is_some_and(|expires| expires <= now)
251                    || (entry.expires_at.is_none()
252                        && now.saturating_duration_since(entry.created_at) >= in_flight_ttl);
253                expired.then(|| key.clone())
254            })
255            .collect();
256        for key in stale {
257            if let Some(entry) = self.remove(&key)
258                && entry.expires_at.is_none()
259            {
260                entry.sender.send_replace(EntryOutcome::Failed(
261                    BridgeError::new(ErrorCode::Timeout, "idempotent operation expired")
262                        .retryable(true),
263                ));
264            }
265        }
266    }
267
268    fn make_capacity(&mut self, maximum: usize) -> Result<(), BridgeError> {
269        if self.entries.len() < maximum {
270            return Ok(());
271        }
272        let oldest_completed = self
273            .entries
274            .iter()
275            .filter(|(_, entry)| entry.expires_at.is_some() && !entry.uncertain)
276            .min_by_key(|(_, entry)| entry.last_access)
277            .map(|(key, _)| key.clone());
278        if let Some(key) = oldest_completed {
279            self.remove(&key);
280            return Ok(());
281        }
282        Err(
283            BridgeError::new(ErrorCode::Overloaded, "idempotency capacity is exhausted")
284                .retryable(true),
285        )
286    }
287
288    fn complete(
289        &mut self,
290        token: &IdempotencyToken,
291        response: ImageResponse,
292        retained_bytes: usize,
293        config: &IdempotencyConfig,
294    ) {
295        if self
296            .entries
297            .get(&token.record_key)
298            .is_none_or(|entry| entry.fingerprint != token.fingerprint)
299        {
300            return;
301        }
302        if retained_bytes > config.max_completed_bytes
303            || self
304                .reserve_completed_bytes(
305                    retained_bytes,
306                    config.max_completed_bytes,
307                    &token.record_key,
308                )
309                .is_err()
310        {
311            if let Some(entry) = self.entries.get_mut(&token.record_key) {
312                entry.expires_at = Some(Instant::now() + config.completed_ttl);
313                entry.uncertain = true;
314                entry.retained_bytes = 0;
315                entry.sender.send_replace(EntryOutcome::Unknown(
316                    BridgeError::new(
317                        ErrorCode::Overloaded,
318                        "completed response is too large for idempotent replay",
319                    )
320                    .retryable(false)
321                    .with_detail("outcome", "completed_unreplayable"),
322                ));
323            }
324            return;
325        }
326        self.access_clock = self.access_clock.saturating_add(1);
327        if let Some(entry) = self.entries.get_mut(&token.record_key) {
328            entry.expires_at = Some(Instant::now() + config.completed_ttl);
329            entry.uncertain = false;
330            entry.retained_bytes = retained_bytes;
331            entry.last_access = self.access_clock;
332            self.completed_bytes = self.completed_bytes.saturating_add(retained_bytes);
333            entry
334                .sender
335                .send_replace(EntryOutcome::Completed(Arc::new(response)));
336        }
337    }
338
339    fn reserve_completed_bytes(
340        &mut self,
341        incoming: usize,
342        maximum: usize,
343        exclude: &RecordKey,
344    ) -> Result<(), BridgeError> {
345        while self.completed_bytes.saturating_add(incoming) > maximum {
346            let victim = self
347                .entries
348                .iter()
349                .filter(|(key, entry)| {
350                    *key != exclude && entry.retained_bytes > 0 && !entry.uncertain
351                })
352                .min_by_key(|(_, entry)| entry.last_access)
353                .map(|(key, _)| key.clone())
354                .ok_or_else(|| {
355                    BridgeError::new(
356                        ErrorCode::Overloaded,
357                        "idempotency completed-response byte capacity is exhausted",
358                    )
359                })?;
360            self.remove(&victim);
361        }
362        Ok(())
363    }
364
365    fn remove(&mut self, key: &RecordKey) -> Option<Entry> {
366        let entry = self.entries.remove(key)?;
367        self.completed_bytes = self.completed_bytes.saturating_sub(entry.retained_bytes);
368        Some(entry)
369    }
370}
371
372fn response_weight(response: &ImageResponse) -> Result<usize, BridgeError> {
373    struct Counter(usize);
374
375    impl io::Write for Counter {
376        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
377            self.0 = self.0.checked_add(bytes.len()).ok_or_else(|| {
378                io::Error::new(io::ErrorKind::FileTooLarge, "response weight overflowed")
379            })?;
380            Ok(bytes.len())
381        }
382
383        fn flush(&mut self) -> io::Result<()> {
384            Ok(())
385        }
386    }
387
388    let mut counter = Counter(256);
389    serde_json::to_writer(&mut counter, response).map_err(|_| {
390        BridgeError::new(
391            ErrorCode::Internal,
392            "could not measure idempotency response weight",
393        )
394    })?;
395    Ok(counter.0)
396}
397
398fn request_fingerprint(request: &ImageRequest) -> Result<[u8; 32], BridgeError> {
399    let mut request = request.clone();
400    request.idempotency_key = None;
401    request.timeout_ms = None;
402    let encoded = serde_json::to_vec(&request).map_err(|_| {
403        BridgeError::new(
404            ErrorCode::Internal,
405            "could not fingerprint the normalized request",
406        )
407    })?;
408    Ok(Sha256::digest(encoded).into())
409}
410
411fn validate_scope(scope: &str) -> Result<(), BridgeError> {
412    if scope.is_empty() || scope.len() > MAX_SCOPE_BYTES || scope.chars().any(char::is_control) {
413        Err(BridgeError::new(
414            ErrorCode::InvalidRequest,
415            "invalid idempotency scope",
416        ))
417    } else {
418        Ok(())
419    }
420}
421
422#[cfg(test)]
423mod tests {
424    #![allow(clippy::manual_let_else, clippy::panic, clippy::unwrap_used)]
425
426    use imagegen_bridge_core::{GenerationParameters, Timings};
427
428    use super::*;
429
430    fn response(id: &str) -> ImageResponse {
431        ImageResponse {
432            id: id.to_owned(),
433            created: 0,
434            provider: "fake".to_owned(),
435            model: "fake".to_owned(),
436            requested: GenerationParameters::default(),
437            effective: GenerationParameters::default(),
438            normalizations: Vec::new(),
439            attempts: Vec::new(),
440            data: Vec::new(),
441            failures: Vec::new(),
442            revised_prompt: None,
443            usage: None,
444            session: None,
445            timings: Timings::default(),
446            warnings: Vec::new(),
447        }
448    }
449
450    #[tokio::test]
451    async fn replays_completed_response_and_rejects_conflicts() {
452        let coordinator = IdempotencyCoordinator::new(IdempotencyConfig::default()).unwrap();
453        let request = ImageRequest::generate("same");
454        let IdempotencyAction::Leader(token) =
455            coordinator.begin("tenant", "key", &request).await.unwrap()
456        else {
457            panic!("first caller must lead");
458        };
459        coordinator.complete(token, response("original")).await;
460        let IdempotencyAction::Cached(cached) =
461            coordinator.begin("tenant", "key", &request).await.unwrap()
462        else {
463            panic!("completed call must be cached");
464        };
465        assert_eq!(cached.id, "original");
466        let error = match coordinator
467            .begin("tenant", "key", &ImageRequest::generate("different"))
468            .await
469        {
470            Err(error) => error,
471            Ok(_) => panic!("a conflicting request must fail"),
472        };
473        assert_eq!(error.code, ErrorCode::IdempotencyConflict);
474    }
475
476    #[tokio::test]
477    async fn followers_receive_the_leader_result() {
478        let coordinator =
479            Arc::new(IdempotencyCoordinator::new(IdempotencyConfig::default()).unwrap());
480        let request = ImageRequest::generate("same");
481        let IdempotencyAction::Leader(token) =
482            coordinator.begin("tenant", "key", &request).await.unwrap()
483        else {
484            panic!("first caller must lead");
485        };
486        let IdempotencyAction::Wait(receiver) =
487            coordinator.begin("tenant", "key", &request).await.unwrap()
488        else {
489            panic!("second caller must wait");
490        };
491        let follower = tokio::spawn(async move {
492            IdempotencyAction::wait(
493                receiver,
494                Instant::now() + Duration::from_secs(5),
495                &CancellationToken::new(),
496            )
497            .await
498        });
499        coordinator.complete(token, response("leader")).await;
500        assert_eq!(follower.await.unwrap().unwrap().id, "leader");
501    }
502
503    #[tokio::test]
504    async fn unknown_outcome_is_retained_and_cannot_elect_a_second_leader() {
505        let coordinator = IdempotencyCoordinator::new(IdempotencyConfig::default()).unwrap();
506        let request = ImageRequest::generate("same");
507        let IdempotencyAction::Leader(token) = coordinator
508            .begin("tenant", "unknown", &request)
509            .await
510            .unwrap()
511        else {
512            panic!("first caller must lead");
513        };
514        coordinator
515            .fail(
516                token,
517                BridgeError::new(ErrorCode::Timeout, "provider result is uncertain")
518                    .retryable(true)
519                    .with_detail("outcome", "unknown"),
520            )
521            .await;
522
523        let repeated = match coordinator.begin("tenant", "unknown", &request).await {
524            Err(error) => error,
525            Ok(_) => panic!("unknown outcome must remain a tombstone"),
526        };
527        assert_eq!(repeated.code, ErrorCode::Timeout);
528        assert!(!repeated.retryable);
529        assert_eq!(repeated.details["outcome"], "unknown");
530
531        let conflict = match coordinator
532            .begin("tenant", "unknown", &ImageRequest::generate("different"))
533            .await
534        {
535            Err(error) => error,
536            Ok(_) => panic!("different request must still conflict"),
537        };
538        assert_eq!(conflict.code, ErrorCode::IdempotencyConflict);
539    }
540
541    #[tokio::test]
542    async fn completed_response_bytes_are_bounded_and_oversize_keeps_a_tombstone() {
543        let mut sample = response("sample");
544        sample.revised_prompt = Some("x".repeat(700));
545        let weight = response_weight(&sample).unwrap();
546        let coordinator = IdempotencyCoordinator::new(IdempotencyConfig {
547            max_entries: 10,
548            max_completed_bytes: weight + 64,
549            ..IdempotencyConfig::default()
550        })
551        .unwrap();
552        let request = ImageRequest::generate("same");
553
554        let IdempotencyAction::Leader(first) = coordinator
555            .begin("tenant", "first", &request)
556            .await
557            .unwrap()
558        else {
559            panic!("first caller must lead");
560        };
561        coordinator.complete(first, sample.clone()).await;
562        let IdempotencyAction::Leader(second) = coordinator
563            .begin("tenant", "second", &request)
564            .await
565            .unwrap()
566        else {
567            panic!("second caller must lead");
568        };
569        coordinator.complete(second, sample).await;
570        let state = coordinator.inner.lock().await;
571        assert!(state.completed_bytes <= coordinator.config.max_completed_bytes);
572        assert_eq!(
573            state
574                .entries
575                .values()
576                .filter(|entry| entry.retained_bytes > 0)
577                .count(),
578            1
579        );
580        drop(state);
581
582        let mut oversized = response("oversized");
583        oversized.revised_prompt = Some("y".repeat(weight * 2));
584        let IdempotencyAction::Leader(oversize_token) = coordinator
585            .begin("tenant", "oversized", &request)
586            .await
587            .unwrap()
588        else {
589            panic!("oversized caller must lead once");
590        };
591        coordinator.complete(oversize_token, oversized).await;
592        let error = match coordinator.begin("tenant", "oversized", &request).await {
593            Err(error) => error,
594            Ok(_) => panic!("unreplayable completion must retain a tombstone"),
595        };
596        assert_eq!(error.details["outcome"], "completed_unreplayable");
597        assert!(!error.retryable);
598    }
599}
imagegen_bridge_runtime/idempotency.rs

imagegen_bridge_runtime/
idempotency.rs
