Skip to main content

ursula_runtime/engine/
mod.rs

1pub mod in_memory;
2pub mod wal;
3
4use std::borrow::Cow;
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::Arc;
8
9use serde::Deserialize;
10use serde::Serialize;
11use ursula_shard::BucketStreamId;
12use ursula_shard::ShardPlacement;
13use ursula_stream::ColdFlushCandidate;
14use ursula_stream::ColdGcEntry;
15use ursula_stream::StreamErrorCode;
16use ursula_stream::StreamErrorContext;
17
18use crate::command::GroupSnapshot;
19use crate::command::GroupWriteCommand;
20use crate::metrics::RaftWriteManySample;
21use crate::metrics::RuntimeMetricsInner;
22use crate::request::AckColdGcResponse;
23use crate::request::AppendBatchRequest;
24use crate::request::AppendExternalRequest;
25use crate::request::AppendRequest;
26use crate::request::AppendResponse;
27use crate::request::BootstrapStreamRequest;
28use crate::request::BootstrapStreamResponse;
29use crate::request::CloseStreamRequest;
30use crate::request::CloseStreamResponse;
31use crate::request::ColdHotBacklog;
32use crate::request::ColdWriteAdmission;
33use crate::request::CreateStreamExternalRequest;
34use crate::request::CreateStreamRequest;
35use crate::request::CreateStreamResponse;
36use crate::request::DeleteSnapshotRequest;
37use crate::request::DeleteStreamRequest;
38use crate::request::DeleteStreamResponse;
39use crate::request::FlushColdRequest;
40use crate::request::FlushColdResponse;
41use crate::request::ForkRefResponse;
42use crate::request::GroupReadStreamParts;
43use crate::request::HeadStreamRequest;
44use crate::request::HeadStreamResponse;
45use crate::request::PlanColdFlushRequest;
46use crate::request::PlanGroupColdFlushRequest;
47use crate::request::PublishSnapshotRequest;
48use crate::request::PublishSnapshotResponse;
49use crate::request::ReadSnapshotRequest;
50use crate::request::ReadSnapshotResponse;
51use crate::request::ReadStreamRequest;
52use crate::request::ReadStreamResponse;
53use crate::request::TouchStreamAccessResponse;
54
55pub type GroupAppendFuture<'a> =
56    Pin<Box<dyn Future<Output = Result<AppendResponse, GroupEngineError>> + Send + 'a>>;
57pub type GroupAppendBatchFuture<'a> =
58    Pin<Box<dyn Future<Output = Result<GroupAppendBatchResponse, GroupEngineError>> + Send + 'a>>;
59pub type GroupFlushColdFuture<'a> =
60    Pin<Box<dyn Future<Output = Result<FlushColdResponse, GroupEngineError>> + Send + 'a>>;
61pub type GroupPlanColdFlushFuture<'a> =
62    Pin<Box<dyn Future<Output = Result<Option<ColdFlushCandidate>, GroupEngineError>> + Send + 'a>>;
63pub type GroupPlanNextColdFlushFuture<'a> =
64    Pin<Box<dyn Future<Output = Result<Option<ColdFlushCandidate>, GroupEngineError>> + Send + 'a>>;
65pub type GroupPlanNextColdFlushBatchFuture<'a> =
66    Pin<Box<dyn Future<Output = Result<Vec<ColdFlushCandidate>, GroupEngineError>> + Send + 'a>>;
67pub type GroupColdHotBacklogFuture<'a> =
68    Pin<Box<dyn Future<Output = Result<ColdHotBacklog, GroupEngineError>> + Send + 'a>>;
69pub type GroupCreateStreamFuture<'a> =
70    Pin<Box<dyn Future<Output = Result<CreateStreamResponse, GroupEngineError>> + Send + 'a>>;
71pub type GroupHeadStreamFuture<'a> =
72    Pin<Box<dyn Future<Output = Result<HeadStreamResponse, GroupEngineError>> + Send + 'a>>;
73pub type GroupReadStreamFuture<'a> =
74    Pin<Box<dyn Future<Output = Result<ReadStreamResponse, GroupEngineError>> + Send + 'a>>;
75pub type GroupReadStreamPartsFuture<'a> =
76    Pin<Box<dyn Future<Output = Result<GroupReadStreamParts, GroupEngineError>> + Send + 'a>>;
77pub type GroupRequireLiveReadOwnerFuture<'a> =
78    Pin<Box<dyn Future<Output = Result<(), GroupEngineError>> + Send + 'a>>;
79pub type GroupPublishSnapshotFuture<'a> =
80    Pin<Box<dyn Future<Output = Result<PublishSnapshotResponse, GroupEngineError>> + Send + 'a>>;
81pub type GroupReadSnapshotFuture<'a> =
82    Pin<Box<dyn Future<Output = Result<ReadSnapshotResponse, GroupEngineError>> + Send + 'a>>;
83pub type GroupDeleteSnapshotFuture<'a> =
84    Pin<Box<dyn Future<Output = Result<(), GroupEngineError>> + Send + 'a>>;
85pub type GroupBootstrapStreamFuture<'a> =
86    Pin<Box<dyn Future<Output = Result<BootstrapStreamResponse, GroupEngineError>> + Send + 'a>>;
87pub type GroupTouchStreamAccessFuture<'a> =
88    Pin<Box<dyn Future<Output = Result<TouchStreamAccessResponse, GroupEngineError>> + Send + 'a>>;
89pub type GroupCloseStreamFuture<'a> =
90    Pin<Box<dyn Future<Output = Result<CloseStreamResponse, GroupEngineError>> + Send + 'a>>;
91pub type GroupDeleteStreamFuture<'a> =
92    Pin<Box<dyn Future<Output = Result<DeleteStreamResponse, GroupEngineError>> + Send + 'a>>;
93pub type GroupAckColdGcFuture<'a> =
94    Pin<Box<dyn Future<Output = Result<AckColdGcResponse, GroupEngineError>> + Send + 'a>>;
95pub type GroupPlanColdGcFuture<'a> =
96    Pin<Box<dyn Future<Output = Result<Vec<ColdGcEntry>, GroupEngineError>> + Send + 'a>>;
97pub type GroupForkRefFuture<'a> =
98    Pin<Box<dyn Future<Output = Result<ForkRefResponse, GroupEngineError>> + Send + 'a>>;
99pub type GroupSnapshotFuture<'a> =
100    Pin<Box<dyn Future<Output = Result<GroupSnapshot, GroupEngineError>> + Send + 'a>>;
101pub type GroupInstallSnapshotFuture<'a> =
102    Pin<Box<dyn Future<Output = Result<(), GroupEngineError>> + Send + 'a>>;
103pub type GroupShutdownFuture<'a> =
104    Pin<Box<dyn Future<Output = Result<(), GroupEngineError>> + Send + 'a>>;
105pub type GroupWriteBatchFuture<'a> = Pin<
106    Box<
107        dyn Future<
108                Output = Result<
109                    Vec<Result<GroupWriteResponse, GroupEngineError>>,
110                    GroupEngineError,
111                >,
112            > + Send
113            + 'a,
114    >,
115>;
116pub type GroupEngineCreateFuture<'a> =
117    Pin<Box<dyn Future<Output = Result<Box<dyn GroupEngine>, GroupEngineError>> + Send + 'a>>;
118
119#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
120pub struct GroupAppendBatchResponse {
121    pub placement: ShardPlacement,
122    pub items: Vec<Result<AppendResponse, GroupEngineError>>,
123}
124
125#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
126pub enum GroupWriteResponse {
127    CreateStream(CreateStreamResponse),
128    Append(AppendResponse),
129    AppendBatch(GroupAppendBatchResponse),
130    PublishSnapshot(PublishSnapshotResponse),
131    TouchStreamAccess(TouchStreamAccessResponse),
132    AddForkRef(ForkRefResponse),
133    ReleaseForkRef(ForkRefResponse),
134    FlushCold(FlushColdResponse),
135    CloseStream(CloseStreamResponse),
136    DeleteStream(DeleteStreamResponse),
137    AckColdGc(AckColdGcResponse),
138    Batch(Vec<Result<GroupWriteResponse, GroupEngineError>>),
139}
140
141pub trait GroupEngine: Send + 'static {
142    fn accepts_local_writes(&self) -> bool {
143        true
144    }
145
146    fn create_stream<'a>(
147        &'a mut self,
148        request: CreateStreamRequest,
149        placement: ShardPlacement,
150    ) -> GroupCreateStreamFuture<'a>;
151
152    fn create_stream_external<'a>(
153        &'a mut self,
154        request: CreateStreamExternalRequest,
155        _placement: ShardPlacement,
156    ) -> GroupCreateStreamFuture<'a> {
157        Box::pin(async move {
158            Err(GroupEngineError::new(format!(
159                "external stream create is not supported for stream '{}'",
160                request.stream_id
161            )))
162        })
163    }
164
165    fn head_stream<'a>(
166        &'a mut self,
167        request: HeadStreamRequest,
168        placement: ShardPlacement,
169    ) -> GroupHeadStreamFuture<'a>;
170
171    fn read_stream<'a>(
172        &'a mut self,
173        request: ReadStreamRequest,
174        placement: ShardPlacement,
175    ) -> GroupReadStreamFuture<'a>;
176
177    fn read_stream_parts<'a>(
178        &'a mut self,
179        request: ReadStreamRequest,
180        placement: ShardPlacement,
181    ) -> GroupReadStreamPartsFuture<'a> {
182        Box::pin(async move {
183            let response = self.read_stream(request, placement).await?;
184            Ok(GroupReadStreamParts::from_response(response))
185        })
186    }
187
188    fn require_local_live_read_owner<'a>(
189        &'a mut self,
190        _placement: ShardPlacement,
191    ) -> GroupRequireLiveReadOwnerFuture<'a> {
192        Box::pin(async { Ok(()) })
193    }
194
195    fn publish_snapshot<'a>(
196        &'a mut self,
197        request: PublishSnapshotRequest,
198        _placement: ShardPlacement,
199    ) -> GroupPublishSnapshotFuture<'a> {
200        Box::pin(async move {
201            Err(GroupEngineError::new(format!(
202                "snapshot publish is not supported for stream '{}'",
203                request.stream_id
204            )))
205        })
206    }
207
208    fn read_snapshot<'a>(
209        &'a mut self,
210        request: ReadSnapshotRequest,
211        _placement: ShardPlacement,
212    ) -> GroupReadSnapshotFuture<'a> {
213        Box::pin(async move {
214            Err(GroupEngineError::new(format!(
215                "snapshot read is not supported for stream '{}'",
216                request.stream_id
217            )))
218        })
219    }
220
221    fn delete_snapshot<'a>(
222        &'a mut self,
223        request: DeleteSnapshotRequest,
224        _placement: ShardPlacement,
225    ) -> GroupDeleteSnapshotFuture<'a> {
226        Box::pin(async move {
227            Err(GroupEngineError::new(format!(
228                "snapshot delete is not supported for stream '{}'",
229                request.stream_id
230            )))
231        })
232    }
233
234    fn bootstrap_stream<'a>(
235        &'a mut self,
236        request: BootstrapStreamRequest,
237        _placement: ShardPlacement,
238    ) -> GroupBootstrapStreamFuture<'a> {
239        Box::pin(async move {
240            Err(GroupEngineError::new(format!(
241                "bootstrap is not supported for stream '{}'",
242                request.stream_id
243            )))
244        })
245    }
246
247    fn touch_stream_access<'a>(
248        &'a mut self,
249        stream_id: BucketStreamId,
250        now_ms: u64,
251        renew_ttl: bool,
252        placement: ShardPlacement,
253    ) -> GroupTouchStreamAccessFuture<'a>;
254
255    fn add_fork_ref<'a>(
256        &'a mut self,
257        stream_id: BucketStreamId,
258        now_ms: u64,
259        placement: ShardPlacement,
260    ) -> GroupForkRefFuture<'a>;
261
262    fn release_fork_ref<'a>(
263        &'a mut self,
264        stream_id: BucketStreamId,
265        placement: ShardPlacement,
266    ) -> GroupForkRefFuture<'a>;
267
268    fn close_stream<'a>(
269        &'a mut self,
270        request: CloseStreamRequest,
271        placement: ShardPlacement,
272    ) -> GroupCloseStreamFuture<'a>;
273
274    fn delete_stream<'a>(
275        &'a mut self,
276        request: DeleteStreamRequest,
277        placement: ShardPlacement,
278    ) -> GroupDeleteStreamFuture<'a>;
279
280    /// Replicated confirmation that cold-GC entries up to `up_to_seq` have been
281    /// physically reclaimed; pops them from the queue. Default unsupported.
282    fn ack_cold_gc<'a>(
283        &'a mut self,
284        _up_to_seq: u64,
285        _placement: ShardPlacement,
286    ) -> GroupAckColdGcFuture<'a> {
287        Box::pin(async { Err(GroupEngineError::new("cold GC ack is not supported")) })
288    }
289
290    /// Leader-local read of the front of the cold-GC queue for the background
291    /// worker to reclaim. Default returns an empty batch.
292    fn plan_cold_gc<'a>(
293        &'a mut self,
294        _max: usize,
295        _placement: ShardPlacement,
296    ) -> GroupPlanColdGcFuture<'a> {
297        Box::pin(async { Ok(Vec::new()) })
298    }
299
300    fn append<'a>(
301        &'a mut self,
302        request: AppendRequest,
303        placement: ShardPlacement,
304    ) -> GroupAppendFuture<'a>;
305
306    fn append_external<'a>(
307        &'a mut self,
308        request: AppendExternalRequest,
309        _placement: ShardPlacement,
310    ) -> GroupAppendFuture<'a> {
311        Box::pin(async move {
312            Err(GroupEngineError::new(format!(
313                "external append is not supported for stream '{}'",
314                request.stream_id
315            )))
316        })
317    }
318
319    fn append_batch<'a>(
320        &'a mut self,
321        request: AppendBatchRequest,
322        placement: ShardPlacement,
323    ) -> GroupAppendBatchFuture<'a>;
324
325    fn create_stream_with_cold_admission<'a>(
326        &'a mut self,
327        request: CreateStreamRequest,
328        placement: ShardPlacement,
329        _admission: ColdWriteAdmission,
330    ) -> GroupCreateStreamFuture<'a> {
331        self.create_stream(request, placement)
332    }
333
334    fn append_with_cold_admission<'a>(
335        &'a mut self,
336        request: AppendRequest,
337        placement: ShardPlacement,
338        _admission: ColdWriteAdmission,
339    ) -> GroupAppendFuture<'a> {
340        self.append(request, placement)
341    }
342
343    fn append_batch_with_cold_admission<'a>(
344        &'a mut self,
345        request: AppendBatchRequest,
346        placement: ShardPlacement,
347        _admission: ColdWriteAdmission,
348    ) -> GroupAppendBatchFuture<'a> {
349        self.append_batch(request, placement)
350    }
351
352    fn append_batch_many_with_cold_admission<'a>(
353        &'a mut self,
354        requests: Vec<AppendBatchRequest>,
355        placement: ShardPlacement,
356        admission: ColdWriteAdmission,
357    ) -> GroupWriteBatchFuture<'a> {
358        Box::pin(async move {
359            let mut responses = Vec::with_capacity(requests.len());
360            for request in requests {
361                let response = self
362                    .append_batch_with_cold_admission(request, placement, admission)
363                    .await
364                    .map(GroupWriteResponse::AppendBatch);
365                responses.push(response);
366            }
367            Ok(responses)
368        })
369    }
370
371    fn flush_cold<'a>(
372        &'a mut self,
373        request: FlushColdRequest,
374        _placement: ShardPlacement,
375    ) -> GroupFlushColdFuture<'a> {
376        Box::pin(async move {
377            Err(GroupEngineError::new(format!(
378                "cold flush is not supported for stream '{}'",
379                request.stream_id
380            )))
381        })
382    }
383
384    fn plan_cold_flush<'a>(
385        &'a mut self,
386        request: PlanColdFlushRequest,
387        _placement: ShardPlacement,
388    ) -> GroupPlanColdFlushFuture<'a> {
389        Box::pin(async move {
390            Err(GroupEngineError::new(format!(
391                "cold flush planning is not supported for stream '{}'",
392                request.stream_id
393            )))
394        })
395    }
396
397    fn plan_next_cold_flush<'a>(
398        &'a mut self,
399        _request: PlanGroupColdFlushRequest,
400        _placement: ShardPlacement,
401    ) -> GroupPlanNextColdFlushFuture<'a> {
402        Box::pin(async move {
403            Err(GroupEngineError::new(
404                "group cold flush planning is not supported",
405            ))
406        })
407    }
408
409    fn plan_next_cold_flush_batch<'a>(
410        &'a mut self,
411        request: PlanGroupColdFlushRequest,
412        placement: ShardPlacement,
413        max_candidates: usize,
414    ) -> GroupPlanNextColdFlushBatchFuture<'a> {
415        Box::pin(async move {
416            match self.plan_next_cold_flush(request, placement).await? {
417                Some(candidate) if max_candidates > 0 => Ok(vec![candidate]),
418                _ => Ok(Vec::new()),
419            }
420        })
421    }
422
423    fn cold_hot_backlog<'a>(
424        &'a mut self,
425        stream_id: BucketStreamId,
426        _placement: ShardPlacement,
427    ) -> GroupColdHotBacklogFuture<'a> {
428        Box::pin(async move {
429            Err(GroupEngineError::new(format!(
430                "cold hot backlog is not supported for stream '{stream_id}'"
431            )))
432        })
433    }
434
435    fn snapshot<'a>(&'a mut self, placement: ShardPlacement) -> GroupSnapshotFuture<'a>;
436
437    fn install_snapshot<'a>(
438        &'a mut self,
439        snapshot: GroupSnapshot,
440    ) -> GroupInstallSnapshotFuture<'a>;
441
442    fn shutdown<'a>(&'a mut self) -> GroupShutdownFuture<'a> {
443        Box::pin(async { Ok(()) })
444    }
445
446    fn write_batch<'a>(
447        &'a mut self,
448        commands: Vec<GroupWriteCommand>,
449        placement: ShardPlacement,
450    ) -> GroupWriteBatchFuture<'a> {
451        Box::pin(async move {
452            let mut responses = Vec::with_capacity(commands.len());
453            for command in commands {
454                let response = match command {
455                    GroupWriteCommand::CreateStream {
456                        stream_id,
457                        content_type,
458                        initial_payload,
459                        close_after,
460                        stream_seq,
461                        producer,
462                        stream_ttl_seconds,
463                        stream_expires_at_ms,
464                        forked_from,
465                        fork_offset,
466                        now_ms,
467                    } => self
468                        .create_stream(
469                            CreateStreamRequest {
470                                stream_id,
471                                content_type,
472                                content_type_explicit: true,
473                                initial_payload,
474                                close_after,
475                                stream_seq,
476                                producer,
477                                stream_ttl_seconds,
478                                stream_expires_at_ms,
479                                forked_from,
480                                fork_offset,
481                                now_ms,
482                            },
483                            placement,
484                        )
485                        .await
486                        .map(GroupWriteResponse::CreateStream),
487                    GroupWriteCommand::CreateExternal {
488                        stream_id,
489                        content_type,
490                        initial_payload,
491                        close_after,
492                        stream_seq,
493                        producer,
494                        stream_ttl_seconds,
495                        stream_expires_at_ms,
496                        forked_from,
497                        fork_offset,
498                        now_ms,
499                    } => self
500                        .create_stream_external(
501                            CreateStreamExternalRequest {
502                                stream_id,
503                                content_type,
504                                initial_payload,
505                                close_after,
506                                stream_seq,
507                                producer,
508                                stream_ttl_seconds,
509                                stream_expires_at_ms,
510                                forked_from,
511                                fork_offset,
512                                now_ms,
513                            },
514                            placement,
515                        )
516                        .await
517                        .map(GroupWriteResponse::CreateStream),
518                    GroupWriteCommand::Append {
519                        stream_id,
520                        content_type,
521                        payload,
522                        close_after,
523                        stream_seq,
524                        producer,
525                        now_ms,
526                    } => self
527                        .append(
528                            AppendRequest {
529                                stream_id,
530                                content_type,
531                                payload,
532                                close_after,
533                                stream_seq,
534                                producer,
535                                now_ms,
536                            },
537                            placement,
538                        )
539                        .await
540                        .map(GroupWriteResponse::Append),
541                    GroupWriteCommand::AppendExternal {
542                        stream_id,
543                        content_type,
544                        payload,
545                        close_after,
546                        stream_seq,
547                        producer,
548                        now_ms,
549                    } => self
550                        .append_external(
551                            AppendExternalRequest {
552                                stream_id,
553                                content_type,
554                                payload,
555                                close_after,
556                                stream_seq,
557                                producer,
558                                now_ms,
559                            },
560                            placement,
561                        )
562                        .await
563                        .map(GroupWriteResponse::Append),
564                    GroupWriteCommand::AppendBatch {
565                        stream_id,
566                        content_type,
567                        payloads,
568                        producer,
569                        now_ms,
570                    } => self
571                        .append_batch(
572                            AppendBatchRequest {
573                                stream_id,
574                                content_type,
575                                payloads,
576                                producer,
577                                now_ms,
578                            },
579                            placement,
580                        )
581                        .await
582                        .map(GroupWriteResponse::AppendBatch),
583                    GroupWriteCommand::PublishSnapshot {
584                        stream_id,
585                        snapshot_offset,
586                        content_type,
587                        payload,
588                        now_ms,
589                    } => self
590                        .publish_snapshot(
591                            PublishSnapshotRequest {
592                                stream_id,
593                                snapshot_offset,
594                                content_type,
595                                payload,
596                                now_ms,
597                            },
598                            placement,
599                        )
600                        .await
601                        .map(GroupWriteResponse::PublishSnapshot),
602                    GroupWriteCommand::TouchStreamAccess {
603                        stream_id,
604                        now_ms,
605                        renew_ttl,
606                    } => self
607                        .touch_stream_access(stream_id, now_ms, renew_ttl, placement)
608                        .await
609                        .map(GroupWriteResponse::TouchStreamAccess),
610                    GroupWriteCommand::AddForkRef { stream_id, now_ms } => self
611                        .add_fork_ref(stream_id, now_ms, placement)
612                        .await
613                        .map(GroupWriteResponse::AddForkRef),
614                    GroupWriteCommand::ReleaseForkRef { stream_id } => self
615                        .release_fork_ref(stream_id, placement)
616                        .await
617                        .map(GroupWriteResponse::ReleaseForkRef),
618                    GroupWriteCommand::FlushCold { stream_id, chunk } => self
619                        .flush_cold(FlushColdRequest { stream_id, chunk }, placement)
620                        .await
621                        .map(GroupWriteResponse::FlushCold),
622                    GroupWriteCommand::CloseStream {
623                        stream_id,
624                        stream_seq,
625                        producer,
626                        now_ms,
627                    } => self
628                        .close_stream(
629                            CloseStreamRequest {
630                                stream_id,
631                                stream_seq,
632                                producer,
633                                now_ms,
634                            },
635                            placement,
636                        )
637                        .await
638                        .map(GroupWriteResponse::CloseStream),
639                    GroupWriteCommand::DeleteStream { stream_id } => self
640                        .delete_stream(DeleteStreamRequest { stream_id }, placement)
641                        .await
642                        .map(GroupWriteResponse::DeleteStream),
643                    GroupWriteCommand::AckColdGc { up_to_seq } => self
644                        .ack_cold_gc(up_to_seq, placement)
645                        .await
646                        .map(GroupWriteResponse::AckColdGc),
647                    GroupWriteCommand::Batch { commands } => self
648                        .write_batch(commands, placement)
649                        .await
650                        .map(GroupWriteResponse::Batch),
651                };
652                responses.push(response);
653            }
654            Ok(responses)
655        })
656    }
657}
658
659pub trait GroupEngineFactory: Send + Sync + 'static {
660    fn hosts_group(&self, _placement: ShardPlacement) -> bool {
661        true
662    }
663
664    fn create<'a>(
665        &'a self,
666        placement: ShardPlacement,
667        metrics: GroupEngineMetrics,
668    ) -> GroupEngineCreateFuture<'a>;
669}
670
671#[derive(Debug, Clone)]
672pub struct GroupEngineMetrics {
673    pub(crate) inner: Arc<RuntimeMetricsInner>,
674}
675
676impl GroupEngineMetrics {
677    pub fn record_wal_batch(
678        &self,
679        placement: ShardPlacement,
680        record_count: usize,
681        write_ns: u64,
682        sync_ns: u64,
683    ) {
684        self.inner.record_wal_batch(
685            placement.core_id,
686            placement.raft_group_id,
687            u64::try_from(record_count).expect("record count fits u64"),
688            write_ns,
689            sync_ns,
690        );
691    }
692
693    pub fn record_raft_write_many(
694        &self,
695        placement: ShardPlacement,
696        command_count: usize,
697        logical_command_count: usize,
698        response_count: usize,
699        submit_ns: u64,
700        response_ns: u64,
701    ) {
702        self.inner.record_raft_write_many(
703            placement.core_id,
704            placement.raft_group_id,
705            RaftWriteManySample {
706                command_count: u64::try_from(command_count).expect("command count fits u64"),
707                logical_command_count: u64::try_from(logical_command_count)
708                    .expect("logical command count fits u64"),
709                response_count: u64::try_from(response_count).expect("response count fits u64"),
710                submit_ns,
711                response_ns,
712            },
713        );
714    }
715
716    pub fn record_raft_apply_batch(
717        &self,
718        placement: ShardPlacement,
719        entry_count: usize,
720        apply_ns: u64,
721    ) {
722        self.inner.record_raft_apply_batch(
723            placement.core_id,
724            placement.raft_group_id,
725            u64::try_from(entry_count).expect("entry count fits u64"),
726            apply_ns,
727        );
728    }
729}
730
731#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
732pub struct GroupLeaderHint {
733    pub node_id: Option<u64>,
734    pub address: Option<String>,
735}
736
737#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
738pub struct StreamEngineError {
739    message: String,
740    code: StreamErrorCode,
741    next_offset: Option<u64>,
742    #[serde(default, skip_serializing_if = "Vec::is_empty")]
743    context: Vec<StreamErrorContext>,
744}
745
746/// Infra error variants with structured fields render their human message on
747/// demand (`message`) instead of storing a denormalized copy alongside the
748/// fields. `Internal` is the exception: it carries free-form text with no
749/// structured source, so it keeps an owned `message`.
750#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
751pub enum GroupInfraError {
752    Internal {
753        message: String,
754    },
755    ProtoDecode {
756        field: String,
757    },
758    ColdBackpressure {
759        stream_id: BucketStreamId,
760        before_group_hot_bytes: u64,
761        after_group_hot_bytes: u64,
762        limit: u64,
763    },
764    RaftUncommittedBackpressure {
765        current: u64,
766        incoming: u64,
767        limit: u64,
768    },
769}
770
771impl GroupInfraError {
772    pub fn internal(message: impl Into<String>) -> Self {
773        Self::Internal {
774            message: message.into(),
775        }
776    }
777
778    pub fn proto_decode(field: impl Into<String>) -> Self {
779        Self::ProtoDecode {
780            field: field.into(),
781        }
782    }
783
784    pub fn cold_backpressure(
785        stream_id: BucketStreamId,
786        before_group_hot_bytes: u64,
787        after_group_hot_bytes: u64,
788        limit: u64,
789    ) -> Self {
790        Self::ColdBackpressure {
791            stream_id,
792            before_group_hot_bytes,
793            after_group_hot_bytes,
794            limit,
795        }
796    }
797
798    pub fn raft_uncommitted_backpressure(current: u64, incoming: u64, limit: u64) -> Self {
799        Self::RaftUncommittedBackpressure {
800            current,
801            incoming,
802            limit,
803        }
804    }
805
806    pub fn message(&self) -> Cow<'_, str> {
807        match self {
808            Self::Internal { message } => Cow::Borrowed(message),
809            Self::ProtoDecode { field } => Cow::Owned(format!(
810                "ProtoDecode: protobuf raft payload missing {field}"
811            )),
812            Self::ColdBackpressure {
813                stream_id,
814                before_group_hot_bytes,
815                after_group_hot_bytes,
816                limit,
817            } => Cow::Owned(format!(
818                "ColdBackpressure: stream '{stream_id}' would raise group hot bytes from {before_group_hot_bytes} to {after_group_hot_bytes}, above limit {limit}"
819            )),
820            Self::RaftUncommittedBackpressure {
821                current,
822                incoming,
823                limit,
824            } => Cow::Owned(format!(
825                "RaftUncommittedBackpressure: group uncommitted bytes {current} plus incoming {incoming} would exceed limit {limit}"
826            )),
827        }
828    }
829
830    pub fn is_cold_backpressure(&self) -> bool {
831        matches!(self, Self::ColdBackpressure { .. })
832    }
833
834    pub fn is_backpressure(&self) -> bool {
835        matches!(
836            self,
837            Self::ColdBackpressure { .. } | Self::RaftUncommittedBackpressure { .. }
838        )
839    }
840}
841
842#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
843pub enum GroupEngineError {
844    Stream(StreamEngineError),
845    Infra(GroupInfraError),
846    ForwardToLeader {
847        message: String,
848        leader_hint: GroupLeaderHint,
849    },
850}
851
852impl GroupEngineError {
853    pub fn new(message: impl Into<String>) -> Self {
854        Self::Infra(GroupInfraError::internal(message))
855    }
856
857    pub fn cold_backpressure(
858        stream_id: BucketStreamId,
859        before_group_hot_bytes: u64,
860        after_group_hot_bytes: u64,
861        limit: u64,
862    ) -> Self {
863        Self::Infra(GroupInfraError::cold_backpressure(
864            stream_id,
865            before_group_hot_bytes,
866            after_group_hot_bytes,
867            limit,
868        ))
869    }
870
871    pub fn raft_uncommitted_backpressure(current: u64, incoming: u64, limit: u64) -> Self {
872        Self::Infra(GroupInfraError::raft_uncommitted_backpressure(
873            current, incoming, limit,
874        ))
875    }
876
877    pub fn stream(code: StreamErrorCode, message: impl Into<String>) -> Self {
878        Self::stream_with_next_offset(code, message, None)
879    }
880
881    pub fn stream_with_next_offset(
882        code: StreamErrorCode,
883        message: impl Into<String>,
884        next_offset: Option<u64>,
885    ) -> Self {
886        Self::stream_with_context(code, message, next_offset, vec![])
887    }
888
889    pub fn stream_with_context(
890        code: StreamErrorCode,
891        message: impl Into<String>,
892        next_offset: Option<u64>,
893        context: Vec<StreamErrorContext>,
894    ) -> Self {
895        Self::Stream(StreamEngineError {
896            message: format!("{code:?}: {}", message.into()),
897            code,
898            next_offset,
899            context,
900        })
901    }
902
903    pub fn stream_from_replicated(
904        message: impl Into<String>,
905        code: StreamErrorCode,
906        next_offset: Option<u64>,
907        context: Vec<StreamErrorContext>,
908    ) -> Self {
909        Self::Stream(StreamEngineError {
910            message: message.into(),
911            code,
912            next_offset,
913            context,
914        })
915    }
916
917    pub fn forward_to_leader(
918        message: impl Into<String>,
919        node_id: Option<u64>,
920        address: Option<String>,
921    ) -> Self {
922        Self::ForwardToLeader {
923            message: message.into(),
924            leader_hint: GroupLeaderHint { node_id, address },
925        }
926    }
927
928    pub fn message(&self) -> Cow<'_, str> {
929        match self {
930            Self::Stream(err) => Cow::Borrowed(&err.message),
931            Self::Infra(err) => err.message(),
932            Self::ForwardToLeader { message, .. } => Cow::Borrowed(message),
933        }
934    }
935
936    pub fn code(&self) -> Option<StreamErrorCode> {
937        match self {
938            Self::Stream(err) => Some(err.code),
939            Self::Infra(_) | Self::ForwardToLeader { .. } => None,
940        }
941    }
942
943    pub fn stream_parts(
944        &self,
945    ) -> Option<(&str, StreamErrorCode, Option<u64>, &[StreamErrorContext])> {
946        match self {
947            Self::Stream(err) => Some((&err.message, err.code, err.next_offset, &err.context)),
948            Self::Infra(_) | Self::ForwardToLeader { .. } => None,
949        }
950    }
951
952    pub fn next_offset(&self) -> Option<u64> {
953        match self {
954            Self::Stream(err) => err.next_offset,
955            Self::Infra(_) | Self::ForwardToLeader { .. } => None,
956        }
957    }
958
959    pub fn context(&self) -> &[StreamErrorContext] {
960        match self {
961            Self::Stream(err) => &err.context,
962            Self::Infra(_) | Self::ForwardToLeader { .. } => &[],
963        }
964    }
965
966    pub fn leader_hint(&self) -> Option<&GroupLeaderHint> {
967        match self {
968            Self::ForwardToLeader { leader_hint, .. } => Some(leader_hint),
969            Self::Stream(_) | Self::Infra(_) => None,
970        }
971    }
972
973    pub fn infra(&self) -> Option<&GroupInfraError> {
974        match self {
975            Self::Infra(err) => Some(err),
976            Self::Stream(_) | Self::ForwardToLeader { .. } => None,
977        }
978    }
979
980    pub fn is_cold_backpressure(&self) -> bool {
981        self.infra()
982            .is_some_and(GroupInfraError::is_cold_backpressure)
983    }
984
985    pub fn is_backpressure(&self) -> bool {
986        self.infra().is_some_and(GroupInfraError::is_backpressure)
987    }
988}
989
990impl std::fmt::Display for GroupEngineError {
991    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
992        f.write_str(&self.message())
993    }
994}
995
996impl std::error::Error for GroupEngineError {}