1use std::cmp::Ordering;
17use std::cmp::Reverse;
18use std::collections::BinaryHeap;
19use std::collections::HashMap;
20use std::collections::HashSet;
21use std::collections::VecDeque;
22
23use slotmap::Key;
24use slotmap::new_key_type;
25use ursula_shard::BucketStreamId;
26
27use self::cold_gc::ColdGcQueue;
28use self::cold_state::StreamColdState;
29use self::hot_buffer::HotBuffer;
30use self::registry::StreamRegistry;
31use self::ttl::TtlEntry;
32use self::ttl::TtlIndex;
33use crate::command::StreamCommand;
34use crate::integrity::StreamIntegrity;
35use crate::model::AppendExternalInput;
36use crate::model::AppendStreamInput;
37use crate::model::COLD_INDEX_PAGE_SPAN_BYTES;
38use crate::model::ColdChunkRef;
39use crate::model::ColdFlushCandidate;
40use crate::model::ColdGcEntry;
41use crate::model::ColdGcTarget;
42use crate::model::ExternalPayloadRef;
43use crate::model::HotPayloadSegment;
44use crate::model::MAX_STREAM_ATTRS_BYTES;
45use crate::model::ObjectPayloadRef;
46use crate::model::ProducerAppendRecord;
47use crate::model::ProducerRequest;
48use crate::model::ProducerSnapshot;
49use crate::model::ProducerState;
50use crate::model::StreamAttrs;
51use crate::model::StreamBatchAppend;
52use crate::model::StreamBatchAppendItem;
53use crate::model::StreamBootstrapPlan;
54use crate::model::StreamMessageRecord;
55use crate::model::StreamMetadata;
56use crate::model::StreamRead;
57use crate::model::StreamReadColdIndexSegment;
58use crate::model::StreamReadObjectSegment;
59use crate::model::StreamReadPlan;
60use crate::model::StreamReadSegment;
61use crate::model::StreamStatus;
62use crate::model::StreamVisibleSnapshot;
63use crate::response::StreamErrorCode;
64use crate::response::StreamErrorContext;
65use crate::response::StreamResponse;
66use crate::snapshot::StreamSnapshot;
67use crate::snapshot::StreamSnapshotEntry;
68use crate::snapshot::StreamSnapshotError;
69use crate::validate::validate_bucket_id;
70use crate::validate::validate_stream_id;
71
72mod append;
73mod cold;
74mod cold_gc;
75mod cold_state;
76mod hot_buffer;
77mod lifecycle;
78mod persist;
79mod query;
80mod registry;
81mod ttl;
82
83const TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE: usize = 256;
84
85new_key_type! {
86 struct StreamKey;
87}
88
89#[derive(Debug, Clone, Default)]
90pub struct StreamStateMachine {
91 buckets: HashSet<String>,
92 registry: StreamRegistry,
93 cold_gc: ColdGcQueue,
94}
95
96#[derive(Debug, Clone)]
97struct StreamSlot {
98 metadata: StreamMetadata,
99 attrs: Option<StreamAttrs>,
100 hot_buffer: HotBuffer,
101 cold: StreamColdState,
102 message_records: Vec<StreamMessageRecord>,
103 integrity: StreamIntegrity,
104 visible_snapshot: Option<StreamVisibleSnapshot>,
105 producers: HashMap<String, ProducerState>,
106}
107
108impl StreamStateMachine {
109 pub fn new() -> Self {
110 Self::default()
111 }
112
113 fn stream_slot(&self, stream_id: &BucketStreamId) -> Option<&StreamSlot> {
114 self.registry.slot(stream_id)
115 }
116
117 fn stream_slot_mut(&mut self, stream_id: &BucketStreamId) -> Option<&mut StreamSlot> {
118 self.registry.slot_mut(stream_id)
119 }
120
121 fn stream_metadata(&self, stream_id: &BucketStreamId) -> Option<&StreamMetadata> {
122 self.registry.metadata(stream_id)
123 }
124
125 fn stream_metadata_mut(&mut self, stream_id: &BucketStreamId) -> Option<&mut StreamMetadata> {
126 self.registry.metadata_mut(stream_id)
127 }
128
129 fn insert_stream_slot(&mut self, slot: StreamSlot) -> Option<StreamKey> {
130 self.registry.insert(slot)
131 }
132
133 fn refresh_ttl_entry(&mut self, stream_id: &BucketStreamId) {
134 self.registry.refresh_ttl(stream_id);
135 }
136
137 pub fn apply(&mut self, command: StreamCommand) -> StreamResponse {
138 match command {
139 StreamCommand::CreateBucket { bucket_id } => self.create_bucket(bucket_id),
140 StreamCommand::DeleteBucket { bucket_id } => self.delete_bucket(&bucket_id),
141 StreamCommand::CreateStream {
142 stream_id,
143 content_type,
144 initial_payload,
145 close_after,
146 stream_seq,
147 producer,
148 stream_ttl_seconds,
149 stream_expires_at_ms,
150 forked_from,
151 fork_offset,
152 attrs,
153 now_ms,
154 } => {
155 let response = self.create_stream(CreateStreamInput {
156 stream_id,
157 content_type,
158 initial_payload,
159 close_after,
160 stream_seq,
161 producer,
162 stream_ttl_seconds,
163 stream_expires_at_ms,
164 forked_from,
165 fork_offset,
166 attrs,
167 now_ms,
168 });
169 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
170 response
171 }
172 StreamCommand::CreateExternal {
173 stream_id,
174 content_type,
175 initial_payload,
176 close_after,
177 stream_seq,
178 producer,
179 stream_ttl_seconds,
180 stream_expires_at_ms,
181 forked_from,
182 fork_offset,
183 attrs,
184 now_ms,
185 } => {
186 let response = self.create_external_stream(CreateExternalStreamInput {
187 stream_id,
188 content_type,
189 initial_payload,
190 close_after,
191 stream_seq,
192 producer,
193 stream_ttl_seconds,
194 stream_expires_at_ms,
195 forked_from,
196 fork_offset,
197 attrs,
198 now_ms,
199 });
200 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
201 response
202 }
203 StreamCommand::Append {
204 stream_id,
205 content_type,
206 payload,
207 close_after,
208 stream_seq,
209 producer,
210 now_ms,
211 } => {
212 let response = self.append_borrowed(AppendStreamInput {
213 stream_id,
214 content_type: content_type.as_deref(),
215 payload: &payload,
216 close_after,
217 stream_seq,
218 producer,
219 now_ms,
220 });
221 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
222 response
223 }
224 StreamCommand::AppendExternal {
225 stream_id,
226 content_type,
227 payload,
228 close_after,
229 stream_seq,
230 producer,
231 now_ms,
232 } => {
233 let response = self.append_external(AppendExternalInput {
234 stream_id,
235 content_type: content_type.as_deref(),
236 payload,
237 close_after,
238 stream_seq,
239 producer,
240 now_ms,
241 });
242 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
243 response
244 }
245 StreamCommand::AppendBatch {
246 stream_id,
247 content_type,
248 payloads,
249 producer,
250 now_ms,
251 } => {
252 let response = match self.append_batch_borrowed(
253 stream_id,
254 content_type.as_deref(),
255 &payloads.iter().map(Vec::as_slice).collect::<Vec<_>>(),
256 producer,
257 now_ms,
258 ) {
259 Ok(batch) => batch
260 .items
261 .last()
262 .map(|item| StreamResponse::Appended {
263 offset: item.offset,
264 next_offset: item.next_offset,
265 closed: item.closed,
266 deduplicated: item.deduplicated,
267 producer: None,
268 })
269 .unwrap_or_else(|| {
270 StreamResponse::error(
271 StreamErrorCode::EmptyAppend,
272 "append batch must contain at least one payload",
273 )
274 }),
275 Err(response) => response,
276 };
277 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
278 response
279 }
280 StreamCommand::PublishSnapshot {
281 stream_id,
282 snapshot_offset,
283 content_type,
284 payload,
285 now_ms,
286 } => {
287 let response = self.publish_snapshot(
288 stream_id,
289 snapshot_offset,
290 content_type,
291 payload,
292 now_ms,
293 );
294 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
295 response
296 }
297 StreamCommand::TouchStreamAccess {
298 stream_id,
299 now_ms,
300 renew_ttl,
301 } => {
302 let response = self.touch_stream_access(&stream_id, now_ms, renew_ttl);
303 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
304 response
305 }
306 StreamCommand::UpdateStreamAttrs {
307 stream_id,
308 attrs,
309 now_ms,
310 } => {
311 let response = self.update_stream_attrs(&stream_id, attrs, now_ms);
312 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
313 response
314 }
315 StreamCommand::AddForkRef { stream_id, now_ms } => {
316 let response = self.add_fork_ref(&stream_id, now_ms);
317 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
318 response
319 }
320 StreamCommand::ReleaseForkRef { stream_id } => self.release_fork_ref(&stream_id),
321 StreamCommand::FlushCold { stream_id, chunk } => self.flush_cold(stream_id, chunk),
322 StreamCommand::Close {
323 stream_id,
324 stream_seq,
325 producer,
326 now_ms,
327 } => {
328 let response = self.close(stream_id, stream_seq, producer, now_ms);
329 self.sweep_expired_streams(now_ms, TTL_EXPIRY_SWEEP_MAX_STREAMS_PER_WRITE);
330 response
331 }
332 StreamCommand::DeleteStream { stream_id } => self.delete_stream(&stream_id),
333 StreamCommand::AckColdGc { up_to_seq } => self.ack_cold_gc(up_to_seq),
334 }
335 }
336}
337
338#[derive(Debug)]
339struct CreateStreamInput {
340 stream_id: BucketStreamId,
341 content_type: String,
342 initial_payload: Vec<u8>,
343 close_after: bool,
344 stream_seq: Option<String>,
345 producer: Option<ProducerRequest>,
346 stream_ttl_seconds: Option<u64>,
347 stream_expires_at_ms: Option<u64>,
348 forked_from: Option<BucketStreamId>,
349 fork_offset: Option<u64>,
350 attrs: Option<StreamAttrs>,
351 now_ms: u64,
352}
353
354#[derive(Debug)]
355struct CreateExternalStreamInput {
356 stream_id: BucketStreamId,
357 content_type: String,
358 initial_payload: ExternalPayloadRef,
359 close_after: bool,
360 stream_seq: Option<String>,
361 producer: Option<ProducerRequest>,
362 stream_ttl_seconds: Option<u64>,
363 stream_expires_at_ms: Option<u64>,
364 forked_from: Option<BucketStreamId>,
365 fork_offset: Option<u64>,
366 attrs: Option<StreamAttrs>,
367 now_ms: u64,
368}
369
370impl CreateStreamInput {
371 fn initial_len(&self) -> u64 {
372 u64::try_from(self.initial_payload.len()).expect("payload len fits u64")
373 }
374}
375
376fn is_soft_deleted(stream: &StreamMetadata) -> bool {
377 stream.status == StreamStatus::SoftDeleted
378}
379
380fn normalize_stream_attrs(attrs: Option<StreamAttrs>) -> Option<StreamAttrs> {
381 attrs.filter(|attrs| !attrs.is_empty())
382}
383
384fn stream_expiry_at_ms(stream: &StreamMetadata) -> Option<u64> {
385 if let Some(expires_at_ms) = stream.stream_expires_at_ms {
386 return Some(expires_at_ms);
387 }
388 stream.stream_ttl_seconds.map(|ttl_seconds| {
389 stream
390 .last_ttl_touch_at_ms
391 .saturating_add(ttl_seconds.saturating_mul(1000))
392 })
393}
394
395fn stream_is_expired(stream: &StreamMetadata, now_ms: u64) -> bool {
396 stream_expiry_at_ms(stream).is_some_and(|expires_at_ms| now_ms >= expires_at_ms)
397}
398
399fn renew_stream_ttl(stream: &mut StreamMetadata, now_ms: u64) {
400 if stream.stream_ttl_seconds.is_some() && stream.stream_expires_at_ms.is_none() {
401 stream.last_ttl_touch_at_ms = now_ms;
402 }
403}
404
405fn validate_producer_request(producer: Option<&ProducerRequest>) -> Result<(), StreamResponse> {
406 let Some(producer) = producer else {
407 return Ok(());
408 };
409 if producer.producer_id.trim().is_empty() {
410 return Err(StreamResponse::error(
411 StreamErrorCode::InvalidProducer,
412 "producer id must not be empty",
413 ));
414 }
415 const MAX_JS_SAFE_INTEGER: u64 = 9_007_199_254_740_991;
416 if producer.producer_epoch > MAX_JS_SAFE_INTEGER {
417 return Err(StreamResponse::error(
418 StreamErrorCode::InvalidProducer,
419 format!(
420 "producer epoch {} exceeds maximum {}",
421 producer.producer_epoch, MAX_JS_SAFE_INTEGER
422 ),
423 ));
424 }
425 if producer.producer_seq > MAX_JS_SAFE_INTEGER {
426 return Err(StreamResponse::error(
427 StreamErrorCode::InvalidProducer,
428 format!(
429 "producer sequence {} exceeds maximum {}",
430 producer.producer_seq, MAX_JS_SAFE_INTEGER
431 ),
432 ));
433 }
434 Ok(())
435}
436
437fn validate_external_payload_ref(payload: &ExternalPayloadRef) -> Result<(), StreamResponse> {
438 if payload.s3_path.trim().is_empty() {
439 return Err(StreamResponse::error(
440 StreamErrorCode::InvalidColdFlush,
441 "external payload S3 path must not be empty",
442 ));
443 }
444 if payload.payload_len == 0 {
445 return Err(StreamResponse::error(
446 StreamErrorCode::EmptyAppend,
447 "external payload length must be greater than zero",
448 ));
449 }
450 if payload.object_size < payload.payload_len {
451 return Err(StreamResponse::error(
452 StreamErrorCode::InvalidColdFlush,
453 "external payload object size must cover payload length",
454 ));
455 }
456 Ok(())
457}
458
459fn compare_stream_ids(left: &BucketStreamId, right: &BucketStreamId) -> std::cmp::Ordering {
460 left.bucket_id
461 .cmp(&right.bucket_id)
462 .then_with(|| left.stream_id.cmp(&right.stream_id))
463}
464
465#[cfg(test)]
466mod tests;