1use crate::durability::settle::append_settled;
2use std::collections::{HashMap, HashSet};
3use std::fmt;
4use std::sync::Arc;
5
6use super::{DurabilityError, DurableStore};
7
8mod codec;
9#[cfg(test)]
10mod tests;
11
12const READ_BATCH_SIZE: usize = 1_024;
13
14#[derive(Clone, Debug, PartialEq, Eq)]
16pub enum ConversationEvent {
17 MessageReceived {
19 message_id: String,
21 received_at: u64,
23 },
24 ProcessingStarted {
26 message_id: String,
28 },
29 StepCompleted {
31 message_id: String,
33 step_index: u32,
35 output: Vec<u8>,
37 },
38 ProcessingFinished {
40 message_id: String,
42 },
43 ErrorOccurred {
45 message_id: String,
47 error: String,
49 },
50}
51
52impl ConversationEvent {
53 #[must_use]
55 pub fn message_id(&self) -> &str {
56 match self {
57 Self::MessageReceived { message_id, .. }
58 | Self::ProcessingStarted { message_id }
59 | Self::StepCompleted { message_id, .. }
60 | Self::ProcessingFinished { message_id }
61 | Self::ErrorOccurred { message_id, .. } => message_id,
62 }
63 }
64}
65
66#[derive(Clone, Debug, Default, PartialEq, Eq)]
68pub struct ConversationState {
69 pub received_messages: HashSet<String>,
71 pub in_progress: HashSet<String>,
73 pub completed_steps: HashMap<(String, u32), Vec<u8>>,
75 pub finished_messages: HashSet<String>,
77 pub errored_messages: HashMap<String, String>,
79}
80
81impl ConversationState {
82 #[must_use]
84 pub fn replay(events: &[ConversationEvent]) -> Self {
85 let mut state = Self::default();
86 for event in events {
87 state.apply(event);
88 }
89 state
90 }
91
92 pub fn apply(&mut self, event: &ConversationEvent) {
94 match event {
95 ConversationEvent::MessageReceived { message_id, .. } => {
96 self.received_messages.insert(message_id.clone());
97 }
98 ConversationEvent::ProcessingStarted { message_id } => {
99 self.in_progress.insert(message_id.clone());
100 }
101 ConversationEvent::StepCompleted {
102 message_id,
103 step_index,
104 output,
105 } => {
106 self.completed_steps
107 .insert((message_id.clone(), *step_index), output.clone());
108 }
109 ConversationEvent::ProcessingFinished { message_id } => {
110 self.finished_messages.insert(message_id.clone());
111 self.in_progress.remove(message_id);
112 }
113 ConversationEvent::ErrorOccurred { message_id, error } => {
114 self.errored_messages
115 .insert(message_id.clone(), error.clone());
116 self.in_progress.remove(message_id);
117 }
118 }
119 }
120
121 #[must_use]
123 pub fn is_fully_processed(&self, message_id: &str) -> bool {
124 self.finished_messages.contains(message_id)
125 }
126
127 #[must_use]
129 pub fn last_completed_step(&self, message_id: &str) -> Option<u32> {
130 self.completed_steps
131 .keys()
132 .filter(|(stored_id, _)| stored_id.as_str() == message_id)
133 .map(|(_, step_index)| *step_index)
134 .max()
135 }
136
137 fn next_step_index(&self, message_id: &str) -> Result<u32, DurabilityError> {
138 self.last_completed_step(message_id).map_or(Ok(0), |step| {
139 step.checked_add(1).ok_or_else(|| {
140 DurabilityError::ConfigError("conversation step index overflow".to_owned())
141 })
142 })
143 }
144}
145
146#[derive(Clone, Debug, PartialEq, Eq)]
148pub enum RedeliveryDecision {
149 Skip,
151 ResumeFrom(u32),
153 Start,
155}
156
157#[derive(Clone)]
159pub struct DurableConversation {
160 conversation_id: String,
161 store: Arc<dyn DurableStore>,
162 state: ConversationState,
163 expected_seq: u64,
164}
165
166impl DurableConversation {
167 #[must_use]
169 pub fn new(conversation_id: impl Into<String>, store: Arc<dyn DurableStore>) -> Self {
170 Self {
171 conversation_id: conversation_id.into(),
172 store,
173 state: ConversationState::default(),
174 expected_seq: 0,
175 }
176 }
177
178 pub async fn recover(
184 conversation_id: impl Into<String>,
185 store: Arc<dyn DurableStore>,
186 ) -> Result<Self, DurabilityError> {
187 let conversation_id = conversation_id.into();
188 let (state, expected_seq) = replay_stream(store.as_ref(), &conversation_id).await?;
189 Ok(Self {
190 conversation_id,
191 store,
192 state,
193 expected_seq,
194 })
195 }
196
197 #[must_use]
199 pub fn conversation_id(&self) -> &str {
200 &self.conversation_id
201 }
202
203 #[must_use]
205 pub const fn state(&self) -> &ConversationState {
206 &self.state
207 }
208
209 #[must_use]
211 pub const fn expected_seq(&self) -> u64 {
212 self.expected_seq
213 }
214
215 pub async fn receive_message(
222 &mut self,
223 message_id: impl Into<String>,
224 received_at: u64,
225 ) -> Result<RedeliveryDecision, DurabilityError> {
226 let message_id = message_id.into();
227 if self.state.is_fully_processed(&message_id) {
228 return Ok(RedeliveryDecision::Skip);
229 }
230 if self.state.received_messages.contains(&message_id) {
231 return Ok(RedeliveryDecision::ResumeFrom(
232 self.state.next_step_index(&message_id)?,
233 ));
234 }
235 self.record_message_received(message_id, received_at)
236 .await?;
237 Ok(RedeliveryDecision::Start)
238 }
239
240 pub async fn record_message_received(
247 &mut self,
248 message_id: impl Into<String>,
249 received_at: u64,
250 ) -> Result<u64, DurabilityError> {
251 self.append_event(ConversationEvent::MessageReceived {
252 message_id: message_id.into(),
253 received_at,
254 })
255 .await
256 }
257
258 pub async fn record_processing_started(
265 &mut self,
266 message_id: impl Into<String>,
267 ) -> Result<u64, DurabilityError> {
268 self.append_event(ConversationEvent::ProcessingStarted {
269 message_id: message_id.into(),
270 })
271 .await
272 }
273
274 pub async fn record_step_completed(
281 &mut self,
282 message_id: impl Into<String>,
283 step_index: u32,
284 output: Vec<u8>,
285 ) -> Result<u64, DurabilityError> {
286 self.append_event(ConversationEvent::StepCompleted {
287 message_id: message_id.into(),
288 step_index,
289 output,
290 })
291 .await
292 }
293
294 pub async fn record_processing_finished(
301 &mut self,
302 message_id: impl Into<String>,
303 ) -> Result<u64, DurabilityError> {
304 self.append_event(ConversationEvent::ProcessingFinished {
305 message_id: message_id.into(),
306 })
307 .await
308 }
309
310 pub async fn record_error(
317 &mut self,
318 message_id: impl Into<String>,
319 error: impl Into<String>,
320 ) -> Result<u64, DurabilityError> {
321 self.append_event(ConversationEvent::ErrorOccurred {
322 message_id: message_id.into(),
323 error: error.into(),
324 })
325 .await
326 }
327
328 async fn append_event(&mut self, event: ConversationEvent) -> Result<u64, DurabilityError> {
329 let payload = event.serialize()?;
330 let assigned_seq = append_settled(
331 self.store.as_ref(),
332 &self.conversation_id,
333 payload,
334 self.expected_seq,
335 )
336 .await?;
337 self.expected_seq = assigned_seq.checked_add(1).ok_or_else(|| {
338 DurabilityError::ConfigError(
339 "sequence number overflow after conversation append".to_owned(),
340 )
341 })?;
342 self.state.apply(&event);
343 Ok(assigned_seq)
344 }
345}
346
347impl fmt::Debug for DurableConversation {
348 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
349 formatter
350 .debug_struct("DurableConversation")
351 .field("conversation_id", &self.conversation_id)
352 .field("state", &self.state)
353 .field("expected_seq", &self.expected_seq)
354 .field("store", &self.store)
355 .finish()
356 }
357}
358
359async fn replay_stream(
360 store: &dyn DurableStore,
361 conversation_id: &str,
362) -> Result<(ConversationState, u64), DurabilityError> {
363 let mut state = ConversationState::default();
364 let mut offset = 0;
365 let mut last_sequence = None;
366 loop {
367 let batch = store
368 .read_from(conversation_id, offset, READ_BATCH_SIZE)
369 .await?;
370 let batch_len = batch.len();
371 if batch_len == 0 {
372 break;
373 }
374 for stored in &batch {
375 let event = ConversationEvent::deserialize(&stored.payload)?;
376 state.apply(&event);
377 last_sequence = Some(stored.sequence);
378 }
379 offset = offset.checked_add(len_to_u64(batch_len)?).ok_or_else(|| {
380 DurabilityError::ConfigError("conversation read offset overflow".to_owned())
381 })?;
382 if batch_len < READ_BATCH_SIZE {
383 break;
384 }
385 }
386 let expected_seq = last_sequence.map_or(Ok(0), |sequence| {
387 sequence.checked_add(1).ok_or_else(|| {
388 DurabilityError::ConfigError(
389 "sequence number overflow after conversation replay".to_owned(),
390 )
391 })
392 })?;
393 Ok((state, expected_seq))
394}
395
396fn len_to_u64(len: usize) -> Result<u64, DurabilityError> {
397 u64::try_from(len).map_err(|error| {
398 DurabilityError::ConfigError(format!("conversation entry count cannot fit u64: {error}"))
399 })
400}