1use std::collections::{HashMap, VecDeque};
8use std::sync::atomic::{AtomicUsize, Ordering};
9use std::sync::Arc;
10
11use arrow_array::RecordBatch;
12use laminar_core::checkpoint::CheckpointAttempt;
13use parking_lot::Mutex;
14use tokio::sync::watch;
15
16mod lifecycle;
17mod reader;
18
19pub(crate) use lifecycle::{
20 process_subscription_log_bytes, SubscriptionRegistry, PROCESS_SUBSCRIPTION_BYTES,
21};
22pub(super) use reader::SubscriptionRead;
23pub(crate) use reader::SubscriptionReader;
24#[cfg(test)]
25use reader::TryRead;
26
27pub(super) const MAX_LIVE_BATCH_BYTES: usize = 8 * 1024 * 1024;
28
29const INTERNAL_LIVE_LOG_BYTES: usize = MAX_LIVE_BATCH_BYTES * 2;
33const ENTRY_ACCOUNTING_BYTES: usize = 64;
34const BARRIER_ENTRY_BYTES: usize = ENTRY_ACCOUNTING_BYTES + std::mem::size_of::<u64>() * 3;
35
36struct SubscriptionMemoryBudget {
37 limit: usize,
38 used: AtomicUsize,
39}
40
41struct ChargedUpdateInner {
42 update: Option<MvUpdate>,
43 bytes: usize,
44 budget: Arc<SubscriptionMemoryBudget>,
45}
46
47impl Drop for ChargedUpdateInner {
48 fn drop(&mut self) {
49 {
50 let _update = self.update.take();
51 }
52 self.budget.release(self.bytes);
53 }
54}
55
56#[derive(Clone)]
57pub(super) struct ChargedUpdate(Arc<ChargedUpdateInner>);
58
59impl ChargedUpdate {
60 fn from_reserved(
62 update: MvUpdate,
63 bytes: usize,
64 budget: Arc<SubscriptionMemoryBudget>,
65 ) -> Self {
66 Self(Arc::new(ChargedUpdateInner {
67 update: Some(update),
68 bytes,
69 budget,
70 }))
71 }
72
73 pub(super) fn as_ref(&self) -> &MvUpdate {
74 self.0
75 .update
76 .as_ref()
77 .expect("charged subscription update is present until final drop")
78 }
79}
80
81impl SubscriptionMemoryBudget {
82 fn new(limit: usize) -> Self {
83 Self {
84 limit,
85 used: AtomicUsize::new(0),
86 }
87 }
88
89 fn try_reserve(&self, bytes: usize) -> bool {
90 self.used
91 .try_update(Ordering::AcqRel, Ordering::Acquire, |used| {
92 (bytes <= self.limit.saturating_sub(used)).then(|| used.saturating_add(bytes))
93 })
94 .is_ok()
95 }
96
97 fn release(&self, bytes: usize) {
98 if bytes == 0 {
99 return;
100 }
101 let released = self
102 .used
103 .try_update(Ordering::AcqRel, Ordering::Acquire, |used| {
104 used.checked_sub(bytes)
105 });
106 debug_assert!(released.is_ok(), "subscription memory released twice");
107 }
108
109 fn used(&self) -> usize {
110 self.used.load(Ordering::Acquire)
111 }
112}
113
114#[derive(Debug)]
115pub(crate) enum MvUpdate {
116 Batch(RecordBatch),
117 Barrier {
118 epoch: u64,
119 checkpoint_id: u64,
120 through_sequence: u64,
122 },
123 Error(String),
124}
125
126#[derive(Clone, Copy, Debug)]
128pub enum SubscribeStart {
129 Tail,
131 AsOfEpoch(u64),
133 AfterSequence(u64),
138}
139
140#[derive(Debug, Clone, Copy)]
141pub(crate) enum SubscriptionOpenError {
142 ReplayPruned {
143 earliest_retained: u64,
145 },
146 EpochNotCommitted {
147 requested: u64,
148 latest_committed: Option<u64>,
149 },
150 SequencePruned {
151 requested: u64,
152 earliest_retained: Option<u64>,
153 },
154 SequenceNotPublished {
155 requested: u64,
156 next_sequence: u64,
157 },
158 Capacity {
159 attached: usize,
160 limit: usize,
161 },
162}
163
164struct SequencedUpdate {
165 sequence: u64,
166 bytes: usize,
167 update: ChargedUpdate,
168}
169
170struct StreamLog {
171 inner: Mutex<StreamLogInner>,
172 wake: watch::Sender<()>,
173 budget: Arc<SubscriptionMemoryBudget>,
174}
175
176struct StreamLogInner {
177 entries: VecDeque<SequencedUpdate>,
178 bytes: usize,
179 retention_cap: usize,
180 retention_floor: u64,
181 retention_bytes: usize,
182 next_sequence: u64,
183 next_reader_id: u64,
184 readers: HashMap<u64, u64>,
185 latest_committed_epoch: Option<u64>,
186 terminal_error: Option<String>,
187 reserved_marker: Option<CheckpointAttempt>,
188}
189
190#[must_use]
191enum AppendOutcome {
192 Stored,
193 Unobserved,
194 Rejected(String),
195}
196
197#[must_use]
198enum ReservedAppendOutcome {
199 Stored,
200 Unobserved,
201 Rejected(String),
202}
203
204impl StreamLog {
205 fn new(
206 retention_cap: usize,
207 budget: Arc<SubscriptionMemoryBudget>,
208 latest_committed_epoch: Option<u64>,
209 ) -> Self {
210 Self::new_at(retention_cap, budget, latest_committed_epoch, 0)
211 }
212
213 fn new_at(
214 retention_cap: usize,
215 budget: Arc<SubscriptionMemoryBudget>,
216 latest_committed_epoch: Option<u64>,
217 next_sequence: u64,
218 ) -> Self {
219 let (wake, _) = watch::channel(());
220 Self {
221 inner: Mutex::new(StreamLogInner {
222 entries: VecDeque::new(),
223 bytes: 0,
224 retention_cap,
225 retention_floor: next_sequence,
226 retention_bytes: 0,
227 next_sequence,
228 next_reader_id: 0,
229 readers: HashMap::new(),
230 latest_committed_epoch,
231 terminal_error: None,
232 reserved_marker: None,
233 }),
234 wake,
235 budget,
236 }
237 }
238
239 fn append(&self, update: MvUpdate) -> AppendOutcome {
240 let entry_bytes = approx_size(&update);
241 let mut inner = self.inner.lock();
242 debug_assert!(!matches!(&update, MvUpdate::Barrier { .. }));
243 if let Some(message) = &inner.terminal_error {
244 return AppendOutcome::Rejected(message.clone());
245 }
246 if inner.retention_cap == 0 && inner.readers.is_empty() {
247 return AppendOutcome::Unobserved;
248 }
249
250 let storage_cap = storage_cap(inner.retention_cap);
251 if entry_bytes > storage_cap {
252 let message = format!(
253 "subscription entry requires {entry_bytes} bytes, exceeding the internal shared-log limit of {storage_cap} bytes"
254 );
255 return self.reject_append(inner, message);
256 }
257
258 let required_sequence_space = if inner.reserved_marker.is_some() {
259 2
260 } else {
261 1
262 };
263 if inner
264 .next_sequence
265 .checked_add(required_sequence_space)
266 .is_none()
267 {
268 return self.reject_append(inner, "subscription sequence space exhausted".into());
269 }
270 let Some(next_sequence) = inner.next_sequence.checked_add(1) else {
271 unreachable!("sequence-space preflight admitted the next sequence");
272 };
273
274 reclaim_consumed_prefix(&mut inner);
275 let available_locally = storage_cap.saturating_sub(entry_bytes);
276 evict_to_byte_target(&mut inner, available_locally);
277
278 if !self.budget.try_reserve(entry_bytes) {
279 let message = format!(
280 "subscription process memory budget exhausted ({} byte limit); update was not delivered",
281 self.budget.limit
282 );
283 return self.reject_append(inner, message);
284 }
285
286 let sequence = inner.next_sequence;
287 inner.entries.push_back(SequencedUpdate {
288 sequence,
289 bytes: entry_bytes,
290 update: ChargedUpdate::from_reserved(update, entry_bytes, Arc::clone(&self.budget)),
291 });
292 inner.next_sequence = next_sequence;
293 inner.bytes = inner.bytes.saturating_add(entry_bytes);
294 retain_appended_entry(&mut inner, sequence, entry_bytes);
295 reclaim_consumed_prefix(&mut inner);
296 drop(inner);
297 self.wake.send_modify(|()| {});
298 AppendOutcome::Stored
299 }
300
301 fn reject_append(
302 &self,
303 mut inner: parking_lot::MutexGuard<'_, StreamLogInner>,
304 message: String,
305 ) -> AppendOutcome {
306 if inner.reserved_marker.is_some() {
307 return AppendOutcome::Rejected(message);
308 }
309 Self::terminate_locked(&mut inner, message.clone());
310 drop(inner);
311 self.wake.send_modify(|()| {});
312 AppendOutcome::Rejected(message)
313 }
314
315 fn reserve_marker(&self, attempt: CheckpointAttempt) -> Result<Option<u64>, String> {
316 let mut inner = self.inner.lock();
317 if inner.terminal_error.is_some() || (inner.retention_cap == 0 && inner.readers.is_empty())
318 {
319 return Ok(None);
320 }
321 if let Some(existing) = inner.reserved_marker {
322 return Err(format!(
323 "subscription marker already reserved for epoch={} checkpoint_id={}",
324 existing.epoch, existing.checkpoint_id
325 ));
326 }
327 if inner.next_sequence.checked_add(1).is_none() {
328 return Err("subscription sequence space cannot admit a checkpoint marker".into());
329 }
330 let through_sequence = inner.next_sequence;
331 inner.reserved_marker = Some(attempt);
332 Ok(Some(through_sequence))
333 }
334
335 fn cancel_marker(&self, attempt: CheckpointAttempt) -> bool {
336 let mut inner = self.inner.lock();
337 if inner.reserved_marker != Some(attempt) {
338 return false;
339 }
340 inner.reserved_marker = None;
341 true
342 }
343
344 fn append_reserved_marker(
345 &self,
346 attempt: CheckpointAttempt,
347 through_sequence: u64,
348 ) -> ReservedAppendOutcome {
349 let mut inner = self.inner.lock();
350 if inner.reserved_marker != Some(attempt) {
351 return ReservedAppendOutcome::Rejected(format!(
352 "subscription marker reservation mismatch for epoch={} checkpoint_id={}",
353 attempt.epoch, attempt.checkpoint_id
354 ));
355 }
356 inner.reserved_marker = None;
357
358 if let Some(message) = &inner.terminal_error {
359 return ReservedAppendOutcome::Rejected(format!(
360 "subscription log terminated before checkpoint marker publication: {message}"
361 ));
362 }
363 inner.latest_committed_epoch = Some(
364 inner
365 .latest_committed_epoch
366 .map_or(attempt.epoch, |latest| latest.max(attempt.epoch)),
367 );
368 if inner.retention_cap == 0 && inner.readers.is_empty() {
369 return ReservedAppendOutcome::Unobserved;
370 }
371 if through_sequence > inner.next_sequence {
372 return ReservedAppendOutcome::Rejected(format!(
373 "subscription checkpoint cut cursor {through_sequence} is ahead of log cursor {}",
374 inner.next_sequence
375 ));
376 }
377 let Some(next_sequence) = inner.next_sequence.checked_add(1) else {
378 return ReservedAppendOutcome::Rejected(
379 "subscription sequence space exhausted before checkpoint marker publication".into(),
380 );
381 };
382
383 let storage_cap = storage_cap(inner.retention_cap);
384 if BARRIER_ENTRY_BYTES > storage_cap {
385 return ReservedAppendOutcome::Rejected(format!(
386 "subscription checkpoint marker requires {BARRIER_ENTRY_BYTES} bytes, exceeding the internal shared-log limit of {storage_cap} bytes"
387 ));
388 }
389 reclaim_consumed_prefix(&mut inner);
390 evict_to_byte_target(&mut inner, storage_cap.saturating_sub(BARRIER_ENTRY_BYTES));
391
392 let sequence = inner.next_sequence;
393 inner.entries.push_back(SequencedUpdate {
394 sequence,
395 bytes: BARRIER_ENTRY_BYTES,
396 update: ChargedUpdate::from_reserved(
397 MvUpdate::Barrier {
398 epoch: attempt.epoch,
399 checkpoint_id: attempt.checkpoint_id,
400 through_sequence,
401 },
402 BARRIER_ENTRY_BYTES,
403 Arc::clone(&self.budget),
404 ),
405 });
406 inner.next_sequence = next_sequence;
407 inner.bytes = inner.bytes.saturating_add(BARRIER_ENTRY_BYTES);
408 retain_appended_entry(&mut inner, sequence, BARRIER_ENTRY_BYTES);
409 reclaim_consumed_prefix(&mut inner);
410 drop(inner);
411 self.wake.send_modify(|()| {});
412 ReservedAppendOutcome::Stored
413 }
414
415 fn subscribe(
416 self: &Arc<Self>,
417 start: SubscribeStart,
418 ) -> Result<SubscriptionReader, SubscriptionOpenError> {
419 let mut inner = self.inner.lock();
420 let attached = inner.readers.len();
421 if attached >= super::MAX_SUBSCRIBERS_PER_MV {
422 return Err(SubscriptionOpenError::Capacity {
423 attached,
424 limit: super::MAX_SUBSCRIBERS_PER_MV,
425 });
426 }
427
428 let (cursor, skip_barrier) = match (inner.terminal_error.is_some(), start) {
429 (true, _) | (false, SubscribeStart::Tail) => (inner.next_sequence, None),
430 (false, SubscribeStart::AsOfEpoch(epoch)) => {
431 let (cursor, barrier_sequence) = cursor_after_retained_epoch(&inner, epoch)?;
432 (cursor, Some((epoch, barrier_sequence)))
433 }
434 (false, SubscribeStart::AfterSequence(sequence)) => {
435 (cursor_after_sequence(&inner, sequence)?, None)
436 }
437 };
438 let wake = self.wake.subscribe();
439 let mut reader_id = inner.next_reader_id;
440 while inner.readers.contains_key(&reader_id) {
441 reader_id = reader_id.wrapping_add(1);
442 }
443 inner.next_reader_id = reader_id.wrapping_add(1);
444 inner.readers.insert(reader_id, cursor);
445 drop(inner);
446
447 Ok(SubscriptionReader::attached(
448 Arc::clone(self),
449 reader_id,
450 cursor,
451 skip_barrier,
452 wake,
453 ))
454 }
455
456 fn set_retention_cap(&self, retention_cap: usize) {
457 let mut inner = self.inner.lock();
458 let previous_head = head_sequence(&inner);
459 inner.retention_cap = retention_cap;
460 recompute_retention_suffix(&mut inner);
461 reclaim_consumed_prefix(&mut inner);
462 evict_to_byte_target(&mut inner, storage_cap(retention_cap));
463 let head_changed = head_sequence(&inner) != previous_head;
464 drop(inner);
465 if head_changed {
466 self.wake.send_modify(|()| {});
467 }
468 }
469
470 fn terminate(&self, message: &str) {
471 let mut inner = self.inner.lock();
472 Self::terminate_locked(&mut inner, message.to_owned());
473 drop(inner);
474 self.wake.send_modify(|()| {});
475 }
476
477 fn terminate_and_replace(&self, message: &str, latest_committed_epoch: Option<u64>) -> Self {
478 let mut inner = self.inner.lock();
479 debug_assert!(
480 inner.reserved_marker.is_none(),
481 "subscription generation invalidated before its checkpoint marker was released"
482 );
483 let retention_cap = inner.retention_cap;
484 let next_sequence = inner.next_sequence;
485 Self::terminate_locked(&mut inner, message.to_owned());
486 drop(inner);
487 self.wake.send_modify(|()| {});
488 Self::new_at(
489 retention_cap,
490 Arc::clone(&self.budget),
491 latest_committed_epoch,
492 next_sequence,
493 )
494 }
495
496 fn subscriber_count(&self) -> usize {
497 self.inner.lock().readers.len()
498 }
499
500 fn observe_committed_epoch(&self, epoch: u64) {
501 let mut inner = self.inner.lock();
502 inner.latest_committed_epoch = Some(
503 inner
504 .latest_committed_epoch
505 .map_or(epoch, |latest| latest.max(epoch)),
506 );
507 }
508
509 fn terminate_locked(inner: &mut StreamLogInner, message: String) {
510 if inner.terminal_error.is_none() {
511 inner.terminal_error = Some(message);
512 }
513 clear_entries(inner);
514 }
515}
516
517impl Drop for StreamLog {
518 fn drop(&mut self) {
519 let inner = self.inner.get_mut();
520 debug_assert!(
521 inner.reserved_marker.is_none(),
522 "subscription log dropped with a live checkpoint marker reservation"
523 );
524 }
525}
526
527fn storage_cap(retention_cap: usize) -> usize {
528 retention_cap.max(INTERNAL_LIVE_LOG_BYTES)
529}
530
531fn head_sequence(inner: &StreamLogInner) -> u64 {
532 inner
533 .entries
534 .front()
535 .map_or(inner.next_sequence, |entry| entry.sequence)
536}
537
538fn clear_entries(inner: &mut StreamLogInner) {
539 inner.entries.clear();
540 inner.bytes = 0;
541 inner.retention_floor = inner.next_sequence;
542 inner.retention_bytes = 0;
543}
544
545fn pop_front(inner: &mut StreamLogInner) -> Option<SequencedUpdate> {
546 let entry = inner.entries.pop_front()?;
547 inner.bytes = inner.bytes.saturating_sub(entry.bytes);
548 if entry.sequence >= inner.retention_floor {
549 debug_assert_eq!(
550 entry.sequence, inner.retention_floor,
551 "retention floor skipped a retained entry"
552 );
553 inner.retention_bytes = inner.retention_bytes.saturating_sub(entry.bytes);
554 inner.retention_floor = entry.sequence.saturating_add(1);
555 if inner.retention_bytes == 0 {
556 inner.retention_floor = inner.next_sequence;
557 }
558 }
559 Some(entry)
560}
561
562fn evict_to_byte_target(inner: &mut StreamLogInner, target: usize) {
563 while inner.bytes > target {
564 let Some(_evicted) = pop_front(inner) else {
565 inner.bytes = 0;
566 inner.retention_floor = inner.next_sequence;
567 inner.retention_bytes = 0;
568 break;
569 };
570 }
571}
572
573fn reclaim_consumed_prefix(inner: &mut StreamLogInner) {
574 let reader_floor = inner
575 .readers
576 .values()
577 .copied()
578 .min()
579 .unwrap_or(inner.next_sequence);
580 let protected_floor = inner.retention_floor.min(reader_floor);
581
582 while inner
583 .entries
584 .front()
585 .is_some_and(|entry| entry.sequence < protected_floor)
586 {
587 let Some(_evicted) = pop_front(inner) else {
588 break;
589 };
590 }
591}
592
593fn calculate_retention_suffix(inner: &StreamLogInner) -> (u64, usize) {
594 if inner.retention_cap == 0 {
595 return (inner.next_sequence, 0);
596 }
597
598 let mut bytes = 0usize;
599 let mut floor = inner.next_sequence;
600 for entry in inner.entries.iter().rev() {
601 let with_entry = bytes.saturating_add(entry.bytes);
602 if with_entry > inner.retention_cap {
603 break;
604 }
605 bytes = with_entry;
606 floor = entry.sequence;
607 }
608 (floor, bytes)
609}
610
611fn recompute_retention_suffix(inner: &mut StreamLogInner) {
612 let (floor, bytes) = calculate_retention_suffix(inner);
613 inner.retention_floor = floor;
614 inner.retention_bytes = bytes;
615}
616
617fn retained_entry_bytes(inner: &StreamLogInner, sequence: u64) -> Option<usize> {
618 let head = head_sequence(inner);
619 if sequence < head || sequence >= inner.next_sequence {
620 return None;
621 }
622 let index = usize::try_from(sequence.saturating_sub(head)).ok()?;
623 inner.entries.get(index).map(|entry| entry.bytes)
624}
625
626fn retain_appended_entry(inner: &mut StreamLogInner, sequence: u64, bytes: usize) {
627 if inner.retention_cap == 0 {
628 inner.retention_floor = inner.next_sequence;
629 inner.retention_bytes = 0;
630 return;
631 }
632
633 if inner.retention_floor == sequence {
634 debug_assert_eq!(inner.retention_bytes, 0);
635 inner.retention_floor = sequence;
636 }
637 inner.retention_bytes = inner.retention_bytes.saturating_add(bytes);
638 while inner.retention_bytes > inner.retention_cap {
639 let Some(expired_bytes) = retained_entry_bytes(inner, inner.retention_floor) else {
640 debug_assert!(false, "retention floor is outside the shared log");
641 inner.retention_floor = inner.next_sequence;
642 inner.retention_bytes = 0;
643 return;
644 };
645 inner.retention_bytes = inner.retention_bytes.saturating_sub(expired_bytes);
646 inner.retention_floor = inner.retention_floor.saturating_add(1);
647 }
648 if inner.retention_bytes == 0 {
649 inner.retention_floor = inner.next_sequence;
650 }
651}
652
653fn cursor_after_sequence(
654 inner: &StreamLogInner,
655 requested: u64,
656) -> Result<u64, SubscriptionOpenError> {
657 if requested >= inner.next_sequence {
658 return Err(SubscriptionOpenError::SequenceNotPublished {
659 requested,
660 next_sequence: inner.next_sequence,
661 });
662 }
663 let cursor = requested + 1;
664 if cursor < inner.retention_floor {
666 return Err(SubscriptionOpenError::SequencePruned {
667 requested,
668 earliest_retained: (inner.retention_floor < inner.next_sequence)
669 .then_some(inner.retention_floor),
670 });
671 }
672 Ok(cursor)
673}
674
675fn cursor_after_retained_epoch(
676 inner: &StreamLogInner,
677 requested_epoch: u64,
678) -> Result<(u64, u64), SubscriptionOpenError> {
679 let Some(latest_committed) = inner.latest_committed_epoch else {
680 return Err(SubscriptionOpenError::EpochNotCommitted {
681 requested: requested_epoch,
682 latest_committed: None,
683 });
684 };
685 if requested_epoch > latest_committed {
686 return Err(SubscriptionOpenError::EpochNotCommitted {
687 requested: requested_epoch,
688 latest_committed: Some(latest_committed),
689 });
690 }
691 if inner.retention_cap == 0 {
692 return Err(SubscriptionOpenError::ReplayPruned {
693 earliest_retained: 0,
694 });
695 }
696
697 let mut cursor = None;
698 let mut earliest_retained = u64::MAX;
699 for entry in inner
700 .entries
701 .iter()
702 .filter(|entry| entry.sequence >= inner.retention_floor)
703 {
704 if let MvUpdate::Barrier {
705 epoch,
706 through_sequence,
707 ..
708 } = entry.update.as_ref()
709 {
710 if *through_sequence < inner.retention_floor {
711 continue;
712 }
713 earliest_retained = earliest_retained.min(*epoch);
714 if *epoch == requested_epoch {
715 cursor = Some((*through_sequence, entry.sequence));
716 }
717 }
718 }
719
720 if let Some(cursor) = cursor {
721 return Ok(cursor);
722 }
723 if earliest_retained == u64::MAX || requested_epoch < earliest_retained {
724 return Err(SubscriptionOpenError::ReplayPruned {
725 earliest_retained: if earliest_retained == u64::MAX {
726 0
727 } else {
728 earliest_retained
729 },
730 });
731 }
732 Err(SubscriptionOpenError::EpochNotCommitted {
733 requested: requested_epoch,
734 latest_committed: Some(latest_committed),
735 })
736}
737
738pub(super) fn approx_size(update: &MvUpdate) -> usize {
739 let payload = match update {
740 MvUpdate::Batch(batch) => batch.get_array_memory_size(),
741 MvUpdate::Barrier { .. } => return BARRIER_ENTRY_BYTES,
742 MvUpdate::Error(message) => message.len(),
743 };
744 payload.saturating_add(ENTRY_ACCOUNTING_BYTES)
745}
746
747#[cfg(test)]
748mod tests;