1#![allow(refining_impl_trait)]
4
5use std::collections::HashMap;
6use std::pin::Pin;
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::Arc;
9
10use bytes::Bytes;
11use connectrpc::{
12 Chain, ConnectError, ConnectRpcService, Limits, PreEncoded, RequestContext as Context,
13};
14use exoware_proto::common::Entry;
15use exoware_proto::compact::{
16 PruneResponse, Service as CompactApi, ServiceServer as CompactServiceServer,
17};
18use exoware_proto::google::rpc::{ErrorInfo, RetryInfo};
19use exoware_proto::ingest::{
20 PutResponse as ProtoPutResponse, Service as IngestApi, ServiceServer as IngestServiceServer,
21};
22use exoware_proto::log::stream::v1::{
23 GetRequestView, GetResponse as StreamGetResponse, Service as StreamApi,
24 ServiceServer as StreamServiceServer, SubscribeRequestView, SubscribeResponse,
25};
26use exoware_proto::query::{
27 Detail, GetManyEntry, GetManyFrame, GetResponse, RangeFrame, ReduceResponse,
28 Service as QueryApi, ServiceServer as QueryServiceServer,
29};
30use exoware_proto::stream_filter::{Filter, StreamFilter};
31use exoware_proto::{
32 connect_compression_registry, parse_range_traversal_direction,
33 to_domain_reduce_request_from_view, to_proto_optional_reduced_value, to_proto_reduced_value,
34 with_error_info_detail, with_query_detail, with_retry_info_detail, RangeTraversalDirection,
35};
36use exoware_sdk as exoware_proto;
37use exoware_sdk::common::kv::v1::filter::KindView as ProtoFilterKindView;
38use exoware_sdk::keys::Key;
39use exoware_sdk::selector::Selector;
40use futures::{stream as stream_util, Stream};
41use tokio::sync::Notify;
42
43use crate::reduce::RangeReducer;
44use crate::stream::{StreamHub, StreamNotifier};
45use crate::validate::{self, IngestLimits};
46use crate::{Ingest, IngestError, Log, LogBatch, Prune, Query, QueryExtra, RangeScan, StoreEngine};
47
48const MAX_CONNECTRPC_BODY_BYTES: usize = 256 * 1024 * 1024;
50const RANGE_STREAM_MAX_FRAME_ROWS: usize = 4096;
51const REDUCE_SCAN_BATCH_SIZE: usize = 4096;
52
53fn query_detail(sequence_number: u64, extra: QueryExtra) -> Detail {
54 Detail {
55 sequence_number,
56 extra,
57 ..Default::default()
58 }
59}
60
61struct RangeStreamRequest {
62 start_key: Key,
63 end_key: Key,
64 limit: usize,
65 batch_size: usize,
66 forward: bool,
67 sequence_number: u64,
68}
69
70async fn range_stream<Q>(
71 query: Arc<Q>,
72 request: RangeStreamRequest,
73) -> Result<Pin<Box<dyn Stream<Item = Result<RangeFrame, ConnectError>> + Send>>, ConnectError>
74where
75 Q: Query,
76{
77 let RangeStreamRequest {
78 start_key,
79 end_key,
80 limit,
81 batch_size,
82 forward,
83 sequence_number,
84 } = request;
85 let entries = query
86 .range_scan(start_key, end_key, limit, forward)
87 .await
88 .map_err(ConnectError::internal)?;
89
90 Ok(Box::pin(stream_util::unfold(
91 Some((entries, false)),
92 move |state| async move {
93 let (mut entries, emitted_frame) = state?;
94 let batch = match entries.next_batch(batch_size).await {
95 Ok(batch) => batch,
96 Err(e) => return Some((Err(ConnectError::internal(e)), None)),
97 };
98 let detail = query_detail(sequence_number, batch.extra);
99 if batch.rows.is_empty() {
100 if emitted_frame && detail.extra.is_empty() {
101 return None;
102 }
103 return Some((
104 Ok(RangeFrame {
105 detail: Some(detail).into(),
106 ..Default::default()
107 }),
108 None,
109 ));
110 }
111
112 let mut chunk = Vec::with_capacity(batch.rows.len());
113 for (key, value) in batch.rows {
114 chunk.push(Entry {
115 key: key.into(),
116 value,
117 ..Default::default()
118 });
119 }
120 Some((
121 Ok(RangeFrame {
122 results: chunk,
123 detail: Some(detail).into(),
124 ..Default::default()
125 }),
126 Some((entries, true)),
127 ))
128 },
129 )))
130}
131
132pub struct AppState<E> {
135 pub engine: Arc<E>,
137 pub ingest_limits: IngestLimits,
139 pub ready: Arc<AtomicBool>,
142 pub stream: Arc<StreamHub>,
144}
145
146impl<E> Clone for AppState<E> {
147 fn clone(&self) -> Self {
148 Self {
149 engine: self.engine.clone(),
150 ingest_limits: self.ingest_limits,
151 ready: self.ready.clone(),
152 stream: self.stream.clone(),
153 }
154 }
155}
156
157impl<E> AppState<E>
158where
159 E: StoreEngine,
160{
161 pub fn new(engine: Arc<E>) -> Self {
162 let current_sequence = engine.current_sequence();
163 Self {
164 engine,
165 ingest_limits: IngestLimits::default(),
166 ready: Arc::new(AtomicBool::new(true)),
167 stream: Arc::new(StreamHub::new(current_sequence)),
168 }
169 }
170
171 pub fn with_ingest_limits(mut self, limits: IngestLimits) -> Self {
172 self.ingest_limits = limits;
173 self
174 }
175}
176
177pub struct IngestState<I> {
179 pub ingest: Arc<I>,
181 pub limits: IngestLimits,
183 pub ready: Arc<AtomicBool>,
185 pub notifier: Option<Arc<dyn StreamNotifier>>,
187}
188
189impl<I> Clone for IngestState<I> {
190 fn clone(&self) -> Self {
191 Self {
192 ingest: self.ingest.clone(),
193 limits: self.limits,
194 ready: self.ready.clone(),
195 notifier: self.notifier.clone(),
196 }
197 }
198}
199
200impl<I> IngestState<I>
201where
202 I: Ingest,
203{
204 pub fn new(ingest: Arc<I>) -> Self {
205 Self {
206 ingest,
207 limits: IngestLimits::default(),
208 ready: Arc::new(AtomicBool::new(true)),
209 notifier: None,
210 }
211 }
212
213 pub fn with_notifier(ingest: Arc<I>, notifier: Arc<dyn StreamNotifier>) -> Self {
214 Self {
215 ingest,
216 limits: IngestLimits::default(),
217 ready: Arc::new(AtomicBool::new(true)),
218 notifier: Some(notifier),
219 }
220 }
221
222 pub fn with_limits(mut self, limits: IngestLimits) -> Self {
223 self.limits = limits;
224 self
225 }
226}
227
228impl<E> From<AppState<E>> for IngestState<E> {
229 fn from(state: AppState<E>) -> Self {
230 Self {
231 ingest: state.engine,
232 limits: state.ingest_limits,
233 ready: state.ready,
234 notifier: Some(state.stream),
235 }
236 }
237}
238
239pub struct QueryState<Q> {
241 pub query: Arc<Q>,
243}
244
245impl<Q> Clone for QueryState<Q> {
246 fn clone(&self) -> Self {
247 Self {
248 query: self.query.clone(),
249 }
250 }
251}
252
253impl<Q> QueryState<Q>
254where
255 Q: Query,
256{
257 pub fn new(query: Arc<Q>) -> Self {
258 Self { query }
259 }
260}
261
262impl<E> From<AppState<E>> for QueryState<E> {
263 fn from(state: AppState<E>) -> Self {
264 Self {
265 query: state.engine,
266 }
267 }
268}
269
270pub struct CompactState<P> {
272 pub prune: Arc<P>,
274}
275
276impl<P> Clone for CompactState<P> {
277 fn clone(&self) -> Self {
278 Self {
279 prune: self.prune.clone(),
280 }
281 }
282}
283
284impl<P> CompactState<P>
285where
286 P: Prune,
287{
288 pub fn new(prune: Arc<P>) -> Self {
289 Self { prune }
290 }
291}
292
293impl<E> From<AppState<E>> for CompactState<E> {
294 fn from(state: AppState<E>) -> Self {
295 Self {
296 prune: state.engine,
297 }
298 }
299}
300
301pub struct StreamState<L> {
303 pub log: Arc<L>,
305 pub notifier: Arc<dyn StreamNotifier>,
307}
308
309impl<L> Clone for StreamState<L> {
310 fn clone(&self) -> Self {
311 Self {
312 log: self.log.clone(),
313 notifier: self.notifier.clone(),
314 }
315 }
316}
317
318impl<L> StreamState<L>
319where
320 L: Log,
321{
322 pub fn new(log: Arc<L>, notifier: Arc<dyn StreamNotifier>) -> Self {
323 Self { log, notifier }
324 }
325}
326
327impl<E> From<AppState<E>> for StreamState<E> {
328 fn from(state: AppState<E>) -> Self {
329 Self {
330 log: state.engine,
331 notifier: state.stream,
332 }
333 }
334}
335
336pub struct IngestConnect<I> {
337 state: IngestState<I>,
338}
339
340impl<I> Clone for IngestConnect<I> {
341 fn clone(&self) -> Self {
342 Self {
343 state: self.state.clone(),
344 }
345 }
346}
347
348impl<I> IngestConnect<I>
349where
350 I: Ingest,
351{
352 pub fn new(state: impl Into<IngestState<I>>) -> Self {
353 Self {
354 state: state.into(),
355 }
356 }
357}
358
359const RETRY_HINT_DELAY: std::time::Duration = std::time::Duration::from_secs(1);
361const REASON_WORKER_NOT_READY: &str = "WORKER_NOT_READY";
363const REASON_INGEST_UNAVAILABLE: &str = "INGEST_UNAVAILABLE";
365
366fn with_retry_hint(err: ConnectError, retry_delay: std::time::Duration) -> ConnectError {
368 with_retry_info_detail(
369 err,
370 RetryInfo {
371 retry_delay: Some(buffa_types::google::protobuf::Duration::from(retry_delay)).into(),
372 ..Default::default()
373 },
374 )
375}
376
377fn ingest_error_to_connect(err: IngestError) -> ConnectError {
378 match err {
379 IngestError::Unavailable { message } => with_retry_hint(
380 with_error_info_detail(
381 ConnectError::unavailable(message),
382 ErrorInfo {
383 reason: REASON_INGEST_UNAVAILABLE.to_string(),
384 domain: crate::validate::INGEST_ERROR_DOMAIN.to_string(),
385 ..Default::default()
386 },
387 ),
388 RETRY_HINT_DELAY,
389 ),
390 IngestError::Internal { message } => ConnectError::internal(message),
391 }
392}
393
394impl<I> IngestApi for IngestConnect<I>
395where
396 I: Ingest,
397{
398 async fn put(
399 &self,
400 _ctx: Context,
401 request: buffa::view::OwnedView<exoware_proto::log::ingest::v1::PutRequestView<'static>>,
402 ) -> connectrpc::ServiceResult<ProtoPutResponse> {
403 if !self.state.ready.load(Ordering::SeqCst) {
404 return Err(with_retry_hint(
405 with_error_info_detail(
406 ConnectError::unavailable("ingest is not ready"),
407 ErrorInfo {
408 reason: REASON_WORKER_NOT_READY.to_string(),
409 domain: crate::validate::INGEST_ERROR_DOMAIN.to_string(),
410 ..Default::default()
411 },
412 ),
413 RETRY_HINT_DELAY,
414 ));
415 }
416
417 validate::validate_put_request(&request, self.state.limits)?;
418
419 let wire = request.bytes();
420 let mut batch = Vec::with_capacity(request.kvs.len());
421 for kv in request.kvs.iter() {
422 let key: Key = wire.slice_ref(kv.key);
423 let value = wire.slice_ref(kv.value);
424 batch.push((key, value));
425 }
426
427 let seq = self
428 .state
429 .ingest
430 .put_batch(batch)
431 .await
432 .map_err(ingest_error_to_connect)?;
433
434 if let Some(notifier) = &self.state.notifier {
436 notifier.advance(seq);
437 }
438
439 connectrpc::Response::ok(ProtoPutResponse {
440 sequence_number: seq,
441 ..Default::default()
442 })
443 }
444}
445
446pub struct QueryConnect<Q> {
447 state: QueryState<Q>,
448}
449
450impl<Q> Clone for QueryConnect<Q> {
451 fn clone(&self) -> Self {
452 Self {
453 state: self.state.clone(),
454 }
455 }
456}
457
458impl<Q> QueryConnect<Q>
459where
460 Q: Query,
461{
462 pub fn new(state: impl Into<QueryState<Q>>) -> Self {
463 Self {
464 state: state.into(),
465 }
466 }
467
468 fn current_sequence_number(&self) -> u64 {
469 self.state.query.current_sequence()
470 }
471
472 fn error_detail(&self) -> Detail {
473 Detail {
474 sequence_number: self.current_sequence_number(),
475 ..Default::default()
476 }
477 }
478
479 fn consistency_not_ready_error(&self, required: u64, current: u64) -> ConnectError {
480 let err = with_retry_hint(
481 ConnectError::aborted("minimum consistency token is not yet visible"),
482 RETRY_HINT_DELAY,
483 );
484 with_query_detail(
485 with_error_info_detail(
486 err,
487 ErrorInfo {
488 reason: "CONSISTENCY_NOT_READY".to_string(),
489 domain: "store.query".to_string(),
490 metadata: [
491 ("required_sequence_number".to_string(), required.to_string()),
492 ("current_sequence_number".to_string(), current.to_string()),
493 ]
494 .into_iter()
495 .collect(),
496 ..Default::default()
497 },
498 ),
499 self.error_detail(),
500 )
501 }
502
503 fn ensure_min_sequence_number(&self, required: Option<u64>) -> Result<u64, ConnectError> {
504 let current = self.current_sequence_number();
505 if let Some(required) = required {
506 if current < required {
507 return Err(self.consistency_not_ready_error(required, current));
508 }
509 }
510 Ok(current)
511 }
512}
513
514impl<Q> QueryApi for QueryConnect<Q>
515where
516 Q: Query,
517{
518 async fn get(
519 &self,
520 _ctx: Context,
521 request: buffa::view::OwnedView<exoware_proto::store::query::v1::GetRequestView<'static>>,
522 ) -> connectrpc::ServiceResult<GetResponse> {
523 validate::validate_get_request(&request)?;
524 let token = self.ensure_min_sequence_number(request.min_sequence_number)?;
525 let wire = request.bytes();
526 let key: Key = wire.slice_ref(request.key);
527 let (value, extra) = self
528 .state
529 .query
530 .get(key)
531 .await
532 .map_err(ConnectError::internal)?;
533 let detail = query_detail(token, extra);
534 connectrpc::Response::ok(GetResponse {
535 value,
536 detail: Some(detail).into(),
537 ..Default::default()
538 })
539 }
540
541 async fn get_many(
542 &self,
543 _ctx: Context,
544 request: buffa::view::OwnedView<
545 exoware_proto::store::query::v1::GetManyRequestView<'static>,
546 >,
547 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<GetManyFrame>> {
548 validate::validate_get_many_request(&request)?;
549 let sequence_number = self.ensure_min_sequence_number(request.min_sequence_number)?;
550
551 let wire = request.bytes();
552 let keys: Vec<Key> = request.keys.iter().map(|key| wire.slice_ref(key)).collect();
553 let (entries, extra) = self
554 .state
555 .query
556 .get_many(keys)
557 .await
558 .map_err(ConnectError::internal)?;
559 let detail = query_detail(sequence_number, extra);
560 let batch_size = (request.batch_size as usize).min(RANGE_STREAM_MAX_FRAME_ROWS);
561 let mut frames = Vec::new();
562 let mut chunk = Vec::new();
563 for (key, value) in entries {
564 chunk.push(GetManyEntry {
565 key: key.to_vec(),
566 value,
567 ..Default::default()
568 });
569 if chunk.len() >= batch_size {
570 frames.push(Ok(GetManyFrame {
571 results: std::mem::take(&mut chunk),
572 detail: Some(detail.clone()).into(),
573 ..Default::default()
574 }));
575 }
576 }
577 if !chunk.is_empty() {
578 frames.push(Ok(GetManyFrame {
579 results: chunk,
580 detail: Some(detail).into(),
581 ..Default::default()
582 }));
583 } else if frames.is_empty() {
584 frames.push(Ok(GetManyFrame {
585 detail: Some(detail).into(),
586 ..Default::default()
587 }));
588 }
589
590 Ok(connectrpc::Response::stream(stream_util::iter(frames)))
591 }
592
593 async fn range(
594 &self,
595 _ctx: Context,
596 request: buffa::view::OwnedView<exoware_proto::store::query::v1::RangeRequestView<'static>>,
597 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<RangeFrame>> {
598 validate::validate_range_request(&request)?;
599 let sequence_number = self.ensure_min_sequence_number(request.min_sequence_number)?;
600 let wire = request.bytes();
601 let start_key: Key = wire.slice_ref(request.start);
602 let end_key: Key = wire.slice_ref(request.end);
603 let limit = request.limit.map(|v| v as usize).unwrap_or(usize::MAX);
604 let batch_size = (request.batch_size as usize).min(RANGE_STREAM_MAX_FRAME_ROWS);
605 let forward = match parse_range_traversal_direction(request.mode) {
606 Ok(RangeTraversalDirection::Forward) => true,
607 Ok(RangeTraversalDirection::Reverse) => false,
608 Err(e) => return Err(ConnectError::internal(format!("traversal mode: {e:?}"))),
609 };
610 Ok(connectrpc::Response::stream(
611 range_stream(
612 self.state.query.clone(),
613 RangeStreamRequest {
614 start_key,
615 end_key,
616 limit,
617 batch_size,
618 forward,
619 sequence_number,
620 },
621 )
622 .await?,
623 ))
624 }
625
626 async fn reduce(
627 &self,
628 _ctx: Context,
629 request: buffa::view::OwnedView<
630 exoware_proto::store::query::v1::ReduceRequestView<'static>,
631 >,
632 ) -> connectrpc::ServiceResult<ReduceResponse> {
633 validate::validate_reduce_request(&request)?;
634 let token = self.ensure_min_sequence_number(request.min_sequence_number)?;
635 let wire = request.bytes();
636 let start_key: Key = wire.slice_ref(request.start);
637 let end_key: Key = wire.slice_ref(request.end);
638 let domain = to_domain_reduce_request_from_view(&request.params)
639 .map_err(validate::reduce_params_error)?;
640
641 let mut rows = self
642 .state
643 .query
644 .range_scan(start_key, end_key, usize::MAX, true)
645 .await
646 .map_err(ConnectError::internal)?;
647
648 let mut reducer = RangeReducer::new(&domain)
649 .map_err(|e: crate::RangeError| ConnectError::internal(e.to_string()))?;
650 let mut latest_extra = None;
651 let final_extra = loop {
652 let batch = rows
653 .next_batch(REDUCE_SCAN_BATCH_SIZE)
654 .await
655 .map_err(ConnectError::internal)?;
656 if batch.rows.is_empty() {
657 break if batch.extra.is_empty() {
658 latest_extra.unwrap_or_default()
659 } else {
660 batch.extra
661 };
662 }
663 latest_extra = Some(batch.extra);
664 for (key, value) in batch.rows {
665 reducer
666 .update(&key, &value)
667 .map_err(|e: crate::RangeError| ConnectError::internal(e.to_string()))?;
668 }
669 };
670 let response = reducer.finish();
671
672 let detail = query_detail(token, final_extra);
673
674 connectrpc::Response::ok(ReduceResponse {
675 results: response
676 .results
677 .into_iter()
678 .map(|result| exoware_proto::query::RangeReduceResult {
679 value: result.value.map(to_proto_reduced_value).into(),
680 ..Default::default()
681 })
682 .collect(),
683 groups: response
684 .groups
685 .into_iter()
686 .map(|group| {
687 let group_values_present =
688 group.group_values.iter().map(Option::is_some).collect();
689 exoware_proto::query::RangeReduceGroup {
690 group_values: group
691 .group_values
692 .into_iter()
693 .map(to_proto_optional_reduced_value)
694 .collect(),
695 group_values_present,
696 results: group
697 .results
698 .into_iter()
699 .map(|result| exoware_proto::query::RangeReduceResult {
700 value: result.value.map(to_proto_reduced_value).into(),
701 ..Default::default()
702 })
703 .collect(),
704 ..Default::default()
705 }
706 })
707 .collect(),
708 detail: Some(detail).into(),
709 ..Default::default()
710 })
711 }
712}
713
714pub struct CompactConnect<P> {
715 state: CompactState<P>,
716}
717
718impl<P> Clone for CompactConnect<P> {
719 fn clone(&self) -> Self {
720 Self {
721 state: self.state.clone(),
722 }
723 }
724}
725
726impl<P> CompactConnect<P>
727where
728 P: Prune,
729{
730 pub fn new(state: impl Into<CompactState<P>>) -> Self {
731 Self {
732 state: state.into(),
733 }
734 }
735}
736
737impl<P> CompactApi for CompactConnect<P>
738where
739 P: Prune,
740{
741 async fn prune(
742 &self,
743 _ctx: Context,
744 request: buffa::view::OwnedView<
745 exoware_proto::store::compact::v1::PruneRequestView<'static>,
746 >,
747 ) -> connectrpc::ServiceResult<PruneResponse> {
748 validate::validate_prune_request(&request)?;
749 let document = exoware_proto::parse_and_validate_policy_document(&request)
750 .map_err(|e| ConnectError::invalid_argument(e.to_string()))?;
751
752 self.state
753 .prune
754 .apply_prune_policies(document)
755 .await
756 .map_err(ConnectError::internal)?;
757 connectrpc::Response::ok(PruneResponse::default())
758 }
759}
760
761pub struct StreamConnect<B> {
762 state: StreamState<B>,
763}
764
765impl<B> Clone for StreamConnect<B> {
766 fn clone(&self) -> Self {
767 Self {
768 state: self.state.clone(),
769 }
770 }
771}
772
773impl<B> StreamConnect<B>
774where
775 B: Log,
776{
777 pub fn new(state: impl Into<StreamState<B>>) -> Self {
778 Self {
779 state: state.into(),
780 }
781 }
782
783 fn batch_evicted_connect_error(oldest_retained: Option<u64>) -> ConnectError {
784 let mut metadata = HashMap::new();
785 if let Some(v) = oldest_retained {
786 metadata.insert(
787 crate::stream::METADATA_OLDEST_RETAINED.to_string(),
788 v.to_string(),
789 );
790 }
791 with_error_info_detail(
792 ConnectError::out_of_range("batch has been evicted from the log"),
793 ErrorInfo {
794 reason: crate::stream::REASON_BATCH_EVICTED.to_string(),
795 domain: crate::stream::STREAM_ERROR_DOMAIN.to_string(),
796 metadata,
797 ..Default::default()
798 },
799 )
800 }
801
802 fn batch_evicted_error(&self, oldest_retained: Option<u64>) -> ConnectError {
803 Self::batch_evicted_connect_error(oldest_retained)
804 }
805
806 fn batch_not_found_error(&self) -> ConnectError {
807 with_error_info_detail(
808 ConnectError::not_found("batch not found"),
809 ErrorInfo {
810 reason: crate::stream::REASON_BATCH_NOT_FOUND.to_string(),
811 domain: crate::stream::STREAM_ERROR_DOMAIN.to_string(),
812 ..Default::default()
813 },
814 )
815 }
816}
817
818fn filtered_subscribe_response(
819 batch: &LogBatch,
820 matchers: &crate::stream::CompiledMatchers,
821) -> Result<Option<SubscribeResponse>, ConnectError> {
822 let response = batch.decode_response().map_err(ConnectError::internal)?;
823 let entries = crate::stream::apply_filter(matchers, &response.entries);
824 Ok((!entries.is_empty()).then_some(SubscribeResponse {
825 sequence_number: batch.sequence_number(),
826 entries,
827 ..Default::default()
828 }))
829}
830
831struct ReplayState {
832 next_sequence: u64,
833 bound: u64,
834 first_batch: Option<LogBatch>,
835}
836
837enum ReplayProgress {
838 Frame(SubscribeResponse),
839 Advanced,
840 Done,
841}
842
843enum LiveProgress {
844 Frame(SubscribeResponse),
845 Advanced,
846 NeedWait,
847}
848
849struct SubscriptionState<B> {
850 state: StreamState<B>,
851 matchers: crate::stream::CompiledMatchers,
852 replay: Option<ReplayState>,
853 next_live_sequence: u64,
854 live_notify: Arc<Notify>,
855 terminated: bool,
856}
857
858impl<B> SubscriptionState<B>
859where
860 B: Log,
861{
862 fn new(
863 state: StreamState<B>,
864 matchers: crate::stream::CompiledMatchers,
865 replay: Option<ReplayState>,
866 next_live_sequence: u64,
867 live_notify: Arc<Notify>,
868 ) -> Self {
869 Self {
870 state,
871 matchers,
872 replay,
873 next_live_sequence,
874 live_notify,
875 terminated: false,
876 }
877 }
878
879 fn into_stream(
880 self,
881 ) -> Pin<Box<dyn Stream<Item = Result<SubscribeResponse, ConnectError>> + Send>> {
882 Box::pin(stream_util::unfold(self, |mut state| async move {
883 loop {
884 if state.terminated {
885 return None;
886 }
887
888 if state.replay.is_some() {
889 match state.next_replay_frame().await {
890 Ok(ReplayProgress::Frame(frame)) => return Some((Ok(frame), state)),
891 Ok(ReplayProgress::Advanced) => continue,
892 Ok(ReplayProgress::Done) => {}
893 Err(err) => {
894 state.terminated = true;
895 return Some((Err(err), state));
896 }
897 }
898 }
899
900 match state.next_live_frame().await {
901 Ok(LiveProgress::Frame(frame)) => return Some((Ok(frame), state)),
902 Ok(LiveProgress::Advanced) => continue,
903 Ok(LiveProgress::NeedWait) => {
904 state.wait_for_live().await;
905 }
906 Err(err) => {
907 state.terminated = true;
908 return Some((Err(err), state));
909 }
910 }
911 }
912 }))
913 }
914
915 async fn next_replay_frame(&mut self) -> Result<ReplayProgress, ConnectError> {
916 let Some(replay) = &mut self.replay else {
917 return Ok(ReplayProgress::Done);
918 };
919 let seq = replay.next_sequence;
920 let batch = if let Some(first_batch) = replay.first_batch.take() {
921 Some(first_batch)
922 } else {
923 self.state
924 .log
925 .get_batch(seq)
926 .await
927 .map_err(ConnectError::internal)?
928 };
929 replay.next_sequence += 1;
930 if replay.next_sequence > replay.bound {
931 self.replay = None;
932 }
933 let Some(batch) = batch else {
934 let oldest = self
935 .state
936 .log
937 .oldest_retained_batch()
938 .await
939 .map_err(ConnectError::internal)?;
940 return Err(StreamConnect::<B>::batch_evicted_connect_error(oldest));
941 };
942 Ok(match filtered_subscribe_response(&batch, &self.matchers)? {
943 Some(frame) => ReplayProgress::Frame(frame),
944 None => ReplayProgress::Advanced,
945 })
946 }
947
948 async fn next_live_frame(&mut self) -> Result<LiveProgress, ConnectError> {
949 let current = self.state.notifier.current_sequence();
950 if self.next_live_sequence > current {
951 return Ok(LiveProgress::NeedWait);
952 }
953 let seq = self.next_live_sequence;
954 self.next_live_sequence += 1;
955 let batch = self
956 .state
957 .log
958 .get_batch(seq)
959 .await
960 .map_err(ConnectError::internal)?;
961 let Some(batch) = batch else {
962 let oldest = self
963 .state
964 .log
965 .oldest_retained_batch()
966 .await
967 .map_err(ConnectError::internal)?;
968 return Err(StreamConnect::<B>::batch_evicted_connect_error(oldest));
969 };
970 Ok(match filtered_subscribe_response(&batch, &self.matchers)? {
971 Some(frame) => LiveProgress::Frame(frame),
972 None => LiveProgress::Advanced,
973 })
974 }
975
976 async fn wait_for_live(&self) {
977 if self.next_live_sequence <= self.state.notifier.current_sequence() {
978 return;
979 }
980 let notified = self.live_notify.clone().notified_owned();
981 if self.next_live_sequence <= self.state.notifier.current_sequence() {
982 return;
983 }
984 notified.await;
985 }
986}
987
988fn domain_filter_from_subscribe_view(
989 req: &SubscribeRequestView<'_>,
990) -> Result<StreamFilter, ConnectError> {
991 let mut selectors = Vec::with_capacity(req.selectors.len());
992 for mk in req.selectors.iter() {
993 selectors.push(Selector {
996 prefix: Bytes::copy_from_slice(mk.prefix),
997 payload_regex: exoware_sdk::kv_codec::Utf8::from(mk.payload_regex),
998 });
999 }
1000 let mut value_filters = Vec::with_capacity(req.value_filters.len());
1001 for vf in req.value_filters.iter() {
1002 value_filters.push(match vf.kind {
1003 Some(ProtoFilterKindView::Exact(bytes)) => Filter::Exact(Bytes::copy_from_slice(bytes)),
1004 Some(ProtoFilterKindView::Prefix(bytes)) => {
1005 Filter::Prefix(Bytes::copy_from_slice(bytes))
1006 }
1007 Some(ProtoFilterKindView::Regex(pattern)) => Filter::Regex(pattern.to_string()),
1008 None => {
1009 return Err(ConnectError::invalid_argument(
1010 "each value_filter must set exactly one of exact, prefix, or regex",
1011 ))
1012 }
1013 });
1014 }
1015 Ok(StreamFilter {
1016 selectors,
1017 value_filters,
1018 })
1019}
1020
1021impl<B> StreamApi for StreamConnect<B>
1022where
1023 B: Log,
1024{
1025 async fn subscribe(
1026 &self,
1027 _ctx: Context,
1028 request: buffa::view::OwnedView<SubscribeRequestView<'static>>,
1029 ) -> connectrpc::ServiceResult<connectrpc::ServiceStream<SubscribeResponse>> {
1030 let filter = domain_filter_from_subscribe_view(&request)?;
1031 let since = request.since_sequence_number;
1032
1033 let matchers = crate::stream::compile_matchers(&filter)?;
1038 let subscription = self.state.notifier.subscribe();
1039 let replay_bound = subscription.current_sequence;
1040 let live_notify = subscription.notify;
1041
1042 let replay = match since {
1047 Some(s) if s <= replay_bound && s > 0 => {
1048 let first_batch = self
1049 .state
1050 .log
1051 .get_batch(s)
1052 .await
1053 .map_err(ConnectError::internal)?;
1054 let Some(first_batch) = first_batch else {
1055 let oldest = self
1056 .state
1057 .log
1058 .oldest_retained_batch()
1059 .await
1060 .map_err(ConnectError::internal)?;
1061 return Err(self.batch_evicted_error(oldest));
1062 };
1063 Some(ReplayState {
1064 next_sequence: s,
1065 bound: replay_bound,
1066 first_batch: Some(first_batch),
1067 })
1068 }
1069 _ => None,
1070 };
1071 let next_live_sequence = replay_bound.saturating_add(1);
1072
1073 Ok(connectrpc::Response::stream(
1074 SubscriptionState::new(
1075 self.state.clone(),
1076 matchers,
1077 replay,
1078 next_live_sequence,
1079 live_notify,
1080 )
1081 .into_stream(),
1082 ))
1083 }
1084
1085 async fn get(
1086 &self,
1087 _ctx: Context,
1088 request: buffa::view::OwnedView<GetRequestView<'static>>,
1089 ) -> connectrpc::ServiceResult<PreEncoded<StreamGetResponse>> {
1090 let seq = request.sequence_number;
1091 match self
1092 .state
1093 .log
1094 .get_batch(seq)
1095 .await
1096 .map_err(ConnectError::internal)?
1097 {
1098 Some(batch) => connectrpc::Response::ok(PreEncoded::from_bytes_unchecked(
1099 batch.into_response_bytes(),
1100 )),
1101 None => {
1102 let current = self.state.log.current_sequence();
1103 if seq > current {
1105 Err(self.batch_not_found_error())
1106 } else {
1107 let oldest = self
1108 .state
1109 .log
1110 .oldest_retained_batch()
1111 .await
1112 .map_err(ConnectError::internal)?;
1113 Err(self.batch_evicted_error(oldest))
1114 }
1115 }
1116 }
1117 }
1118}
1119
1120fn connect_limits() -> Limits {
1121 Limits::default()
1122 .max_request_body_size(MAX_CONNECTRPC_BODY_BYTES)
1123 .max_message_size(MAX_CONNECTRPC_BODY_BYTES)
1124}
1125
1126pub(crate) type IngestService<I> = ConnectRpcService<IngestServiceServer<IngestConnect<I>>>;
1127pub(crate) type QueryService<Q> = ConnectRpcService<QueryServiceServer<QueryConnect<Q>>>;
1128pub(crate) type CompactService<P> = ConnectRpcService<CompactServiceServer<CompactConnect<P>>>;
1129pub(crate) type StreamService<B> = ConnectRpcService<StreamServiceServer<StreamConnect<B>>>;
1130pub(crate) type QueryStack<Q, B> = ConnectRpcService<
1131 Chain<QueryServiceServer<QueryConnect<Q>>, StreamServiceServer<StreamConnect<B>>>,
1132>;
1133pub(crate) type ConnectStack<I, Q, P, B> = ConnectRpcService<
1134 Chain<
1135 IngestServiceServer<IngestConnect<I>>,
1136 Chain<
1137 QueryServiceServer<QueryConnect<Q>>,
1138 Chain<CompactServiceServer<CompactConnect<P>>, StreamServiceServer<StreamConnect<B>>>,
1139 >,
1140 >,
1141>;
1142
1143fn ingest_server<I>(state: IngestState<I>) -> IngestServiceServer<IngestConnect<I>>
1144where
1145 I: Ingest,
1146{
1147 IngestServiceServer::new(IngestConnect::new(state))
1148}
1149
1150fn query_server<Q>(state: QueryState<Q>) -> QueryServiceServer<QueryConnect<Q>>
1151where
1152 Q: Query,
1153{
1154 QueryServiceServer::new(QueryConnect::new(state))
1155}
1156
1157fn compact_server<P>(state: CompactState<P>) -> CompactServiceServer<CompactConnect<P>>
1158where
1159 P: Prune,
1160{
1161 CompactServiceServer::new(CompactConnect::new(state))
1162}
1163
1164fn stream_server<B>(state: StreamState<B>) -> StreamServiceServer<StreamConnect<B>>
1165where
1166 B: Log,
1167{
1168 StreamServiceServer::new(StreamConnect::new(state))
1169}
1170
1171pub fn ingest_service<I>(state: IngestState<I>) -> IngestService<I>
1172where
1173 I: Ingest,
1174{
1175 ConnectRpcService::new(ingest_server(state))
1176 .with_limits(connect_limits())
1177 .with_compression(connect_compression_registry())
1178}
1179
1180pub fn query_service<Q>(state: QueryState<Q>) -> QueryService<Q>
1181where
1182 Q: Query,
1183{
1184 ConnectRpcService::new(query_server(state))
1185 .with_limits(connect_limits())
1186 .with_compression(connect_compression_registry())
1187}
1188
1189pub fn compact_service<P>(state: CompactState<P>) -> CompactService<P>
1190where
1191 P: Prune,
1192{
1193 ConnectRpcService::new(compact_server(state))
1194 .with_limits(connect_limits())
1195 .with_compression(connect_compression_registry())
1196}
1197
1198pub fn stream_service<B>(state: StreamState<B>) -> StreamService<B>
1199where
1200 B: Log,
1201{
1202 ConnectRpcService::new(stream_server(state))
1203 .with_limits(connect_limits())
1204 .with_compression(connect_compression_registry())
1205}
1206
1207pub fn query_stack<Q, B>(
1208 query_state: QueryState<Q>,
1209 stream_state: StreamState<B>,
1210) -> QueryStack<Q, B>
1211where
1212 Q: Query,
1213 B: Log,
1214{
1215 ConnectRpcService::new(Chain(
1216 query_server(query_state),
1217 stream_server(stream_state),
1218 ))
1219 .with_limits(connect_limits())
1220 .with_compression(connect_compression_registry())
1221}
1222
1223pub fn connect_stack<E>(state: AppState<E>) -> ConnectStack<E, E, E, E>
1224where
1225 E: StoreEngine,
1226{
1227 ConnectRpcService::new(Chain(
1228 ingest_server(state.clone().into()),
1229 Chain(
1230 query_server(state.clone().into()),
1231 Chain(
1232 compact_server(state.clone().into()),
1233 stream_server(state.into()),
1234 ),
1235 ),
1236 ))
1237 .with_limits(connect_limits())
1238 .with_compression(connect_compression_registry())
1239}
1240
1241#[cfg(test)]
1242mod tests {
1243 use super::*;
1244 use std::collections::{BTreeMap, HashMap};
1245 use std::sync::atomic::AtomicU64;
1246 use std::sync::Mutex;
1247 use std::time::Duration;
1248
1249 use buffa::Message;
1250 use exoware_proto::common::kv::v1::Selector as ProtoSelector;
1251 use exoware_proto::log::stream::v1::{SubscribeRequest, SubscribeRequestView};
1252 use exoware_proto::store::compact::v1::{
1253 policy, policy_retain, Policy as ProtoPolicy, PolicyRetain, PruneRequest, PruneRequestView,
1254 RetainKeepLatest,
1255 };
1256 use exoware_sdk::keys::Prefix;
1257 use exoware_sdk::kv_codec::KvReducedValue;
1258 use exoware_sdk::prune_policy::{PrunePolicyDocument, PRUNE_POLICY_DOCUMENT_VERSION};
1259 use exoware_sdk::{decode_connect_error, to_domain_reduce_response};
1260 use futures::StreamExt;
1261
1262 use crate::{
1263 Ingest, IngestError, Log, Prune, Query, QueryExtra, RangeScan, RangeScanBatch, Sequence,
1264 StreamNotification, StreamNotifier,
1265 };
1266
1267 const TEST_PREFIX: u8 = 1;
1268
1269 #[derive(Clone)]
1270 struct PublishDuringReplay {
1271 hub: Arc<StreamHub>,
1272 sequence_offset: u64,
1273 kvs: Vec<(Bytes, Bytes)>,
1274 }
1275
1276 #[derive(Default)]
1277 struct FakeEngineState {
1278 current_sequence: u64,
1279 batches: BTreeMap<u64, Option<Vec<(Bytes, Bytes)>>>,
1280 oldest_retained: Option<u64>,
1281 publish_on_get_batch: Option<PublishDuringReplay>,
1282 range_rows: Vec<(Bytes, Bytes)>,
1283 range_eof_extra: QueryExtra,
1284 range_next_count: usize,
1285 query_extra: QueryExtra,
1286 prune_policy_counts: Vec<usize>,
1287 put_error: Option<IngestError>,
1288 }
1289
1290 #[derive(Default)]
1291 struct FakeEngine {
1292 state: Arc<Mutex<FakeEngineState>>,
1293 }
1294
1295 struct IteratorRangeScan {
1296 iter: Box<dyn Iterator<Item = Result<(Bytes, Bytes), String>> + Send + 'static>,
1297 eof_extra: Option<QueryExtra>,
1298 }
1299
1300 impl RangeScan for IteratorRangeScan {
1301 async fn next_batch(&mut self, max_items: usize) -> Result<RangeScanBatch, String> {
1302 let mut rows = Vec::new();
1303 for row in self.iter.by_ref().take(max_items) {
1304 rows.push(row?);
1305 }
1306 let extra = if rows.is_empty() {
1307 self.eof_extra.take().unwrap_or_default()
1308 } else {
1309 QueryExtra::default()
1310 };
1311 Ok(RangeScanBatch { rows, extra })
1312 }
1313 }
1314
1315 fn range_scan_from_iter<I>(iter: I) -> IteratorRangeScan
1316 where
1317 I: Iterator<Item = Result<(Bytes, Bytes), String>> + Send + 'static,
1318 {
1319 range_scan_from_iter_with_eof_extra(iter, QueryExtra::default())
1320 }
1321
1322 fn range_scan_from_iter_with_eof_extra<I>(iter: I, eof_extra: QueryExtra) -> IteratorRangeScan
1323 where
1324 I: Iterator<Item = Result<(Bytes, Bytes), String>> + Send + 'static,
1325 {
1326 IteratorRangeScan {
1327 iter: Box::new(iter),
1328 eof_extra: Some(eof_extra),
1329 }
1330 }
1331
1332 impl FakeEngine {
1333 fn set_current_sequence(&self, sequence_number: u64) {
1334 self.state.lock().expect("lock").current_sequence = sequence_number;
1335 }
1336
1337 fn set_put_error(&self, err: IngestError) {
1338 self.state.lock().expect("lock").put_error = Some(err);
1339 }
1340
1341 fn set_batch(&self, sequence_number: u64, kvs: Option<Vec<(Bytes, Bytes)>>) {
1342 self.state
1343 .lock()
1344 .expect("lock")
1345 .batches
1346 .insert(sequence_number, kvs);
1347 }
1348
1349 fn set_oldest_retained(&self, oldest_retained: Option<u64>) {
1350 self.state.lock().expect("lock").oldest_retained = oldest_retained;
1351 }
1352
1353 fn publish_live(
1354 &self,
1355 hub: Arc<StreamHub>,
1356 sequence_number: u64,
1357 kvs: Vec<(Bytes, Bytes)>,
1358 ) {
1359 let mut state = self.state.lock().expect("lock");
1360 state.current_sequence = state.current_sequence.max(sequence_number);
1361 state.batches.insert(sequence_number, Some(kvs.clone()));
1362 drop(state);
1363 hub.publish(sequence_number);
1364 }
1365
1366 fn publish_on_every_get_batch(
1367 &self,
1368 hub: Arc<StreamHub>,
1369 sequence_offset: u64,
1370 kvs: Vec<(Bytes, Bytes)>,
1371 ) {
1372 self.state.lock().expect("lock").publish_on_get_batch = Some(PublishDuringReplay {
1373 hub,
1374 sequence_offset,
1375 kvs,
1376 });
1377 }
1378
1379 fn set_range_rows(&self, rows: Vec<(Bytes, Bytes)>) {
1380 self.state.lock().expect("lock").range_rows = rows;
1381 }
1382
1383 fn set_range_eof_extra(&self, extra: QueryExtra) {
1384 self.state.lock().expect("lock").range_eof_extra = extra;
1385 }
1386
1387 fn range_next_count(&self) -> usize {
1388 self.state.lock().expect("lock").range_next_count
1389 }
1390
1391 fn set_query_extra(&self, extra: QueryExtra) {
1392 self.state.lock().expect("lock").query_extra = extra;
1393 }
1394 }
1395
1396 impl Sequence for FakeEngine {
1397 fn current_sequence(&self) -> u64 {
1398 self.state.lock().expect("lock").current_sequence
1399 }
1400 }
1401
1402 impl Ingest for FakeEngine {
1403 async fn put_batch(&self, kvs: Vec<(Bytes, Bytes)>) -> Result<u64, IngestError> {
1404 let mut state = self.state.lock().map_err(|e| IngestError::Internal {
1405 message: e.to_string(),
1406 })?;
1407 if let Some(err) = state.put_error.take() {
1408 return Err(err);
1409 }
1410 state.current_sequence += 1;
1411 let seq = state.current_sequence;
1412 state.batches.insert(seq, Some(kvs));
1413 Ok(seq)
1414 }
1415 }
1416
1417 impl Query for FakeEngine {
1418 type RangeScan = IteratorRangeScan;
1419
1420 async fn get(&self, _key: Bytes) -> Result<(Option<Bytes>, QueryExtra), String> {
1421 self.state
1422 .lock()
1423 .map(|state| (None, state.query_extra.clone()))
1424 .map_err(|e| e.to_string())
1425 }
1426
1427 async fn get_many(
1428 &self,
1429 keys: Vec<Bytes>,
1430 ) -> Result<(Vec<(Bytes, Option<Bytes>)>, QueryExtra), String> {
1431 self.state
1432 .lock()
1433 .map(|state| {
1434 let entries = keys.into_iter().map(|key| (key, None)).collect();
1435 (entries, state.query_extra.clone())
1436 })
1437 .map_err(|e| e.to_string())
1438 }
1439
1440 async fn range_scan(
1441 &self,
1442 _start: Bytes,
1443 _end: Bytes,
1444 _limit: usize,
1445 _forward: bool,
1446 ) -> Result<Self::RangeScan, String> {
1447 let result = self
1448 .state
1449 .lock()
1450 .map(|state| (state.range_rows.clone(), state.range_eof_extra.clone()))
1451 .map_err(|e| e.to_string());
1452 let state = self.state.clone();
1453 let cursor = result.map(|(rows, eof_extra)| {
1454 range_scan_from_iter_with_eof_extra(
1455 rows.into_iter().map(move |row| {
1456 state.lock().expect("lock").range_next_count += 1;
1457 Ok(row)
1458 }),
1459 eof_extra,
1460 )
1461 });
1462 cursor
1463 }
1464 }
1465
1466 impl Prune for FakeEngine {
1467 async fn apply_prune_policies(&self, document: PrunePolicyDocument) -> Result<(), String> {
1468 self.state
1469 .lock()
1470 .map(|mut state| {
1471 state.prune_policy_counts.push(document.policies.len());
1472 })
1473 .map_err(|e| e.to_string())
1474 }
1475 }
1476
1477 impl Log for FakeEngine {
1478 async fn get_batch(&self, sequence_number: u64) -> Result<Option<LogBatch>, String> {
1479 let result: Result<_, String> = (|| {
1480 let mut state = self.state.lock().map_err(|e| e.to_string())?;
1481 let publish = state.publish_on_get_batch.clone();
1482 if let Some(publish) = publish.as_ref() {
1483 let live_sequence = publish.sequence_offset + sequence_number;
1484 state.current_sequence = state.current_sequence.max(live_sequence);
1485 state
1486 .batches
1487 .entry(live_sequence)
1488 .or_insert_with(|| Some(publish.kvs.clone()));
1489 }
1490 Ok((
1491 publish,
1492 state.batches.get(&sequence_number).cloned().unwrap_or(None),
1493 ))
1494 })();
1495 let (publish, batch) = result?;
1496 if let Some(publish) = publish {
1497 publish
1498 .hub
1499 .publish(publish.sequence_offset + sequence_number);
1500 }
1501 Ok(batch.map(|kvs| LogBatch::from_entries(sequence_number, kvs)))
1502 }
1503
1504 async fn oldest_retained_batch(&self) -> Result<Option<u64>, String> {
1505 self.state
1506 .lock()
1507 .map(|state| state.oldest_retained)
1508 .map_err(|e| e.to_string())
1509 }
1510 }
1511
1512 struct QueryOnlyEngine {
1513 sequence_number: u64,
1514 value: Option<Bytes>,
1515 }
1516
1517 impl Sequence for QueryOnlyEngine {
1518 fn current_sequence(&self) -> u64 {
1519 self.sequence_number
1520 }
1521 }
1522
1523 impl Query for QueryOnlyEngine {
1524 type RangeScan = IteratorRangeScan;
1525
1526 async fn get(&self, _key: Bytes) -> Result<(Option<Bytes>, QueryExtra), String> {
1527 Ok((self.value.clone(), QueryExtra::default()))
1528 }
1529
1530 async fn range_scan(
1531 &self,
1532 _start: Bytes,
1533 _end: Bytes,
1534 _limit: usize,
1535 _forward: bool,
1536 ) -> Result<Self::RangeScan, String> {
1537 Ok(range_scan_from_iter(std::iter::empty()))
1538 }
1539
1540 async fn get_many(
1541 &self,
1542 keys: Vec<Bytes>,
1543 ) -> Result<(Vec<(Bytes, Option<Bytes>)>, QueryExtra), String> {
1544 Ok((
1545 keys.into_iter().map(|key| (key, None)).collect(),
1546 QueryExtra::default(),
1547 ))
1548 }
1549 }
1550
1551 #[derive(Default)]
1552 struct PruneOnlyEngine {
1553 documents: Mutex<Vec<(u32, usize)>>,
1554 }
1555
1556 impl PruneOnlyEngine {
1557 fn applied_count(&self) -> usize {
1558 self.documents.lock().expect("lock").len()
1559 }
1560
1561 fn last_document(&self) -> Option<(u32, usize)> {
1562 self.documents.lock().expect("lock").last().copied()
1563 }
1564 }
1565
1566 impl Prune for PruneOnlyEngine {
1567 async fn apply_prune_policies(&self, document: PrunePolicyDocument) -> Result<(), String> {
1568 self.documents
1569 .lock()
1570 .map(|mut documents| {
1571 documents.push((document.version, document.policies.len()));
1572 })
1573 .map_err(|e| e.to_string())
1574 }
1575 }
1576
1577 struct ManualNotifier {
1578 current_sequence: AtomicU64,
1579 notify: Arc<Notify>,
1580 }
1581
1582 impl ManualNotifier {
1583 fn new(current_sequence: u64) -> Self {
1584 Self {
1585 current_sequence: AtomicU64::new(current_sequence),
1586 notify: Arc::new(Notify::new()),
1587 }
1588 }
1589 }
1590
1591 impl StreamNotifier for ManualNotifier {
1592 fn subscribe(&self) -> StreamNotification {
1593 StreamNotification {
1594 current_sequence: self.current_sequence.load(Ordering::Acquire),
1595 notify: self.notify.clone(),
1596 }
1597 }
1598
1599 fn current_sequence(&self) -> u64 {
1600 self.current_sequence.load(Ordering::Acquire)
1601 }
1602
1603 fn advance(&self, seq: u64) {
1604 self.current_sequence.fetch_max(seq, Ordering::SeqCst);
1605 self.notify.notify_waiters();
1606 }
1607 }
1608
1609 fn matching_kv(payload: &[u8], value: &[u8]) -> (Bytes, Bytes) {
1610 let key = Prefix::from_byte(TEST_PREFIX)
1611 .encode(payload)
1612 .expect("encode key");
1613 (key, Bytes::copy_from_slice(value))
1614 }
1615
1616 fn numeric_query_extra(name: &str, value: f64) -> QueryExtra {
1617 HashMap::from([(
1618 name.to_string(),
1619 buffa_types::google::protobuf::Value::from(value),
1620 )])
1621 }
1622
1623 fn subscribe_request_bytes(since_sequence_number: Option<u64>) -> Vec<u8> {
1624 SubscribeRequest {
1625 selectors: vec![ProtoSelector {
1626 prefix: Bytes::from(vec![TEST_PREFIX]),
1627 payload_regex: "(?s).*".to_string(),
1628 ..Default::default()
1629 }],
1630 since_sequence_number,
1631 ..Default::default()
1632 }
1633 .encode_to_vec()
1634 }
1635
1636 fn put_request(
1637 value_len: usize,
1638 ) -> buffa::view::OwnedView<exoware_proto::log::ingest::v1::PutRequestView<'static>> {
1639 let bytes = exoware_proto::ingest::PutRequest {
1640 kvs: vec![exoware_proto::common::Entry {
1641 key: b"k".to_vec(),
1642 value: Bytes::from(vec![1u8; value_len]),
1643 ..Default::default()
1644 }],
1645 ..Default::default()
1646 }
1647 .encode_to_vec();
1648 buffa::view::OwnedView::<exoware_proto::log::ingest::v1::PutRequestView<'static>>::decode(
1649 bytes.into(),
1650 )
1651 .expect("decode put request")
1652 }
1653
1654 fn sequence_drop_all_policy() -> ProtoPolicy {
1655 ProtoPolicy {
1656 scope: Some(policy::Scope::Sequence(Box::default())),
1657 retain: Some(PolicyRetain {
1658 kind: Some(policy_retain::Kind::DropAll(Box::default())),
1659 ..Default::default()
1660 })
1661 .into(),
1662 ..Default::default()
1663 }
1664 }
1665
1666 fn sequence_keep_latest_policy(count: u64) -> ProtoPolicy {
1667 ProtoPolicy {
1668 scope: Some(policy::Scope::Sequence(Box::default())),
1669 retain: Some(PolicyRetain {
1670 kind: Some(policy_retain::Kind::KeepLatest(Box::new(
1671 RetainKeepLatest {
1672 count,
1673 ..Default::default()
1674 },
1675 ))),
1676 ..Default::default()
1677 })
1678 .into(),
1679 ..Default::default()
1680 }
1681 }
1682
1683 fn prune_request(
1684 policies: Vec<ProtoPolicy>,
1685 ) -> buffa::view::OwnedView<PruneRequestView<'static>> {
1686 let bytes = PruneRequest {
1687 policies,
1688 ..Default::default()
1689 }
1690 .encode_to_vec();
1691 buffa::view::OwnedView::<PruneRequestView<'static>>::decode(bytes.into())
1692 .expect("decode prune request")
1693 }
1694
1695 async fn subscribe_stream<B>(
1696 connect: &StreamConnect<B>,
1697 since_sequence_number: Option<u64>,
1698 ) -> Result<
1699 Pin<Box<dyn Stream<Item = Result<SubscribeResponse, ConnectError>> + Send>>,
1700 ConnectError,
1701 >
1702 where
1703 B: Log,
1704 {
1705 let bytes = subscribe_request_bytes(since_sequence_number);
1706 let request = buffa::view::OwnedView::<SubscribeRequestView<'static>>::decode(bytes.into())
1707 .expect("decode subscribe request");
1708 Ok(StreamApi::subscribe(connect, Context::default(), request)
1709 .await?
1710 .body)
1711 }
1712
1713 #[tokio::test]
1714 async fn compact_connect_accepts_prune_only_engine() {
1715 let prune = Arc::new(PruneOnlyEngine::default());
1716 let connect = CompactConnect::new(CompactState::new(prune.clone()));
1717 let request = prune_request(vec![sequence_drop_all_policy()]);
1718
1719 CompactApi::prune(&connect, Context::default(), request)
1720 .await
1721 .expect("prune");
1722
1723 assert_eq!(prune.applied_count(), 1);
1724 assert_eq!(
1725 prune.last_document(),
1726 Some((PRUNE_POLICY_DOCUMENT_VERSION, 1))
1727 );
1728 }
1729
1730 #[tokio::test]
1731 async fn compact_rejects_unparseable_policy_before_engine_prune() {
1732 let prune = Arc::new(PruneOnlyEngine::default());
1733 let connect = CompactConnect::new(CompactState::new(prune.clone()));
1734 let invalid_policy = ProtoPolicy {
1735 scope: Some(policy::Scope::Sequence(Box::default())),
1736 ..Default::default()
1737 };
1738 let request = prune_request(vec![invalid_policy]);
1739
1740 let err = CompactApi::prune(&connect, Context::default(), request)
1741 .await
1742 .expect_err("invalid prune");
1743
1744 assert_eq!(err.code, connectrpc::ErrorCode::InvalidArgument);
1745 assert_eq!(prune.applied_count(), 0);
1746 }
1747
1748 #[tokio::test]
1749 async fn compact_rejects_invalid_policy_before_engine_prune() {
1750 let prune = Arc::new(PruneOnlyEngine::default());
1751 let connect = CompactConnect::new(CompactState::new(prune.clone()));
1752 let request = prune_request(vec![sequence_keep_latest_policy(0)]);
1753
1754 let err = CompactApi::prune(&connect, Context::default(), request)
1755 .await
1756 .expect_err("invalid prune");
1757
1758 assert_eq!(err.code, connectrpc::ErrorCode::InvalidArgument);
1759 assert_eq!(prune.applied_count(), 0);
1760 }
1761
1762 #[tokio::test]
1763 async fn query_connect_accepts_query_only_engine() {
1764 let query = Arc::new(QueryOnlyEngine {
1765 sequence_number: 9,
1766 value: Some(Bytes::from_static(b"value")),
1767 });
1768 let connect = QueryConnect::new(QueryState { query });
1769 let bytes = exoware_proto::query::GetRequest {
1770 key: b"k".to_vec(),
1771 ..Default::default()
1772 }
1773 .encode_to_vec();
1774 let request = buffa::view::OwnedView::<
1775 exoware_proto::store::query::v1::GetRequestView<'static>,
1776 >::decode(bytes.into())
1777 .expect("decode get request");
1778
1779 let response = QueryApi::get(&connect, Context::default(), request)
1780 .await
1781 .expect("get")
1782 .body;
1783 let detail = response.detail.as_option().expect("query detail");
1784
1785 assert_eq!(response.value.as_deref(), Some(b"value".as_slice()));
1786 assert_eq!(detail.sequence_number, 9);
1787 }
1788
1789 #[tokio::test]
1790 async fn get_includes_engine_query_extra() {
1791 let engine = Arc::new(FakeEngine::default());
1792 engine.set_current_sequence(5);
1793 engine.set_query_extra(HashMap::from([(
1794 "scanned_bytes".to_string(),
1795 buffa_types::google::protobuf::Value::from(123.0),
1796 )]));
1797 let connect = QueryConnect::new(AppState::new(engine));
1798 let bytes = exoware_proto::query::GetRequest {
1799 key: b"k".to_vec(),
1800 ..Default::default()
1801 }
1802 .encode_to_vec();
1803 let request = buffa::view::OwnedView::<
1804 exoware_proto::store::query::v1::GetRequestView<'static>,
1805 >::decode(bytes.into())
1806 .expect("decode get request");
1807
1808 let response = QueryApi::get(&connect, Context::default(), request)
1809 .await
1810 .expect("get")
1811 .body;
1812 let detail = response.detail.as_option().expect("query detail");
1813
1814 assert_eq!(detail.sequence_number, 5);
1815 assert_eq!(
1816 detail
1817 .extra
1818 .get("scanned_bytes")
1819 .and_then(|v| v.as_number()),
1820 Some(123.0)
1821 );
1822 }
1823
1824 #[test]
1825 fn split_service_constructors_build_independent_process_surfaces() {
1826 let engine = Arc::new(FakeEngine::default());
1827 let state = AppState::new(engine);
1828
1829 let _ingest = ingest_service(state.clone().into());
1830 let _query = query_service(state.clone().into());
1831 let _compact = compact_service(state.clone().into());
1832 let _stream = stream_service(state.clone().into());
1833 let _query_stack = query_stack(state.clone().into(), state.into());
1834 }
1835
1836 #[tokio::test]
1837 async fn ingest_uses_configured_value_limit() {
1838 let engine = Arc::new(FakeEngine::default());
1839 let state = IngestState::new(engine).with_limits(IngestLimits { max_value_len: 4 });
1840 let connect = IngestConnect::new(state);
1841
1842 let err = IngestApi::put(&connect, Context::default(), put_request(5))
1843 .await
1844 .expect_err("put should reject oversized value");
1845
1846 assert_eq!(err.code, connectrpc::ErrorCode::InvalidArgument);
1847 }
1848
1849 #[tokio::test]
1850 async fn ingest_unavailable_surfaces_as_unavailable_with_retry_info() {
1851 let engine = Arc::new(FakeEngine::default());
1852 engine.set_put_error(IngestError::Unavailable {
1853 message: "backend bouncing".to_string(),
1854 });
1855 let connect = IngestConnect::new(IngestState::new(engine));
1856
1857 let err = IngestApi::put(&connect, Context::default(), put_request(1))
1858 .await
1859 .expect_err("transient put failure should surface");
1860
1861 assert_eq!(err.code, connectrpc::ErrorCode::Unavailable);
1862 let decoded = decode_connect_error(&err).expect("decode details");
1863 assert_eq!(
1864 decoded.error_info.expect("error info").reason,
1865 REASON_INGEST_UNAVAILABLE
1866 );
1867 let retry_delay = decoded.retry_info.expect("retry info").retry_delay;
1868 let retry_delay = retry_delay.as_option().expect("retry delay");
1869 assert_eq!((retry_delay.seconds, retry_delay.nanos), (1, 0));
1870 }
1871
1872 #[tokio::test]
1873 async fn ingest_internal_surfaces_as_internal_without_error_info() {
1874 let engine = Arc::new(FakeEngine::default());
1875 engine.set_put_error(IngestError::Internal {
1876 message: "invariant violated".to_string(),
1877 });
1878 let connect = IngestConnect::new(IngestState::new(engine));
1879
1880 let err = IngestApi::put(&connect, Context::default(), put_request(1))
1881 .await
1882 .expect_err("fatal put failure should surface");
1883
1884 assert_eq!(err.code, connectrpc::ErrorCode::Internal);
1885 let decoded = decode_connect_error(&err).expect("decode details");
1886 assert!(decoded.error_info.is_none());
1887 }
1888
1889 #[tokio::test]
1890 async fn put_when_not_ready_keeps_worker_not_ready_reason() {
1891 let engine = Arc::new(FakeEngine::default());
1892 let state = IngestState::new(engine);
1893 state.ready.store(false, Ordering::SeqCst);
1894 let connect = IngestConnect::new(state);
1895
1896 let err = IngestApi::put(&connect, Context::default(), put_request(1))
1897 .await
1898 .expect_err("not-ready gate should reject");
1899
1900 assert_eq!(err.code, connectrpc::ErrorCode::Unavailable);
1903 let decoded = decode_connect_error(&err).expect("decode details");
1904 assert_eq!(
1905 decoded.error_info.expect("error info").reason,
1906 REASON_WORKER_NOT_READY
1907 );
1908 }
1909
1910 #[tokio::test]
1911 async fn stream_can_be_advanced_by_external_notifier() {
1912 let engine = Arc::new(FakeEngine::default());
1913 let notifier = Arc::new(ManualNotifier::new(0));
1914 let connect = StreamConnect::new(StreamState::new(engine.clone(), notifier.clone()));
1915 let mut stream = subscribe_stream(&connect, None).await.expect("subscribe");
1916
1917 engine.set_current_sequence(1);
1918 engine.set_batch(1, Some(vec![matching_kv(b"hit", b"v1")]));
1919 notifier.advance(1);
1920
1921 let frame = tokio::time::timeout(Duration::from_secs(1), stream.next())
1922 .await
1923 .expect("stream should yield")
1924 .expect("frame should exist")
1925 .expect("frame should be ok");
1926 assert_eq!(frame.sequence_number, 1);
1927 assert_eq!(frame.entries.len(), 1);
1928 assert_eq!(frame.entries[0].value.as_ref(), b"v1");
1929 }
1930
1931 #[tokio::test]
1932 async fn reduce_consumes_range_iterator_and_returns_detail() {
1933 let engine = Arc::new(FakeEngine::default());
1934 engine.set_current_sequence(7);
1935 engine.set_range_rows(vec![
1936 (Bytes::from_static(b"a"), Bytes::from_static(b"xx")),
1937 (Bytes::from_static(b"bb"), Bytes::from_static(b"yyy")),
1938 ]);
1939 let connect = QueryConnect::new(AppState::new(engine.clone()));
1940 let bytes = exoware_proto::query::ReduceRequest {
1941 start: b"a".to_vec(),
1942 end: b"z".to_vec(),
1943 params: Some(exoware_proto::query::ReduceParams {
1944 reducers: vec![exoware_proto::query::RangeReducerSpec {
1945 op: exoware_proto::query::RangeReduceOp::RANGE_REDUCE_OP_COUNT_ALL.into(),
1946 ..Default::default()
1947 }],
1948 ..Default::default()
1949 })
1950 .into(),
1951 ..Default::default()
1952 }
1953 .encode_to_vec();
1954 let request = buffa::view::OwnedView::<
1955 exoware_proto::store::query::v1::ReduceRequestView<'static>,
1956 >::decode(bytes.into())
1957 .expect("decode reduce request");
1958
1959 let response = QueryApi::reduce(&connect, Context::default(), request)
1960 .await
1961 .expect("reduce")
1962 .body;
1963 let detail = response.detail.as_option().expect("query detail").clone();
1964 let response = to_domain_reduce_response(response).expect("decode reduce response");
1965
1966 assert_eq!(engine.range_next_count(), 2);
1967 assert_eq!(response.results.len(), 1);
1968 assert_eq!(response.results[0].value, Some(KvReducedValue::UInt64(2)));
1969 assert_eq!(detail.sequence_number, 7);
1970 assert!(detail.extra.is_empty());
1971 }
1972
1973 #[tokio::test]
1974 async fn reduce_uses_eof_query_extra() {
1975 let engine = Arc::new(FakeEngine::default());
1976 engine.set_current_sequence(8);
1977 engine.set_range_rows(vec![
1978 (Bytes::from_static(b"a"), Bytes::from_static(b"xx")),
1979 (Bytes::from_static(b"bb"), Bytes::from_static(b"yyy")),
1980 ]);
1981 engine.set_range_eof_extra(numeric_query_extra("final_rows", 2.0));
1982 let connect = QueryConnect::new(AppState::new(engine));
1983 let bytes = exoware_proto::query::ReduceRequest {
1984 start: b"a".to_vec(),
1985 end: b"z".to_vec(),
1986 params: Some(exoware_proto::query::ReduceParams {
1987 reducers: vec![exoware_proto::query::RangeReducerSpec {
1988 op: exoware_proto::query::RangeReduceOp::RANGE_REDUCE_OP_COUNT_ALL.into(),
1989 ..Default::default()
1990 }],
1991 ..Default::default()
1992 })
1993 .into(),
1994 ..Default::default()
1995 }
1996 .encode_to_vec();
1997 let request = buffa::view::OwnedView::<
1998 exoware_proto::store::query::v1::ReduceRequestView<'static>,
1999 >::decode(bytes.into())
2000 .expect("decode reduce request");
2001
2002 let response = QueryApi::reduce(&connect, Context::default(), request)
2003 .await
2004 .expect("reduce")
2005 .body;
2006 let detail = response.detail.as_option().expect("query detail");
2007
2008 assert_eq!(detail.sequence_number, 8);
2009 assert_eq!(
2010 detail.extra.get("final_rows").and_then(|v| v.as_number()),
2011 Some(2.0)
2012 );
2013 }
2014
2015 #[tokio::test]
2016 async fn get_many_populates_detail_on_each_frame() {
2017 let engine = Arc::new(FakeEngine::default());
2018 engine.set_current_sequence(11);
2019 let connect = QueryConnect::new(AppState::new(engine));
2020 let bytes = exoware_proto::query::GetManyRequest {
2021 keys: vec![b"a".to_vec(), b"bb".to_vec(), b"ccc".to_vec()],
2022 batch_size: 2,
2023 ..Default::default()
2024 }
2025 .encode_to_vec();
2026 let request = buffa::view::OwnedView::<
2027 exoware_proto::store::query::v1::GetManyRequestView<'static>,
2028 >::decode(bytes.into())
2029 .expect("decode get_many request");
2030
2031 let mut stream = QueryApi::get_many(&connect, Context::default(), request)
2032 .await
2033 .expect("get_many")
2034 .body;
2035 let mut frame_sizes = Vec::new();
2036 let mut detail_frames = 0usize;
2037 while let Some(frame) = stream.next().await {
2038 let frame = frame.expect("get_many frame");
2039 frame_sizes.push(frame.results.len());
2040 let detail = frame.detail.as_option().expect("query detail");
2041 assert_eq!(detail.sequence_number, 11);
2042 assert!(detail.extra.is_empty());
2043 detail_frames += 1;
2044 }
2045
2046 assert_eq!(frame_sizes, vec![2, 1]);
2047 assert_eq!(detail_frames, 2);
2048 }
2049
2050 #[tokio::test]
2051 async fn range_returns_without_materializing_full_iterator() {
2052 let engine = Arc::new(FakeEngine::default());
2053 engine.set_current_sequence(9);
2054 engine.set_range_rows(
2055 (0..1000)
2056 .map(|i| {
2057 (
2058 Bytes::from(format!("key-{i:04}")),
2059 Bytes::from_static(b"value"),
2060 )
2061 })
2062 .collect(),
2063 );
2064 let connect = QueryConnect::new(AppState::new(engine.clone()));
2065 let bytes = exoware_proto::query::RangeRequest {
2066 start: b"a".to_vec(),
2067 end: b"z".to_vec(),
2068 limit: Some(1000),
2069 batch_size: 1,
2070 ..Default::default()
2071 }
2072 .encode_to_vec();
2073 let request = buffa::view::OwnedView::<
2074 exoware_proto::store::query::v1::RangeRequestView<'static>,
2075 >::decode(bytes.into())
2076 .expect("decode range request");
2077
2078 let mut stream = QueryApi::range(&connect, Context::default(), request)
2079 .await
2080 .expect("range")
2081 .body;
2082
2083 tokio::time::sleep(Duration::from_millis(50)).await;
2084 let consumed = engine.range_next_count();
2085 assert!(
2086 consumed < 1000,
2087 "range should not consume the full iterator before the response stream is read; consumed {consumed}",
2088 );
2089
2090 let mut rows = 0;
2091 let mut latest_detail = None;
2092 let mut detail_frames = 0usize;
2093 while let Some(frame) = stream.next().await {
2094 let frame = frame.expect("range frame");
2095 rows += frame.results.len();
2096 if let Some(detail) = frame.detail.as_option() {
2097 detail_frames += 1;
2098 latest_detail = Some(detail.clone());
2099 }
2100 }
2101
2102 assert_eq!(rows, 1000);
2103 assert_eq!(detail_frames, 1000);
2104 let detail = latest_detail.expect("query detail");
2105 assert_eq!(detail.sequence_number, 9);
2106 assert!(detail.extra.is_empty());
2107 }
2108
2109 #[tokio::test]
2110 async fn range_emits_eof_query_extra_after_rows() {
2111 let engine = Arc::new(FakeEngine::default());
2112 engine.set_current_sequence(10);
2113 engine.set_range_rows(vec![
2114 (Bytes::from_static(b"a"), Bytes::from_static(b"1")),
2115 (Bytes::from_static(b"b"), Bytes::from_static(b"2")),
2116 ]);
2117 engine.set_range_eof_extra(numeric_query_extra("final_rows", 2.0));
2118 let connect = QueryConnect::new(AppState::new(engine));
2119 let bytes = exoware_proto::query::RangeRequest {
2120 start: b"a".to_vec(),
2121 end: b"z".to_vec(),
2122 limit: Some(2),
2123 batch_size: 2,
2124 ..Default::default()
2125 }
2126 .encode_to_vec();
2127 let request = buffa::view::OwnedView::<
2128 exoware_proto::store::query::v1::RangeRequestView<'static>,
2129 >::decode(bytes.into())
2130 .expect("decode range request");
2131
2132 let mut stream = QueryApi::range(&connect, Context::default(), request)
2133 .await
2134 .expect("range")
2135 .body;
2136 let mut frames = Vec::new();
2137 while let Some(frame) = stream.next().await {
2138 frames.push(frame.expect("range frame"));
2139 }
2140
2141 assert_eq!(frames.len(), 2);
2142 assert_eq!(frames[0].results.len(), 2);
2143 let row_detail = frames[0].detail.as_option().expect("row detail");
2144 assert!(row_detail.extra.is_empty());
2145
2146 assert!(frames[1].results.is_empty());
2147 let final_detail = frames[1].detail.as_option().expect("final detail");
2148 assert_eq!(final_detail.sequence_number, 10);
2149 assert_eq!(
2150 final_detail
2151 .extra
2152 .get("final_rows")
2153 .and_then(|v| v.as_number()),
2154 Some(2.0)
2155 );
2156 }
2157
2158 #[tokio::test]
2159 async fn subscribe_without_replay_reads_the_next_live_batch() {
2160 let engine = Arc::new(FakeEngine::default());
2161 let state = AppState::new(engine.clone());
2162 let connect = StreamConnect::new(state.clone());
2163 let mut stream = subscribe_stream(&connect, None).await.expect("subscribe");
2164 engine.publish_live(state.stream.clone(), 1, vec![matching_kv(b"hit", b"v1")]);
2165 let frame = tokio::time::timeout(Duration::from_secs(1), stream.next())
2166 .await
2167 .expect("stream should yield")
2168 .expect("frame should exist")
2169 .expect("frame should be ok");
2170 assert_eq!(frame.sequence_number, 1);
2171 assert_eq!(frame.entries.len(), 1);
2172 assert_eq!(frame.entries[0].value.as_ref(), b"v1");
2173 }
2174
2175 #[tokio::test]
2176 async fn subscribe_past_end_reads_only_future_live_batches() {
2177 let engine = Arc::new(FakeEngine::default());
2178 engine.set_current_sequence(5);
2179 for seq in 1..=5 {
2180 engine.set_batch(seq, Some(vec![matching_kv(b"seed", b"v")]));
2181 }
2182 let state = AppState::new(engine.clone());
2183 let connect = StreamConnect::new(state.clone());
2184 let mut stream = subscribe_stream(&connect, Some(15))
2185 .await
2186 .expect("subscribe");
2187
2188 assert!(
2189 tokio::time::timeout(Duration::from_millis(200), stream.next())
2190 .await
2191 .is_err(),
2192 "past-end cursor should not replay synthetic or historical frames",
2193 );
2194
2195 engine.publish_live(state.stream.clone(), 6, vec![matching_kv(b"live", b"n")]);
2196 let frame = tokio::time::timeout(Duration::from_secs(1), stream.next())
2197 .await
2198 .expect("stream should yield")
2199 .expect("frame should exist")
2200 .expect("frame should be ok");
2201 assert_eq!(frame.sequence_number, 6);
2202 assert_eq!(frame.entries.len(), 1);
2203 assert_eq!(frame.entries[0].value.as_ref(), b"n");
2204 }
2205
2206 #[tokio::test]
2207 async fn replay_hole_returns_batch_evicted_error_instead_of_empty_frame() {
2208 let engine = Arc::new(FakeEngine::default());
2209 engine.set_current_sequence(3);
2210 engine.set_oldest_retained(Some(2));
2211 engine.set_batch(2, Some(vec![matching_kv(b"replay", b"v2")]));
2212
2213 let state = AppState::new(engine);
2214 let connect = StreamConnect::new(state);
2215 let mut stream = subscribe_stream(&connect, Some(2))
2216 .await
2217 .expect("subscribe");
2218
2219 let first = tokio::time::timeout(Duration::from_secs(1), stream.next())
2220 .await
2221 .expect("stream should yield")
2222 .expect("first replay frame should exist")
2223 .expect("first replay frame should be ok");
2224 assert_eq!(first.sequence_number, 2);
2225 assert_eq!(first.entries.len(), 1);
2226
2227 let err = tokio::time::timeout(Duration::from_secs(1), stream.next())
2228 .await
2229 .expect("stream should yield error")
2230 .expect("error item should exist")
2231 .expect_err("replay hole must be surfaced as an error");
2232 let decoded = decode_connect_error(&err).expect("decode connect error");
2233 assert_eq!(
2234 decoded.error_info.expect("error info").reason,
2235 crate::stream::REASON_BATCH_EVICTED,
2236 );
2237 assert!(
2238 tokio::time::timeout(Duration::from_secs(1), stream.next())
2239 .await
2240 .expect("stream should terminate")
2241 .is_none(),
2242 "stream must terminate after surfacing the replay hole",
2243 );
2244 }
2245
2246 #[tokio::test]
2247 async fn replay_with_live_burst_under_capacity_still_delivers_in_order() {
2248 const REPLAY_BATCHES: u64 = 100;
2249
2250 let engine = Arc::new(FakeEngine::default());
2251 engine.set_current_sequence(REPLAY_BATCHES);
2252 engine.set_oldest_retained(Some(1));
2253 for seq in 1..=REPLAY_BATCHES {
2254 engine.set_batch(seq, Some(vec![matching_kv(b"replay", b"v")]));
2255 }
2256
2257 let state = AppState::new(engine.clone());
2258 engine.publish_on_every_get_batch(
2259 state.stream.clone(),
2260 REPLAY_BATCHES,
2261 vec![matching_kv(b"live", b"tail")],
2262 );
2263
2264 let connect = StreamConnect::new(state);
2265 let mut stream = subscribe_stream(&connect, Some(1))
2266 .await
2267 .expect("subscribe");
2268 let mut sequence_numbers = Vec::with_capacity((REPLAY_BATCHES * 2) as usize);
2269 while sequence_numbers.len() < (REPLAY_BATCHES * 2) as usize {
2270 let frame = tokio::time::timeout(Duration::from_secs(2), stream.next())
2271 .await
2272 .expect("stream should keep yielding")
2273 .expect("frame should exist")
2274 .expect("frame should be ok");
2275 sequence_numbers.push(frame.sequence_number);
2276 }
2277
2278 let expected: Vec<u64> = (1..=(REPLAY_BATCHES * 2)).collect();
2279 assert_eq!(sequence_numbers, expected);
2280 }
2281
2282 #[tokio::test]
2283 async fn replay_large_live_burst_is_paced_by_client_reads() {
2284 const REPLAY_BATCHES: u64 = 300;
2285
2286 let engine = Arc::new(FakeEngine::default());
2287 engine.set_current_sequence(REPLAY_BATCHES);
2288 engine.set_oldest_retained(Some(1));
2289 for seq in 1..=REPLAY_BATCHES {
2290 engine.set_batch(seq, Some(vec![matching_kv(b"replay", b"v")]));
2291 }
2292
2293 let state = AppState::new(engine.clone());
2294 engine.publish_on_every_get_batch(
2295 state.stream.clone(),
2296 REPLAY_BATCHES,
2297 vec![matching_kv(b"live", b"tail")],
2298 );
2299
2300 let connect = StreamConnect::new(state);
2301 let mut stream = subscribe_stream(&connect, Some(1))
2302 .await
2303 .expect("subscribe");
2304 let mut sequence_numbers = Vec::with_capacity((REPLAY_BATCHES * 2) as usize);
2305 while sequence_numbers.len() < (REPLAY_BATCHES * 2) as usize {
2306 let frame = tokio::time::timeout(Duration::from_secs(2), stream.next())
2307 .await
2308 .expect("stream should keep yielding")
2309 .expect("frame should exist")
2310 .expect("frame should be ok");
2311 sequence_numbers.push(frame.sequence_number);
2312 }
2313 let expected: Vec<u64> = (1..=(REPLAY_BATCHES * 2)).collect();
2314 assert_eq!(sequence_numbers, expected);
2315 }
2316}