1pub mod keys;
14pub mod kv_codec;
15pub mod proto;
16pub mod prune_policy;
17pub mod selector;
18pub mod stream_filter;
19pub use keys::{Key, KeyMut, KeyValidationError, Prefix, PrefixError, Value, MAX_KEY_LEN};
20pub use proto::*;
21extern crate self as exoware_proto;
22
23use bytes::Bytes;
24use connectrpc::client::{ClientConfig, ServerStream as ConnectServerStream};
25use connectrpc::{ConnectError, ErrorCode};
26use exoware_proto::compact::ServiceClient as CompactServiceClient;
27use exoware_proto::ingest::ServiceClient as IngestServiceClient;
28use exoware_proto::log::ingest::v1::PutRequest as ProtoPutRequest;
29use exoware_proto::query as proto_query;
30use exoware_proto::query::ServiceClient as QueryServiceClient;
31use exoware_proto::store::compact::v1::PruneRequest as ProtoPruneRequest;
32use exoware_proto::store::query::v1::{
33 GetManyRequest as ProtoGetManyRequest, GetRequest as ProtoGetRequest,
34 RangeRequest as ProtoRangeRequest, ReduceRequest as ProtoWireReduceRequest,
35};
36use exoware_proto::RangeReduceRequest as DomainRangeReduceRequest;
37use exoware_proto::{
38 connect_compression_registry as proto_connect_compression_registry,
39 decode_connect_error as proto_decode_connect_error,
40 to_domain_reduce_response as proto_to_domain_reduce_response,
41 to_proto_reduce_params as proto_to_proto_reduce_params,
42 PreferZstdHttpClient as ProtoPreferZstdHttpClient,
43};
44use futures::future::BoxFuture;
45use keys::is_valid_key_size;
46use kv_codec::{KvExpr, KvFieldRef, KvReducedValue};
47use std::collections::HashMap;
48use std::sync::{
49 atomic::{AtomicU64, Ordering},
50 Arc,
51};
52use std::time::Duration;
53
54const DEFAULT_RETRY_MAX_ATTEMPTS: usize = 3;
55const DEFAULT_RETRY_INITIAL_BACKOFF_MS: u64 = 100;
56const DEFAULT_RETRY_MAX_BACKOFF_MS: u64 = 2_000;
57
58pub trait IntoStoreWriteValue {
61 fn into_store_write_value(self) -> Bytes;
62}
63
64impl IntoStoreWriteValue for Bytes {
65 fn into_store_write_value(self) -> Bytes {
66 self
67 }
68}
69
70impl IntoStoreWriteValue for &Bytes {
71 fn into_store_write_value(self) -> Bytes {
72 self.clone()
73 }
74}
75
76impl IntoStoreWriteValue for Vec<u8> {
77 fn into_store_write_value(self) -> Bytes {
78 self.into()
79 }
80}
81
82impl IntoStoreWriteValue for &Vec<u8> {
83 fn into_store_write_value(self) -> Bytes {
84 Bytes::copy_from_slice(self)
85 }
86}
87
88impl IntoStoreWriteValue for &[u8] {
89 fn into_store_write_value(self) -> Bytes {
90 Bytes::copy_from_slice(self)
91 }
92}
93
94impl<const N: usize> IntoStoreWriteValue for &[u8; N] {
95 fn into_store_write_value(self) -> Bytes {
96 Bytes::copy_from_slice(self)
97 }
98}
99
100#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
112pub enum ConnectRequestCompression {
113 #[default]
115 None,
116 Zstd,
118 Gzip,
120}
121
122impl ConnectRequestCompression {
123 fn wire_name(self) -> Option<&'static str> {
124 match self {
125 Self::None => None,
126 Self::Zstd => Some("zstd"),
127 Self::Gzip => Some("gzip"),
128 }
129 }
130}
131
132const STORE_CLIENT_MAX_MESSAGE_BYTES: usize = 256 * 1024 * 1024;
137
138fn store_connect_client_config(
145 base_uri: http::Uri,
146 request_compression: ConnectRequestCompression,
147) -> ClientConfig {
148 let config = ClientConfig::new(base_uri)
149 .with_compression(proto_connect_compression_registry())
150 .with_default_max_message_size(STORE_CLIENT_MAX_MESSAGE_BYTES);
151 match request_compression.wire_name() {
152 Some(name) => config.compress_requests(name),
153 None => config,
154 }
155}
156
157#[derive(Debug, thiserror::Error)]
159pub enum ClientError {
160 #[error("HTTP error: {0}")]
161 Http(#[from] reqwest::Error),
162 #[error("RPC error ({0})")]
163 Rpc(Box<ConnectError>),
164 #[error("store key prefix error: {0}")]
165 Prefix(#[from] StoreKeyPrefixError),
166 #[error("invalid key length: expected {expected}, got {got}")]
167 InvalidKeyLength { expected: usize, got: usize },
168 #[error("wire format error: {0}")]
169 WireFormat(String),
170}
171
172impl ClientError {
173 pub fn rpc_error(&self) -> Option<&ConnectError> {
174 match self {
175 Self::Rpc(err) => Some(err.as_ref()),
176 _ => None,
177 }
178 }
179
180 pub fn rpc_code(&self) -> Option<ErrorCode> {
181 self.rpc_error().map(|err| err.code)
182 }
183
184 pub fn decoded_rpc_error(
185 &self,
186 ) -> Result<Option<exoware_proto::DecodedConnectError>, buffa::DecodeError> {
187 self.rpc_error().map(proto_decode_connect_error).transpose()
188 }
189}
190
191#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
194pub enum StoreKeyPrefixError {
195 #[error("key does not belong to this store prefix")]
196 PrefixMismatch,
197 #[error("key offset {offset} plus store prefix shift {shift} exceeds u16")]
198 KeyOffsetOverflow { offset: u16, shift: u16 },
199 #[error("key prefix error: {0}")]
200 Prefix(#[from] PrefixError),
201}
202
203#[derive(Clone, Debug, PartialEq, Eq, Hash)]
211pub struct StoreKeyPrefix {
212 inner: Prefix,
213}
214
215impl StoreKeyPrefix {
216 pub fn new(prefix: impl Into<Bytes>) -> Result<Self, StoreKeyPrefixError> {
217 Ok(Self {
218 inner: Prefix::new(prefix)?,
219 })
220 }
221
222 pub const fn identity() -> Self {
226 Self {
227 inner: Prefix::empty(),
228 }
229 }
230
231 #[inline]
233 pub fn prefix(&self) -> &Bytes {
234 self.inner.as_bytes()
235 }
236
237 #[inline]
240 pub fn matches(&self, key: &[u8]) -> bool {
241 self.inner.matches(key)
242 }
243
244 #[inline]
246 pub fn max_logical_key_len(&self) -> usize {
247 self.inner.max_payload_len()
248 }
249
250 pub fn encode_key(&self, key: &Key) -> Result<Key, StoreKeyPrefixError> {
253 Ok(self.inner.encode_key(key)?)
254 }
255
256 pub fn decode_key(&self, key: &Key) -> Result<Key, StoreKeyPrefixError> {
259 self.inner
260 .strip(key)
261 .map_err(|_| StoreKeyPrefixError::PrefixMismatch)
262 }
263
264 pub fn encode_range(&self, start: &Key, end: &Key) -> Result<(Key, Key), StoreKeyPrefixError> {
272 let start = self.encode_key(start)?;
273 let end = if end.is_empty() {
274 self.inner.bounds().1
275 } else {
276 let max_len = self.max_logical_key_len();
277 if end.len() > max_len {
278 self.encode_key(&end.slice(..max_len))?
279 } else {
280 self.encode_key(end)?
281 }
282 };
283 Ok((start, end))
284 }
285
286 fn prefix_selector(
287 &self,
288 selector: &crate::selector::Selector,
289 ) -> Result<crate::selector::Selector, StoreKeyPrefixError> {
290 let prefix = self.inner.join(&Prefix::new(selector.prefix.clone())?)?;
291 Ok(crate::selector::Selector {
292 prefix: prefix.as_bytes().clone(),
293 payload_regex: selector.payload_regex.clone(),
294 })
295 }
296}
297
298#[derive(Clone, Debug)]
302pub struct PrefixedStoreClient {
303 client: StoreClient,
304 prefix: StoreKeyPrefix,
305}
306
307impl PrefixedStoreClient {
308 pub fn new(client: StoreClient, prefix: StoreKeyPrefix) -> Self {
310 Self { client, prefix }
311 }
312
313 pub fn empty(client: StoreClient) -> Self {
316 Self::new(client, StoreKeyPrefix::identity())
317 }
318
319 pub fn key_prefix(&self) -> &StoreKeyPrefix {
321 &self.prefix
322 }
323
324 pub fn client(&self) -> &StoreClient {
327 &self.client
328 }
329
330 pub fn encode_store_key(&self, key: &Key) -> Result<Key, ClientError> {
334 Ok(self.prefix.encode_key(key)?)
335 }
336
337 pub fn decode_store_key(&self, key: &Key) -> Result<Key, ClientError> {
339 Ok(self.prefix.decode_key(key)?)
340 }
341
342 fn encode_store_range(&self, start: &Key, end: &Key) -> Result<(Key, Key), ClientError> {
343 Ok(self.prefix.encode_range(start, end)?)
344 }
345
346 fn prefix_prune_policies(
347 &self,
348 policies: &[crate::prune_policy::PrunePolicy],
349 ) -> Result<Vec<crate::prune_policy::PrunePolicy>, ClientError> {
350 policies
351 .iter()
352 .map(|policy| {
353 use crate::prune_policy::{PolicyScope, PrunePolicy};
354 let scope = match &policy.scope {
355 PolicyScope::Keys(scope) => {
356 let mut scope = scope.clone();
357 scope.selector = self.prefix.prefix_selector(&scope.selector)?;
358 PolicyScope::Keys(scope)
359 }
360 PolicyScope::Sequence => PolicyScope::Sequence,
361 };
362 Ok(PrunePolicy {
363 scope,
364 retain: policy.retain.clone(),
365 })
366 })
367 .collect::<Result<Vec<_>, StoreKeyPrefixError>>()
368 .map_err(ClientError::from)
369 }
370
371 fn prefix_stream_filter(
372 &self,
373 filter: crate::stream_filter::StreamFilter,
374 ) -> Result<crate::stream_filter::StreamFilter, ClientError> {
375 let selectors = filter
376 .selectors
377 .iter()
378 .map(|mk| self.prefix.prefix_selector(mk))
379 .collect::<Result<Vec<_>, _>>()?;
380 Ok(crate::stream_filter::StreamFilter {
381 selectors,
382 value_filters: filter.value_filters,
383 })
384 }
385
386 fn prefix_reduce_request(
387 &self,
388 request: &DomainRangeReduceRequest,
389 ) -> Result<DomainRangeReduceRequest, ClientError> {
390 let mut request = request.clone();
391 shift_reduce_request_key_offsets(self.prefix.prefix().len(), &mut request)?;
392 Ok(request)
393 }
394
395 pub fn ingest(&self) -> Ingest<'_> {
399 Ingest { c: self }
400 }
401
402 pub fn query(&self) -> Query<'_> {
404 Query { c: self }
405 }
406
407 pub fn compact(&self) -> Compact<'_> {
409 Compact { c: self }
410 }
411
412 pub fn stream(&self) -> Stream<'_> {
414 Stream { c: self }
415 }
416
417 pub fn create_session(&self) -> SerializableReadSession {
419 self.create_session_with_sequence(0)
420 }
421
422 pub fn create_session_with_sequence(&self, sequence: u64) -> SerializableReadSession {
424 SerializableReadSession {
425 client: self.clone(),
426 state: Arc::new(SessionState {
427 sequence: Arc::new(AtomicU64::new(sequence)),
428 init_gate: tokio::sync::Mutex::new(()),
429 }),
430 }
431 }
432
433 pub(crate) async fn put(&self, kvs: &[(&Key, &[u8])]) -> Result<u64, ClientError> {
436 let keys = kvs
437 .iter()
438 .map(|(key, _)| self.encode_store_key(key))
439 .collect::<Result<Vec<_>, _>>()?;
440 let prefixed: Vec<(&Key, &[u8])> = keys
441 .iter()
442 .zip(kvs.iter())
443 .map(|(key, (_, value))| (key, *value))
444 .collect();
445 self.client.put_physical(&prefixed).await
446 }
447
448 pub(crate) async fn send_get(
451 &self,
452 key: &Key,
453 min_sequence_number: Option<u64>,
454 ) -> Result<
455 (
456 exoware_proto::query::GetResponse,
457 Option<proto_query::Detail>,
458 ),
459 ClientError,
460 > {
461 self.client
462 .send_get(&self.encode_store_key(key)?, min_sequence_number)
463 .await
464 }
465
466 pub(crate) async fn get(&self, key: &Key) -> Result<Option<Bytes>, ClientError> {
467 self.client.get(&self.encode_store_key(key)?).await
468 }
469
470 pub(crate) async fn get_with_min_sequence_number(
471 &self,
472 key: &Key,
473 min_sequence_number: u64,
474 ) -> Result<Option<Bytes>, ClientError> {
475 self.client
476 .get_with_min_sequence_number(&self.encode_store_key(key)?, min_sequence_number)
477 .await
478 }
479
480 pub(crate) async fn get_many(
481 &self,
482 keys: &[&Key],
483 batch_size: u32,
484 ) -> Result<GetManyStream, ClientError> {
485 self.get_many_internal(keys, batch_size, None, None).await
486 }
487
488 pub(crate) async fn get_many_with_min_sequence_number(
489 &self,
490 keys: &[&Key],
491 batch_size: u32,
492 min_sequence_number: u64,
493 ) -> Result<GetManyStream, ClientError> {
494 self.get_many_internal(keys, batch_size, Some(min_sequence_number), None)
495 .await
496 }
497
498 pub(crate) async fn get_many_internal(
499 &self,
500 keys: &[&Key],
501 batch_size: u32,
502 min_sequence_number: Option<u64>,
503 observed_sequence: Option<Arc<AtomicU64>>,
504 ) -> Result<GetManyStream, ClientError> {
505 let mut proto_keys: Vec<Vec<u8>> = Vec::with_capacity(keys.len());
506 for key in keys {
507 let encoded = self.encode_store_key(key)?;
508 if !is_valid_key_size(encoded.len()) {
509 return Err(ClientError::WireFormat(format!(
510 "key length {} is outside valid store key range ({}..={})",
511 encoded.len(),
512 keys::MIN_KEY_LEN,
513 MAX_KEY_LEN
514 )));
515 }
516 proto_keys.push(encoded.to_vec());
517 }
518 let mut stream = self
519 .client
520 .get_many_internal(
521 proto_keys,
522 batch_size,
523 min_sequence_number,
524 observed_sequence,
525 )
526 .await?;
527 stream.key_prefix = Some(self.prefix.clone());
528 Ok(stream)
529 }
530
531 pub(crate) async fn range(
532 &self,
533 start: &Key,
534 end: &Key,
535 limit: usize,
536 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
537 self.range_internal(start, end, limit, RangeMode::Forward, None)
538 .await
539 }
540
541 pub(crate) async fn range_with_mode(
542 &self,
543 start: &Key,
544 end: &Key,
545 limit: usize,
546 mode: RangeMode,
547 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
548 self.range_internal(start, end, limit, mode, None).await
549 }
550
551 pub(crate) async fn range_with_min_sequence_number(
552 &self,
553 start: &Key,
554 end: &Key,
555 limit: usize,
556 min_sequence_number: u64,
557 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
558 self.range_internal(
559 start,
560 end,
561 limit,
562 RangeMode::Forward,
563 Some(min_sequence_number),
564 )
565 .await
566 }
567
568 pub(crate) async fn range_with_mode_and_min_sequence_number(
569 &self,
570 start: &Key,
571 end: &Key,
572 limit: usize,
573 mode: RangeMode,
574 min_sequence_number: u64,
575 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
576 self.range_internal(start, end, limit, mode, Some(min_sequence_number))
577 .await
578 }
579
580 async fn range_internal(
581 &self,
582 start: &Key,
583 end: &Key,
584 limit: usize,
585 mode: RangeMode,
586 min_sequence_number: Option<u64>,
587 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
588 self.range_stream_internal(
589 start,
590 end,
591 limit,
592 limit.max(1),
593 mode,
594 RangeStreamReadOptions {
595 min_sequence_number,
596 observed_sequence: None,
597 },
598 )
599 .await?
600 .collect()
601 .await
602 }
603
604 pub(crate) async fn range_stream(
605 &self,
606 start: &Key,
607 end: &Key,
608 limit: usize,
609 batch_size: usize,
610 ) -> Result<RangeStream, ClientError> {
611 self.range_stream_internal(
612 start,
613 end,
614 limit,
615 batch_size,
616 RangeMode::Forward,
617 RangeStreamReadOptions::default(),
618 )
619 .await
620 }
621
622 pub(crate) async fn range_stream_with_mode(
623 &self,
624 start: &Key,
625 end: &Key,
626 limit: usize,
627 batch_size: usize,
628 mode: RangeMode,
629 ) -> Result<RangeStream, ClientError> {
630 self.range_stream_internal(start, end, limit, batch_size, mode, Default::default())
631 .await
632 }
633
634 pub(crate) async fn range_stream_with_min_sequence_number(
635 &self,
636 start: &Key,
637 end: &Key,
638 limit: usize,
639 batch_size: usize,
640 min_sequence_number: u64,
641 ) -> Result<RangeStream, ClientError> {
642 self.range_stream_internal(
643 start,
644 end,
645 limit,
646 batch_size,
647 RangeMode::Forward,
648 RangeStreamReadOptions {
649 min_sequence_number: Some(min_sequence_number),
650 observed_sequence: None,
651 },
652 )
653 .await
654 }
655
656 pub(crate) async fn range_stream_with_mode_and_min_sequence_number(
657 &self,
658 start: &Key,
659 end: &Key,
660 limit: usize,
661 batch_size: usize,
662 mode: RangeMode,
663 min_sequence_number: u64,
664 ) -> Result<RangeStream, ClientError> {
665 self.range_stream_internal(
666 start,
667 end,
668 limit,
669 batch_size,
670 mode,
671 RangeStreamReadOptions {
672 min_sequence_number: Some(min_sequence_number),
673 observed_sequence: None,
674 },
675 )
676 .await
677 }
678
679 pub(crate) async fn range_stream_internal(
680 &self,
681 start: &Key,
682 end: &Key,
683 limit: usize,
684 batch_size: usize,
685 mode: RangeMode,
686 options: RangeStreamReadOptions,
687 ) -> Result<RangeStream, ClientError> {
688 let (start, end) = self.encode_store_range(start, end)?;
689 let mut stream = self
690 .client
691 .range_stream_internal(&start, &end, limit, batch_size, mode, options)
692 .await?;
693 stream.key_prefix = Some(self.prefix.clone());
694 Ok(stream)
695 }
696
697 pub(crate) async fn range_reduce(
698 &self,
699 start: &Key,
700 end: &Key,
701 request: &DomainRangeReduceRequest,
702 ) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
703 let (response, _) = self
704 .range_reduce_response_internal(start, end, request, None)
705 .await?;
706 scalar_reduce_results(response)
707 }
708
709 pub(crate) async fn range_reduce_with_min_sequence_number(
710 &self,
711 start: &Key,
712 end: &Key,
713 request: &DomainRangeReduceRequest,
714 min_sequence_number: u64,
715 ) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
716 let (response, _) = self
717 .range_reduce_response_internal(start, end, request, Some(min_sequence_number))
718 .await?;
719 scalar_reduce_results(response)
720 }
721
722 pub(crate) async fn range_reduce_response(
723 &self,
724 start: &Key,
725 end: &Key,
726 request: &DomainRangeReduceRequest,
727 ) -> Result<exoware_proto::query::ReduceResponse, ClientError> {
728 let (body, _) = self
729 .range_reduce_response_internal(start, end, request, None)
730 .await?;
731 Ok(body)
732 }
733
734 pub(crate) async fn range_reduce_response_with_min_sequence_number(
735 &self,
736 start: &Key,
737 end: &Key,
738 request: &DomainRangeReduceRequest,
739 min_sequence_number: u64,
740 ) -> Result<exoware_proto::query::ReduceResponse, ClientError> {
741 let (body, _) = self
742 .range_reduce_response_internal(start, end, request, Some(min_sequence_number))
743 .await?;
744 Ok(body)
745 }
746
747 pub(crate) async fn range_reduce_response_internal(
748 &self,
749 start: &Key,
750 end: &Key,
751 request: &DomainRangeReduceRequest,
752 min_sequence_number: Option<u64>,
753 ) -> Result<
754 (
755 exoware_proto::query::ReduceResponse,
756 Option<proto_query::Detail>,
757 ),
758 ClientError,
759 > {
760 let (start, end) = self.encode_store_range(start, end)?;
761 let request = self.prefix_reduce_request(request)?;
762 self.client
763 .range_reduce_response_internal(&start, &end, &request, min_sequence_number)
764 .await
765 }
766
767 pub(crate) async fn prune(
768 &self,
769 policies: &[crate::prune_policy::PrunePolicy],
770 ) -> Result<(), ClientError> {
771 let policies = self.prefix_prune_policies(policies)?;
772 self.client.prune(&policies).await
773 }
774
775 pub(crate) async fn subscribe(
776 &self,
777 filter: crate::stream_filter::StreamFilter,
778 since_sequence_number: Option<u64>,
779 ) -> Result<StreamSubscription, ClientError> {
780 let filter = self.prefix_stream_filter(filter)?;
784 let mut sub = self
785 .client
786 .subscribe_physical(filter, since_sequence_number)
787 .await?;
788 sub.key_prefix = Some(self.prefix.clone());
789 Ok(sub)
790 }
791
792 pub(crate) async fn stream_get(
793 &self,
794 sequence_number: u64,
795 ) -> Result<Option<Vec<(Key, Bytes)>>, ClientError> {
796 let Some(owned) = self.client.stream_get_physical(sequence_number).await? else {
797 return Ok(None);
798 };
799 let mut out = Vec::with_capacity(owned.entries.len());
800 for entry in owned.entries {
801 let key = Bytes::from(entry.key);
802 if !self.prefix.matches(&key) {
803 continue;
804 }
805 out.push((self.decode_store_key(&key)?, entry.value));
806 }
807 Ok(Some(out))
808 }
809}
810
811fn scalar_reduce_results(
813 response: exoware_proto::query::ReduceResponse,
814) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
815 let decoded = proto_to_domain_reduce_response(response).map_err(ClientError::WireFormat)?;
816 if !decoded.groups.is_empty() {
817 return Err(ClientError::WireFormat(
818 "grouped range reduction response returned for scalar request".to_string(),
819 ));
820 }
821 Ok(decoded
822 .results
823 .iter()
824 .map(|result| result.value.clone())
825 .collect())
826}
827
828#[derive(Clone, Debug, Default)]
834pub struct StoreWriteBatch {
835 entries: Vec<(Key, Bytes)>,
836}
837
838impl StoreWriteBatch {
839 pub fn new() -> Self {
840 Self::default()
841 }
842
843 pub fn len(&self) -> usize {
844 self.entries.len()
845 }
846
847 pub fn is_empty(&self) -> bool {
848 self.entries.is_empty()
849 }
850
851 pub fn clear(&mut self) {
852 self.entries.clear();
853 }
854
855 pub fn reserve(&mut self, additional: usize) {
856 self.entries.reserve(additional);
857 }
858
859 pub fn push(
865 &mut self,
866 client: &PrefixedStoreClient,
867 key: &Key,
868 value: impl IntoStoreWriteValue,
869 ) -> Result<&mut Self, ClientError> {
870 self.entries.push((
871 client.encode_store_key(key)?,
872 value.into_store_write_value(),
873 ));
874 Ok(self)
875 }
876
877 pub fn entries(&self) -> &[(Key, Bytes)] {
879 &self.entries
880 }
881
882 pub async fn commit(&self, client: &StoreClient) -> Result<u64, ClientError> {
883 client.put_prepared_physical(&self.entries).await
884 }
885}
886
887pub trait StoreBatchUpload {
899 type Prepared: Send;
900 type Receipt: Send;
901 type Error: std::fmt::Display + Send;
902
903 fn store_client(&self) -> &PrefixedStoreClient;
904
905 fn stage_upload(
906 &self,
907 prepared: &mut Self::Prepared,
908 batch: &mut StoreWriteBatch,
909 ) -> Result<(), Self::Error>;
910
911 fn commit_error(&self, error: ClientError) -> Self::Error;
912
913 fn mark_upload_persisted<'a>(
914 &'a self,
915 prepared: Self::Prepared,
916 sequence_number: u64,
917 ) -> BoxFuture<'a, Self::Receipt>
918 where
919 Self: Sync + 'a,
920 Self::Prepared: 'a;
921
922 fn mark_upload_failed<'a>(
923 &'a self,
924 prepared: Self::Prepared,
925 error: String,
926 ) -> BoxFuture<'a, ()>
927 where
928 Self: Sync + 'a,
929 Self::Prepared: 'a;
930
931 fn commit_upload<'a>(
932 &'a self,
933 prepared: Self::Prepared,
934 ) -> BoxFuture<'a, Result<Self::Receipt, Self::Error>>
935 where
936 Self: Sync + Sized + 'a,
937 Self::Prepared: 'a,
938 Self::Receipt: 'a,
939 Self::Error: 'a,
940 {
941 Box::pin(async move {
942 let mut prepared = prepared;
943 let mut batch = StoreWriteBatch::new();
944 if let Err(err) = self.stage_upload(&mut prepared, &mut batch) {
945 let message = err.to_string();
946 self.mark_upload_failed(prepared, message).await;
947 return Err(err);
948 }
949 match batch.commit(self.store_client().client()).await {
950 Ok(sequence_number) => {
951 Ok(self.mark_upload_persisted(prepared, sequence_number).await)
952 }
953 Err(err) => {
954 let message = err.to_string();
955 self.mark_upload_failed(prepared, message).await;
956 Err(self.commit_error(err))
957 }
958 }
959 })
960 }
961}
962
963pub trait StoreBatchPublication {
971 type PreparedPublication: Send;
972 type PublicationReceipt: Send;
973 type Error: std::fmt::Display + Send;
974
975 fn store_client(&self) -> &PrefixedStoreClient;
978
979 fn stage_publication(
980 &self,
981 prepared: &Self::PreparedPublication,
982 batch: &mut StoreWriteBatch,
983 ) -> Result<(), Self::Error>;
984
985 fn publication_commit_error(&self, error: ClientError) -> Self::Error;
986
987 fn mark_publication_persisted<'a>(
988 &'a self,
989 prepared: Self::PreparedPublication,
990 sequence_number: u64,
991 ) -> BoxFuture<'a, Self::PublicationReceipt>
992 where
993 Self: Sync + 'a,
994 Self::PreparedPublication: 'a;
995
996 fn mark_publication_failed<'a>(
997 &'a self,
998 _prepared: Self::PreparedPublication,
999 _error: String,
1000 ) -> BoxFuture<'a, ()>
1001 where
1002 Self: Sync + 'a,
1003 Self::PreparedPublication: 'a,
1004 {
1005 Box::pin(async {})
1006 }
1007
1008 fn commit_publication<'a>(
1009 &'a self,
1010 prepared: Self::PreparedPublication,
1011 ) -> BoxFuture<'a, Result<Self::PublicationReceipt, Self::Error>>
1012 where
1013 Self: Sync + Sized + 'a,
1014 Self::PreparedPublication: 'a,
1015 Self::PublicationReceipt: 'a,
1016 Self::Error: 'a,
1017 {
1018 Box::pin(async move {
1019 let mut batch = StoreWriteBatch::new();
1020 if let Err(err) = self.stage_publication(&prepared, &mut batch) {
1021 let message = err.to_string();
1022 self.mark_publication_failed(prepared, message).await;
1023 return Err(err);
1024 }
1025 match batch.commit(self.store_client().client()).await {
1026 Ok(sequence_number) => Ok(self
1027 .mark_publication_persisted(prepared, sequence_number)
1028 .await),
1029 Err(err) => {
1030 let message = err.to_string();
1031 self.mark_publication_failed(prepared, message).await;
1032 Err(self.publication_commit_error(err))
1033 }
1034 }
1035 })
1036 }
1037}
1038
1039pub trait StorePublicationFrontierWriter: StoreBatchPublication {
1048 fn latest_publication_receipt<'a>(&'a self) -> BoxFuture<'a, Option<Self::PublicationReceipt>>
1049 where
1050 Self: Sync + 'a,
1051 Self::PublicationReceipt: 'a;
1052
1053 fn prepare_publication<'a>(
1054 &'a self,
1055 ) -> BoxFuture<'a, Result<Option<Self::PreparedPublication>, Self::Error>>
1056 where
1057 Self: Sync + 'a,
1058 Self::PreparedPublication: 'a,
1059 Self::Error: 'a;
1060
1061 fn flush_publication_with_receipt<'a>(
1062 &'a self,
1063 ) -> BoxFuture<'a, Result<Option<Self::PublicationReceipt>, Self::Error>>
1064 where
1065 Self: Sync + 'a,
1066 Self::PublicationReceipt: 'a,
1067 Self::Error: 'a;
1068
1069 fn flush_publication<'a>(&'a self) -> BoxFuture<'a, Result<(), Self::Error>>
1070 where
1071 Self: Sync + 'a,
1072 Self::PublicationReceipt: 'a,
1073 Self::Error: 'a,
1074 {
1075 Box::pin(async move { self.flush_publication_with_receipt().await.map(|_| ()) })
1076 }
1077}
1078
1079#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1081pub enum RangeMode {
1082 Forward,
1083 Reverse,
1084}
1085
1086#[derive(Clone, Debug)]
1087pub struct RangeChunk {
1088 pub rows: Vec<(Key, Bytes)>,
1090 pub detail: Option<proto_query::Detail>,
1092}
1093
1094#[derive(Clone, Debug)]
1095pub struct GetManyChunk {
1096 pub entries: Vec<(Key, Option<Bytes>)>,
1098 pub detail: Option<proto_query::Detail>,
1100}
1101
1102pub struct RangeStream {
1104 stream:
1105 ConnectServerStream<hyper::body::Incoming, exoware_proto::query::RangeFrameView<'static>>,
1106 pending_frame: Option<exoware_proto::query::RangeFrame>,
1107 rows_seen: usize,
1108 final_count: Option<usize>,
1109 finished: bool,
1110 observed_sequence: Option<Arc<AtomicU64>>,
1111 key_prefix: Option<StoreKeyPrefix>,
1112}
1113
1114impl RangeStream {
1115 fn from_connect_stream(
1116 stream: ConnectServerStream<
1117 hyper::body::Incoming,
1118 exoware_proto::query::RangeFrameView<'static>,
1119 >,
1120 observed_sequence: Option<Arc<AtomicU64>>,
1121 key_prefix: Option<StoreKeyPrefix>,
1122 ) -> Self {
1123 Self {
1124 stream,
1125 pending_frame: None,
1126 rows_seen: 0,
1127 final_count: None,
1128 finished: false,
1129 observed_sequence,
1130 key_prefix,
1131 }
1132 }
1133
1134 pub fn final_count(&self) -> Option<usize> {
1135 self.final_count
1136 }
1137
1138 async fn prefetch_first_frame(&mut self) -> Result<(), ConnectError> {
1139 if self.pending_frame.is_some() || self.finished {
1140 return Ok(());
1141 }
1142 match self.stream.message().await? {
1143 Some(frame) => {
1144 self.pending_frame = Some(frame.to_owned_message());
1145 Ok(())
1146 }
1147 None => {
1148 self.finished = true;
1149 if let Some(err) = self.stream.error() {
1150 Err(err.clone())
1151 } else {
1152 self.final_count = Some(self.rows_seen);
1153 Ok(())
1154 }
1155 }
1156 }
1157 }
1158
1159 pub async fn next_chunk(&mut self) -> Result<Option<RangeChunk>, ClientError> {
1160 loop {
1161 if self.finished {
1162 return Ok(None);
1163 }
1164
1165 let frame = if let Some(frame) = self.pending_frame.take() {
1166 frame
1167 } else {
1168 let Some(frame) = self
1169 .stream
1170 .message()
1171 .await
1172 .map_err(client_error_from_connect)?
1173 else {
1174 self.finished = true;
1175 if let Some(err) = self.stream.error() {
1176 return Err(client_error_from_connect(err.clone()));
1177 }
1178 self.final_count = Some(self.rows_seen);
1179 return Ok(None);
1180 };
1181 frame.to_owned_message()
1182 };
1183
1184 let detail = frame.detail.as_option().cloned();
1185 if let (Some(sequence_store), Some(detail)) = (&self.observed_sequence, detail.as_ref())
1186 {
1187 sequence_store.fetch_max(detail.sequence_number, Ordering::SeqCst);
1188 }
1189 let n = frame.results.len();
1190
1191 if n == 0 && detail.is_none() {
1193 continue;
1194 }
1195
1196 let mut out = Vec::with_capacity(n);
1197 for entry in frame.results {
1198 let key = Bytes::from(entry.key);
1199 let key = match &self.key_prefix {
1200 Some(prefix) => prefix.decode_key(&key)?,
1201 None => key,
1202 };
1203 out.push((key, entry.value));
1204 }
1205 self.rows_seen += n;
1206 return Ok(Some(RangeChunk { rows: out, detail }));
1207 }
1208 }
1209
1210 pub async fn collect(mut self) -> Result<Vec<(Key, Bytes)>, ClientError> {
1211 let mut entries = Vec::new();
1212 while let Some(chunk) = self.next_chunk().await? {
1213 entries.extend(chunk.rows);
1214 }
1215 Ok(entries)
1216 }
1217}
1218
1219pub struct GetManyStream {
1220 stream:
1221 ConnectServerStream<hyper::body::Incoming, exoware_proto::query::GetManyFrameView<'static>>,
1222 pending_frame: Option<exoware_proto::query::GetManyFrame>,
1223 finished: bool,
1224 observed_sequence: Option<Arc<AtomicU64>>,
1225 key_prefix: Option<StoreKeyPrefix>,
1226}
1227
1228impl GetManyStream {
1229 fn from_connect_stream(
1230 stream: ConnectServerStream<
1231 hyper::body::Incoming,
1232 exoware_proto::query::GetManyFrameView<'static>,
1233 >,
1234 observed_sequence: Option<Arc<AtomicU64>>,
1235 key_prefix: Option<StoreKeyPrefix>,
1236 ) -> Self {
1237 Self {
1238 stream,
1239 pending_frame: None,
1240 finished: false,
1241 observed_sequence,
1242 key_prefix,
1243 }
1244 }
1245
1246 async fn prefetch_first_frame(&mut self) -> Result<(), ConnectError> {
1247 if self.pending_frame.is_some() || self.finished {
1248 return Ok(());
1249 }
1250 match self.stream.message().await? {
1251 Some(frame) => {
1252 self.pending_frame = Some(frame.to_owned_message());
1253 Ok(())
1254 }
1255 None => {
1256 self.finished = true;
1257 if let Some(err) = self.stream.error() {
1258 Err(err.clone())
1259 } else {
1260 Ok(())
1261 }
1262 }
1263 }
1264 }
1265
1266 pub async fn next_chunk(&mut self) -> Result<Option<GetManyChunk>, ClientError> {
1267 loop {
1268 if self.finished {
1269 return Ok(None);
1270 }
1271 let frame = if let Some(frame) = self.pending_frame.take() {
1272 frame
1273 } else {
1274 let Some(frame) = self
1275 .stream
1276 .message()
1277 .await
1278 .map_err(client_error_from_connect)?
1279 else {
1280 self.finished = true;
1281 if let Some(err) = self.stream.error() {
1282 return Err(client_error_from_connect(err.clone()));
1283 }
1284 return Ok(None);
1285 };
1286 frame.to_owned_message()
1287 };
1288
1289 let detail = frame.detail.as_option().cloned();
1290 if let (Some(sequence_store), Some(detail)) = (&self.observed_sequence, detail.as_ref())
1291 {
1292 sequence_store.fetch_max(detail.sequence_number, Ordering::SeqCst);
1293 }
1294 let n = frame.results.len();
1295
1296 if n == 0 && detail.is_none() {
1298 continue;
1299 }
1300
1301 let mut out = Vec::with_capacity(n);
1302 for entry in frame.results {
1303 let key = Bytes::from(entry.key);
1304 let key = match &self.key_prefix {
1305 Some(prefix) => prefix.decode_key(&key)?,
1306 None => key,
1307 };
1308 out.push((key, entry.value));
1309 }
1310 return Ok(Some(GetManyChunk {
1311 entries: out,
1312 detail,
1313 }));
1314 }
1315 }
1316
1317 pub async fn collect(mut self) -> Result<HashMap<Key, Bytes>, ClientError> {
1318 let mut map = HashMap::new();
1319 while let Some(chunk) = self.next_chunk().await? {
1320 for (key, value) in chunk.entries {
1321 if let Some(v) = value {
1322 map.insert(key, v);
1323 }
1324 }
1325 }
1326 Ok(map)
1327 }
1328}
1329
1330impl RangeMode {
1331 fn to_proto(self) -> proto_query::TraversalMode {
1332 match self {
1333 Self::Forward => proto_query::TraversalMode::TRAVERSAL_MODE_FORWARD,
1334 Self::Reverse => proto_query::TraversalMode::TRAVERSAL_MODE_REVERSE,
1335 }
1336 }
1337}
1338
1339#[derive(Clone, Debug)]
1343pub struct StreamSubscriptionEntry {
1344 pub key: Key,
1345 pub value: Bytes,
1346}
1347
1348#[derive(Clone, Debug)]
1350pub struct StreamSubscriptionFrame {
1351 pub sequence_number: u64,
1352 pub entries: Vec<StreamSubscriptionEntry>,
1353}
1354
1355pub struct StreamSubscription {
1358 stream: ConnectServerStream<
1359 hyper::body::Incoming,
1360 exoware_proto::log::stream::v1::SubscribeResponseView<'static>,
1361 >,
1362 key_prefix: Option<StoreKeyPrefix>,
1363}
1364
1365impl std::fmt::Debug for StreamSubscription {
1366 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1367 f.debug_struct("StreamSubscription").finish_non_exhaustive()
1368 }
1369}
1370
1371impl StreamSubscription {
1372 pub async fn next(&mut self) -> Result<Option<StreamSubscriptionFrame>, ClientError> {
1374 loop {
1375 match self
1376 .stream
1377 .message()
1378 .await
1379 .map_err(client_error_from_connect)?
1380 {
1381 Some(view) => {
1382 let owned = view.to_owned_message();
1383 let mut entries = Vec::with_capacity(owned.entries.len());
1384 for entry in owned.entries {
1385 let key = Bytes::from(entry.key);
1386 let key = match &self.key_prefix {
1387 Some(prefix) => prefix.decode_key(&key)?,
1388 None => key,
1389 };
1390 entries.push(StreamSubscriptionEntry {
1391 key,
1392 value: entry.value,
1393 });
1394 }
1395 if entries.is_empty() {
1396 continue;
1397 }
1398 let frame = StreamSubscriptionFrame {
1399 sequence_number: owned.sequence_number,
1400 entries,
1401 };
1402 return Ok(Some(frame));
1403 }
1404 None => {
1405 if let Some(err) = self.stream.error() {
1406 return Err(client_error_from_connect(err.clone()));
1407 } else {
1408 return Ok(None);
1409 }
1410 }
1411 }
1412 }
1413 }
1414}
1415
1416fn is_batch_missing_error(err: &ConnectError) -> bool {
1419 match proto_decode_connect_error(err) {
1420 Ok(decoded) => decoded.error_info.is_some_and(|info| {
1421 info.domain == "log.stream"
1422 && matches!(info.reason.as_str(), "BATCH_EVICTED" | "BATCH_NOT_FOUND")
1423 }),
1424 Err(_) => false,
1425 }
1426}
1427
1428#[derive(Clone, Copy, Debug)]
1430pub struct RetryConfig {
1431 max_attempts: usize,
1432 initial_backoff: Duration,
1433 max_backoff: Duration,
1434}
1435
1436impl RetryConfig {
1437 pub fn standard() -> Self {
1438 Self {
1439 max_attempts: DEFAULT_RETRY_MAX_ATTEMPTS,
1440 initial_backoff: Duration::from_millis(DEFAULT_RETRY_INITIAL_BACKOFF_MS),
1441 max_backoff: Duration::from_millis(DEFAULT_RETRY_MAX_BACKOFF_MS),
1442 }
1443 }
1444
1445 pub fn disabled() -> Self {
1446 Self::standard().with_max_attempts(1)
1447 }
1448
1449 pub fn with_max_attempts(mut self, max_attempts: usize) -> Self {
1450 self.max_attempts = max_attempts.max(1);
1451 self
1452 }
1453
1454 pub fn with_initial_backoff(mut self, initial_backoff: Duration) -> Self {
1455 self.initial_backoff = initial_backoff;
1456 self
1457 }
1458
1459 pub fn with_max_backoff(mut self, max_backoff: Duration) -> Self {
1460 self.max_backoff = max_backoff;
1461 self
1462 }
1463
1464 pub(crate) fn sanitized(self) -> Self {
1465 let max_attempts = self.max_attempts.max(1);
1466 let max_backoff = self.max_backoff.max(self.initial_backoff);
1467 Self {
1468 max_attempts,
1469 initial_backoff: self.initial_backoff,
1470 max_backoff,
1471 }
1472 }
1473}
1474
1475impl Default for RetryConfig {
1476 fn default() -> Self {
1477 Self::standard()
1478 }
1479}
1480
1481fn trim_connect_base(url: &str) -> String {
1482 url.trim_end_matches('/').to_string()
1483}
1484
1485fn new_http_client() -> reqwest::Client {
1486 reqwest::Client::builder()
1487 .pool_max_idle_per_host(32)
1488 .timeout(Duration::from_secs(30))
1489 .build()
1490 .expect("failed to build HTTP client")
1491}
1492
1493#[derive(Debug, thiserror::Error)]
1495pub enum ClientBuildError {
1496 #[error("StoreClientBuilder: missing health URL (set health_url or url)")]
1497 MissingHealthUrl,
1498 #[error("StoreClientBuilder: missing ingest URL (set ingest_url or url)")]
1499 MissingIngestUrl,
1500 #[error("StoreClientBuilder: missing query URL (set query_url or url)")]
1501 MissingQueryUrl,
1502 #[error("StoreClientBuilder: missing compact URL (set compact_url or url)")]
1503 MissingCompactUrl,
1504 #[error("StoreClientBuilder: missing stream URL (set stream_url or url)")]
1505 MissingStreamUrl,
1506 #[error("StoreClientBuilder: invalid URL \"{url}\": {source}")]
1507 InvalidUrl {
1508 url: String,
1509 source: http::uri::InvalidUri,
1510 },
1511}
1512
1513#[derive(Debug, Default)]
1518pub struct StoreClientBuilder {
1519 health_url: Option<String>,
1520 ingest_url: Option<String>,
1521 query_url: Option<String>,
1522 compact_url: Option<String>,
1523 stream_url: Option<String>,
1524 retry_config: RetryConfig,
1525 connect_request_compression: ConnectRequestCompression,
1526}
1527
1528impl StoreClientBuilder {
1529 pub fn url(mut self, url: &str) -> Self {
1531 let u = trim_connect_base(url);
1532 self.health_url = Some(u.clone());
1533 self.ingest_url = Some(u.clone());
1534 self.query_url = Some(u.clone());
1535 self.compact_url = Some(u.clone());
1536 self.stream_url = Some(u);
1537 self
1538 }
1539
1540 pub fn health_url(mut self, url: &str) -> Self {
1542 self.health_url = Some(trim_connect_base(url));
1543 self
1544 }
1545
1546 pub fn ingest_url(mut self, url: &str) -> Self {
1548 self.ingest_url = Some(trim_connect_base(url));
1549 self
1550 }
1551
1552 pub fn query_url(mut self, url: &str) -> Self {
1554 self.query_url = Some(trim_connect_base(url));
1555 self
1556 }
1557
1558 pub fn compact_url(mut self, url: &str) -> Self {
1560 self.compact_url = Some(trim_connect_base(url));
1561 self
1562 }
1563
1564 pub fn stream_url(mut self, url: &str) -> Self {
1566 self.stream_url = Some(trim_connect_base(url));
1567 self
1568 }
1569
1570 pub fn retry_config(mut self, retry: RetryConfig) -> Self {
1572 self.retry_config = retry.sanitized();
1573 self
1574 }
1575
1576 pub fn connect_request_compression(mut self, compression: ConnectRequestCompression) -> Self {
1578 self.connect_request_compression = compression;
1579 self
1580 }
1581
1582 pub fn build(self) -> Result<StoreClient, ClientBuildError> {
1584 let health_url = self.health_url.ok_or(ClientBuildError::MissingHealthUrl)?;
1585 let ingest_url = self.ingest_url.ok_or(ClientBuildError::MissingIngestUrl)?;
1586 let query_url = self.query_url.ok_or(ClientBuildError::MissingQueryUrl)?;
1587 let compact_url = self
1588 .compact_url
1589 .ok_or(ClientBuildError::MissingCompactUrl)?;
1590 let stream_url = self.stream_url.ok_or(ClientBuildError::MissingStreamUrl)?;
1591 let ingest_uri: http::Uri =
1592 ingest_url
1593 .parse()
1594 .map_err(|e| ClientBuildError::InvalidUrl {
1595 url: ingest_url.clone(),
1596 source: e,
1597 })?;
1598 let query_uri: http::Uri = query_url
1599 .parse()
1600 .map_err(|e| ClientBuildError::InvalidUrl {
1601 url: query_url.clone(),
1602 source: e,
1603 })?;
1604 let compact_uri: http::Uri =
1605 compact_url
1606 .parse()
1607 .map_err(|e| ClientBuildError::InvalidUrl {
1608 url: compact_url.clone(),
1609 source: e,
1610 })?;
1611 let stream_uri: http::Uri =
1612 stream_url
1613 .parse()
1614 .map_err(|e| ClientBuildError::InvalidUrl {
1615 url: stream_url.clone(),
1616 source: e,
1617 })?;
1618 Ok(StoreClient {
1619 health_url,
1620 ingest_uri,
1621 query_uri,
1622 compact_uri,
1623 stream_uri,
1624 http: new_http_client(),
1625 connect_http: ProtoPreferZstdHttpClient::plaintext(),
1626 retry_config: self.retry_config,
1627 connect_request_compression: self.connect_request_compression,
1628 })
1629 }
1630}
1631
1632#[derive(Clone, Debug)]
1634pub struct StoreClient {
1635 pub(crate) health_url: String,
1637 ingest_uri: http::Uri,
1638 query_uri: http::Uri,
1639 compact_uri: http::Uri,
1640 stream_uri: http::Uri,
1641 http: reqwest::Client,
1642 connect_http: ProtoPreferZstdHttpClient,
1643 retry_config: RetryConfig,
1644 connect_request_compression: ConnectRequestCompression,
1645}
1646
1647#[derive(Clone, Debug)]
1662pub struct SerializableReadSession {
1663 client: PrefixedStoreClient,
1664 state: Arc<SessionState>,
1665}
1666
1667#[derive(Debug)]
1668struct SessionState {
1669 sequence: Arc<AtomicU64>,
1670 init_gate: tokio::sync::Mutex<()>,
1671}
1672
1673impl SessionState {
1674 fn fixed_sequence(&self) -> Option<u64> {
1675 let sequence = self.sequence.load(Ordering::Acquire);
1676 (sequence > 0).then_some(sequence)
1677 }
1678}
1679
1680#[derive(Default)]
1681struct RangeStreamReadOptions {
1682 min_sequence_number: Option<u64>,
1683 observed_sequence: Option<Arc<AtomicU64>>,
1684}
1685
1686impl StoreClient {
1687 pub fn builder() -> StoreClientBuilder {
1689 StoreClientBuilder::default()
1690 }
1691
1692 pub fn new(url: &str) -> Self {
1693 Self::with_retry_config(url, RetryConfig::standard())
1694 }
1695
1696 pub fn with_retry_config(url: &str, retry_config: RetryConfig) -> Self {
1697 Self::builder()
1698 .url(url)
1699 .retry_config(retry_config)
1700 .build()
1701 .expect("url sets all service URLs")
1702 }
1703
1704 pub fn prefixed(&self, prefix: StoreKeyPrefix) -> PrefixedStoreClient {
1707 PrefixedStoreClient::new(self.clone(), prefix)
1708 }
1709
1710 pub fn connect_request_compression(&self) -> ConnectRequestCompression {
1712 self.connect_request_compression
1713 }
1714
1715 pub fn decode_error_details(
1716 err: &ConnectError,
1717 ) -> Result<exoware_proto::DecodedConnectError, buffa::DecodeError> {
1718 proto_decode_connect_error(err)
1719 }
1720
1721 pub(crate) async fn put_physical(&self, kvs: &[(&Key, &[u8])]) -> Result<u64, ClientError> {
1728 let mut proto_kvs = Vec::with_capacity(kvs.len());
1729 for (key, value) in kvs {
1730 if !is_valid_key_size(key.len()) {
1731 return Err(ClientError::WireFormat(format!(
1732 "key length {} is outside valid store key range ({}..={})",
1733 key.len(),
1734 keys::MIN_KEY_LEN,
1735 MAX_KEY_LEN
1736 )));
1737 }
1738 proto_kvs.push(exoware_proto::common::Entry {
1739 key: key.to_vec(),
1740 value: Bytes::copy_from_slice(value),
1741 ..Default::default()
1742 });
1743 }
1744 self.send_put(proto_kvs).await
1745 }
1746
1747 async fn put_prepared_physical(&self, kvs: &[(Key, Bytes)]) -> Result<u64, ClientError> {
1748 let mut proto_kvs = Vec::with_capacity(kvs.len());
1749 for (key, value) in kvs {
1750 if !is_valid_key_size(key.len()) {
1751 return Err(ClientError::WireFormat(format!(
1752 "key length {} is outside valid store key range ({}..={})",
1753 key.len(),
1754 keys::MIN_KEY_LEN,
1755 MAX_KEY_LEN
1756 )));
1757 }
1758 proto_kvs.push(exoware_proto::common::Entry {
1759 key: key.to_vec(),
1760 value: value.clone(),
1761 ..Default::default()
1762 });
1763 }
1764 self.send_put(proto_kvs).await
1765 }
1766
1767 async fn send_put(&self, kvs: Vec<exoware_proto::common::Entry>) -> Result<u64, ClientError> {
1768 let config =
1769 store_connect_client_config(self.ingest_uri.clone(), self.connect_request_compression);
1770 let client = IngestServiceClient::new(self.connect_http.clone(), config);
1771 let response = client
1772 .put(ProtoPutRequest {
1773 kvs,
1774 ..Default::default()
1775 })
1776 .await
1777 .map_err(client_error_from_connect)?;
1778 Ok(response.into_owned().sequence_number)
1779 }
1780
1781 pub(crate) async fn get(&self, key: &Key) -> Result<Option<Bytes>, ClientError> {
1782 self.get_internal(key, None).await
1783 }
1784
1785 pub(crate) async fn get_with_min_sequence_number(
1786 &self,
1787 key: &Key,
1788 min_sequence_number: u64,
1789 ) -> Result<Option<Bytes>, ClientError> {
1790 self.get_internal(key, Some(min_sequence_number)).await
1791 }
1792
1793 async fn get_internal(
1794 &self,
1795 key: &Key,
1796 min_sequence_number: Option<u64>,
1797 ) -> Result<Option<Bytes>, ClientError> {
1798 let (response, _detail) = self
1799 .send_get(key, self.normalize_min_sequence_number(min_sequence_number))
1800 .await?;
1801 Ok(response.value)
1802 }
1803
1804 pub(crate) async fn get_many_internal(
1808 &self,
1809 proto_keys: Vec<Vec<u8>>,
1810 batch_size: u32,
1811 min_sequence_number: Option<u64>,
1812 observed_sequence: Option<Arc<AtomicU64>>,
1813 ) -> Result<GetManyStream, ClientError> {
1814 let config =
1815 store_connect_client_config(self.query_uri.clone(), self.connect_request_compression);
1816 let client = QueryServiceClient::new(self.connect_http.clone(), config);
1817 let effective_min = self.normalize_min_sequence_number(min_sequence_number);
1818 let max_attempts = self.retry_config.max_attempts.max(1);
1819 let mut attempt = 1usize;
1820 loop {
1821 match client
1822 .get_many(ProtoGetManyRequest {
1823 keys: proto_keys.clone(),
1824 min_sequence_number: effective_min,
1825 batch_size,
1826 ..Default::default()
1827 })
1828 .await
1829 {
1830 Ok(stream) => {
1831 let mut gms =
1832 GetManyStream::from_connect_stream(stream, observed_sequence.clone(), None);
1833 if let Err(err) = gms.prefetch_first_frame().await {
1834 if attempt < max_attempts && is_retryable_error(&err) {
1835 let delay = retry_delay_for_error(&err, attempt, self.retry_config);
1836 tokio::time::sleep(delay).await;
1837 attempt += 1;
1838 continue;
1839 }
1840 return Err(client_error_from_connect(err));
1841 }
1842 return Ok(gms);
1843 }
1844 Err(err) => {
1845 if attempt < max_attempts && is_retryable_error(&err) {
1846 let delay = retry_delay_for_error(&err, attempt, self.retry_config);
1847 tokio::time::sleep(delay).await;
1848 attempt += 1;
1849 continue;
1850 }
1851 return Err(client_error_from_connect(err));
1852 }
1853 }
1854 }
1855 }
1856
1857 pub(crate) async fn prune(
1858 &self,
1859 policies: &[crate::prune_policy::PrunePolicy],
1860 ) -> Result<(), ClientError> {
1861 let config =
1862 store_connect_client_config(self.compact_uri.clone(), self.connect_request_compression);
1863 let client = CompactServiceClient::new(self.connect_http.clone(), config);
1864 client
1865 .prune(ProtoPruneRequest {
1866 policies: exoware_proto::prune_policies_to_proto(policies),
1867 ..Default::default()
1868 })
1869 .await
1870 .map_err(client_error_from_connect)?;
1871 Ok(())
1872 }
1873
1874 async fn subscribe_physical(
1878 &self,
1879 filter: crate::stream_filter::StreamFilter,
1880 since_sequence_number: Option<u64>,
1881 ) -> Result<StreamSubscription, ClientError> {
1882 crate::stream_filter::validate_filter(&filter)
1883 .map_err(|e| ClientError::WireFormat(e.to_string()))?;
1884 let selectors = filter
1885 .selectors
1886 .into_iter()
1887 .map(|mk| exoware_proto::common::kv::v1::Selector {
1888 prefix: mk.prefix,
1889 payload_regex: mk.payload_regex.0,
1890 ..Default::default()
1891 })
1892 .collect();
1893 let value_filters = filter
1894 .value_filters
1895 .into_iter()
1896 .map(|vf| {
1897 use crate::stream_filter::Filter;
1898 use exoware_proto::common::kv::v1::filter::Kind as ProtoKind;
1899 let kind = match vf {
1900 Filter::Exact(bytes) => ProtoKind::Exact(bytes),
1901 Filter::Prefix(bytes) => ProtoKind::Prefix(bytes),
1902 Filter::Regex(pattern) => ProtoKind::Regex(pattern),
1903 };
1904 exoware_proto::common::kv::v1::Filter {
1905 kind: Some(kind),
1906 ..Default::default()
1907 }
1908 })
1909 .collect();
1910 let request = exoware_proto::log::stream::v1::SubscribeRequest {
1911 selectors,
1912 value_filters,
1913 since_sequence_number,
1914 ..Default::default()
1915 };
1916 let config =
1917 store_connect_client_config(self.stream_uri.clone(), self.connect_request_compression);
1918 let client =
1919 exoware_proto::log::stream::v1::ServiceClient::new(self.connect_http.clone(), config);
1920 let stream = client
1921 .subscribe(request)
1922 .await
1923 .map_err(client_error_from_connect)?;
1924 Ok(StreamSubscription {
1925 stream,
1926 key_prefix: None,
1927 })
1928 }
1929
1930 async fn stream_get_physical(
1935 &self,
1936 sequence_number: u64,
1937 ) -> Result<Option<exoware_proto::log::stream::v1::GetResponse>, ClientError> {
1938 let config =
1939 store_connect_client_config(self.stream_uri.clone(), self.connect_request_compression);
1940 let client =
1941 exoware_proto::log::stream::v1::ServiceClient::new(self.connect_http.clone(), config);
1942 match client
1943 .get(exoware_proto::log::stream::v1::GetRequest {
1944 sequence_number,
1945 ..Default::default()
1946 })
1947 .await
1948 {
1949 Ok(resp) => Ok(Some(resp.into_owned())),
1950 Err(err) => {
1951 if is_batch_missing_error(&err) {
1952 Ok(None)
1953 } else {
1954 Err(client_error_from_connect(err))
1955 }
1956 }
1957 }
1958 }
1959
1960 pub async fn health(&self) -> Result<bool, ClientError> {
1961 let resp = self
1962 .http
1963 .get(format!("{}/health", self.health_url))
1964 .send()
1965 .await?;
1966 Ok(resp.status().is_success())
1967 }
1968
1969 pub async fn ready(&self) -> Result<bool, ClientError> {
1970 let resp = self
1971 .http
1972 .get(format!("{}/ready", self.health_url))
1973 .send()
1974 .await?;
1975 Ok(resp.status().is_success())
1976 }
1977
1978 fn normalize_min_sequence_number(&self, requested_sequence: Option<u64>) -> Option<u64> {
1979 requested_sequence.filter(|sequence| *sequence > 0)
1980 }
1981
1982 async fn send_get(
1983 &self,
1984 key: &Key,
1985 min_sequence_number: Option<u64>,
1986 ) -> Result<
1987 (
1988 exoware_proto::query::GetResponse,
1989 Option<proto_query::Detail>,
1990 ),
1991 ClientError,
1992 > {
1993 if !is_valid_key_size(key.len()) {
1994 return Err(ClientError::WireFormat(format!(
1995 "key length {} is outside valid store key range ({}..={})",
1996 key.len(),
1997 keys::MIN_KEY_LEN,
1998 MAX_KEY_LEN
1999 )));
2000 }
2001
2002 let config =
2003 store_connect_client_config(self.query_uri.clone(), self.connect_request_compression);
2004 let client = QueryServiceClient::new(self.connect_http.clone(), config);
2005 let response = self
2006 .send_with_retry(|| async {
2007 client
2008 .get(ProtoGetRequest {
2009 key: key.clone().into(),
2010 min_sequence_number,
2011 ..Default::default()
2012 })
2013 .await
2014 })
2015 .await?;
2016 let owned = response.into_owned();
2017 let detail = owned.detail.as_option().cloned();
2018 Ok((owned, detail))
2019 }
2020
2021 #[cfg(test)]
2022 pub async fn send_get_for_tests(
2023 &self,
2024 key: &Key,
2025 min_sequence_number: Option<u64>,
2026 ) -> Result<
2027 (
2028 exoware_proto::query::GetResponse,
2029 Option<proto_query::Detail>,
2030 ),
2031 ClientError,
2032 > {
2033 self.send_get(key, min_sequence_number).await
2034 }
2035
2036 async fn range_stream_internal(
2037 &self,
2038 start: &Key,
2039 end: &Key,
2040 limit: usize,
2041 batch_size: usize,
2042 mode: RangeMode,
2043 options: RangeStreamReadOptions,
2044 ) -> Result<RangeStream, ClientError> {
2045 if !is_valid_key_size(start.len()) || !is_valid_key_size(end.len()) {
2046 return Err(ClientError::WireFormat(
2047 "range start/end key length is outside valid store key range".to_string(),
2048 ));
2049 }
2050 if batch_size == 0 {
2051 return Err(ClientError::WireFormat(
2052 "batch_size must be positive".to_string(),
2053 ));
2054 }
2055
2056 let config =
2057 store_connect_client_config(self.query_uri.clone(), self.connect_request_compression);
2058 let client = QueryServiceClient::new(self.connect_http.clone(), config);
2059 let min_sequence_number = self.normalize_min_sequence_number(options.min_sequence_number);
2060 let max_attempts = self.retry_config.max_attempts.max(1);
2061 let mut attempt = 1usize;
2062 loop {
2063 let response = match client
2068 .range(ProtoRangeRequest {
2069 start: start.clone().into(),
2070 end: end.clone().into(),
2071 limit: Some(u32::try_from(limit).unwrap_or(u32::MAX)),
2072 batch_size: u32::try_from(batch_size).unwrap_or(u32::MAX),
2073 mode: mode.to_proto().into(),
2074 min_sequence_number,
2075 ..Default::default()
2076 })
2077 .await
2078 {
2079 Ok(response) => response,
2080 Err(err) => {
2081 if attempt < max_attempts && is_retryable_error(&err) {
2082 let delay = retry_delay_for_error(&err, attempt, self.retry_config);
2083 tracing::debug!(
2084 attempt,
2085 max_attempts,
2086 code = err.code.as_str(),
2087 delay_ms = delay.as_millis() as u64,
2088 "store client retrying transient range-open error",
2089 );
2090 tokio::time::sleep(delay).await;
2091 attempt += 1;
2092 continue;
2093 }
2094 return Err(client_error_from_connect(err));
2095 }
2096 };
2097
2098 let mut stream =
2099 RangeStream::from_connect_stream(response, options.observed_sequence.clone(), None);
2100 if let Err(err) = stream.prefetch_first_frame().await {
2101 if attempt < max_attempts && is_retryable_error(&err) {
2102 let delay = retry_delay_for_error(&err, attempt, self.retry_config);
2103 tracing::debug!(
2104 attempt,
2105 max_attempts,
2106 code = err.code.as_str(),
2107 delay_ms = delay.as_millis() as u64,
2108 "store client retrying transient stream-open error",
2109 );
2110 tokio::time::sleep(delay).await;
2111 attempt += 1;
2112 continue;
2113 }
2114 return Err(client_error_from_connect(err));
2115 }
2116 return Ok(stream);
2117 }
2118 }
2119
2120 async fn range_reduce_response_internal(
2121 &self,
2122 start: &Key,
2123 end: &Key,
2124 request: &DomainRangeReduceRequest,
2125 min_sequence_number: Option<u64>,
2126 ) -> Result<
2127 (
2128 exoware_proto::query::ReduceResponse,
2129 Option<proto_query::Detail>,
2130 ),
2131 ClientError,
2132 > {
2133 let config =
2134 store_connect_client_config(self.query_uri.clone(), self.connect_request_compression);
2135 let client = QueryServiceClient::new(self.connect_http.clone(), config);
2136 let proto_params = proto_to_proto_reduce_params(request.clone());
2137 let min_sequence_number = self.normalize_min_sequence_number(min_sequence_number);
2138 let response = self
2139 .send_with_retry(|| async {
2140 client
2141 .reduce(ProtoWireReduceRequest {
2142 start: start.clone().into(),
2143 end: end.clone().into(),
2144 params: Some(proto_params.clone()).into(),
2145 min_sequence_number,
2146 ..Default::default()
2147 })
2148 .await
2149 })
2150 .await?;
2151 let owned = response.into_owned();
2152 let detail = owned.detail.as_option().cloned();
2153 Ok((owned, detail))
2154 }
2155
2156 async fn send_with_retry<F, Fut, T>(&self, mut make_request: F) -> Result<T, ClientError>
2157 where
2158 F: FnMut() -> Fut,
2159 Fut: std::future::Future<Output = Result<T, ConnectError>>,
2160 {
2161 let max_attempts = self.retry_config.max_attempts.max(1);
2162 let mut attempt = 1usize;
2163 loop {
2164 match make_request().await {
2165 Ok(response) => return Ok(response),
2166 Err(err) => {
2167 if attempt < max_attempts && is_retryable_error(&err) {
2168 let delay = retry_delay_for_error(&err, attempt, self.retry_config);
2169 tracing::debug!(
2170 attempt,
2171 max_attempts,
2172 code = err.code.as_str(),
2173 delay_ms = delay.as_millis() as u64,
2174 "store client retrying transient RPC error",
2175 );
2176 tokio::time::sleep(delay).await;
2177 attempt += 1;
2178 continue;
2179 }
2180 return Err(client_error_from_connect(err));
2181 }
2182 }
2183 }
2184 }
2185}
2186
2187fn shift_reduce_request_key_offsets(
2188 prefix_len: usize,
2189 request: &mut DomainRangeReduceRequest,
2190) -> Result<(), StoreKeyPrefixError> {
2191 debug_assert!(prefix_len <= MAX_KEY_LEN);
2197 let shift_bytes = prefix_len as u16;
2198 let shift_bits = shift_bytes * 8;
2199 for reducer in &mut request.reducers {
2200 if let Some(expr) = &mut reducer.expr {
2201 shift_expr_key_offsets(shift_bytes, shift_bits, expr)?;
2202 }
2203 }
2204 for expr in &mut request.group_by {
2205 shift_expr_key_offsets(shift_bytes, shift_bits, expr)?;
2206 }
2207 if let Some(filter) = &mut request.filter {
2208 for check in &mut filter.checks {
2209 shift_field_ref_key_offset(shift_bytes, shift_bits, &mut check.field)?;
2210 }
2211 }
2212 Ok(())
2213}
2214
2215fn shift_expr_key_offsets(
2216 shift_bytes: u16,
2217 shift_bits: u16,
2218 expr: &mut KvExpr,
2219) -> Result<(), StoreKeyPrefixError> {
2220 match expr {
2221 KvExpr::Field(field) => shift_field_ref_key_offset(shift_bytes, shift_bits, field),
2222 KvExpr::Literal(_) => Ok(()),
2223 KvExpr::Add(left, right)
2224 | KvExpr::Sub(left, right)
2225 | KvExpr::Mul(left, right)
2226 | KvExpr::Div(left, right) => {
2227 shift_expr_key_offsets(shift_bytes, shift_bits, left)?;
2228 shift_expr_key_offsets(shift_bytes, shift_bits, right)
2229 }
2230 KvExpr::Lower(inner) | KvExpr::DateTruncDay(inner) => {
2231 shift_expr_key_offsets(shift_bytes, shift_bits, inner)
2232 }
2233 }
2234}
2235
2236fn shift_field_ref_key_offset(
2237 shift_bytes: u16,
2238 shift_bits: u16,
2239 field: &mut KvFieldRef,
2240) -> Result<(), StoreKeyPrefixError> {
2241 match field {
2242 KvFieldRef::Key { byte_offset, .. } => {
2243 *byte_offset = byte_offset.checked_add(shift_bytes).ok_or(
2244 StoreKeyPrefixError::KeyOffsetOverflow {
2245 offset: *byte_offset,
2246 shift: shift_bytes,
2247 },
2248 )?;
2249 Ok(())
2250 }
2251 KvFieldRef::ZOrderKey { bit_offset, .. } => {
2252 *bit_offset = bit_offset.checked_add(shift_bits).ok_or(
2253 StoreKeyPrefixError::KeyOffsetOverflow {
2254 offset: *bit_offset,
2255 shift: shift_bits,
2256 },
2257 )?;
2258 Ok(())
2259 }
2260 KvFieldRef::Value { .. } => Ok(()),
2261 }
2262}
2263
2264#[derive(Clone, Copy, Debug)]
2267pub struct Ingest<'a> {
2268 c: &'a PrefixedStoreClient,
2269}
2270
2271#[derive(Clone, Copy, Debug)]
2272pub struct Query<'a> {
2273 c: &'a PrefixedStoreClient,
2274}
2275
2276#[derive(Clone, Copy, Debug)]
2277pub struct Compact<'a> {
2278 c: &'a PrefixedStoreClient,
2279}
2280
2281#[derive(Clone, Copy, Debug)]
2282pub struct Stream<'a> {
2283 c: &'a PrefixedStoreClient,
2284}
2285
2286impl<'a> Ingest<'a> {
2287 pub async fn put(&self, kvs: &[(&Key, &[u8])]) -> Result<u64, ClientError> {
2288 self.c.put(kvs).await
2289 }
2290
2291 pub async fn put_prepared(&self, batch: &StoreWriteBatch) -> Result<u64, ClientError> {
2294 batch.commit(self.c.client()).await
2295 }
2296}
2297
2298impl<'a> Query<'a> {
2299 pub async fn get(&self, key: &Key) -> Result<Option<Bytes>, ClientError> {
2300 self.c.get(key).await
2301 }
2302
2303 pub async fn get_with_min_sequence_number(
2304 &self,
2305 key: &Key,
2306 min_sequence_number: u64,
2307 ) -> Result<Option<Bytes>, ClientError> {
2308 self.c
2309 .get_with_min_sequence_number(key, min_sequence_number)
2310 .await
2311 }
2312
2313 pub async fn get_many(
2314 &self,
2315 keys: &[&Key],
2316 batch_size: u32,
2317 ) -> Result<GetManyStream, ClientError> {
2318 self.c.get_many(keys, batch_size).await
2319 }
2320
2321 pub async fn get_many_with_min_sequence_number(
2322 &self,
2323 keys: &[&Key],
2324 batch_size: u32,
2325 min_sequence_number: u64,
2326 ) -> Result<GetManyStream, ClientError> {
2327 self.c
2328 .get_many_with_min_sequence_number(keys, batch_size, min_sequence_number)
2329 .await
2330 }
2331
2332 pub async fn range(
2334 &self,
2335 start: &Key,
2336 end: &Key,
2337 limit: usize,
2338 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2339 self.c.range(start, end, limit).await
2340 }
2341
2342 pub async fn range_with_mode(
2343 &self,
2344 start: &Key,
2345 end: &Key,
2346 limit: usize,
2347 mode: RangeMode,
2348 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2349 self.c.range_with_mode(start, end, limit, mode).await
2350 }
2351
2352 pub async fn range_with_min_sequence_number(
2353 &self,
2354 start: &Key,
2355 end: &Key,
2356 limit: usize,
2357 min_sequence_number: u64,
2358 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2359 self.c
2360 .range_with_min_sequence_number(start, end, limit, min_sequence_number)
2361 .await
2362 }
2363
2364 pub async fn range_with_mode_and_min_sequence_number(
2365 &self,
2366 start: &Key,
2367 end: &Key,
2368 limit: usize,
2369 mode: RangeMode,
2370 min_sequence_number: u64,
2371 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2372 self.c
2373 .range_with_mode_and_min_sequence_number(start, end, limit, mode, min_sequence_number)
2374 .await
2375 }
2376
2377 pub async fn range_stream(
2378 &self,
2379 start: &Key,
2380 end: &Key,
2381 limit: usize,
2382 batch_size: usize,
2383 ) -> Result<RangeStream, ClientError> {
2384 self.c.range_stream(start, end, limit, batch_size).await
2385 }
2386
2387 pub async fn range_stream_with_mode(
2388 &self,
2389 start: &Key,
2390 end: &Key,
2391 limit: usize,
2392 batch_size: usize,
2393 mode: RangeMode,
2394 ) -> Result<RangeStream, ClientError> {
2395 self.c
2396 .range_stream_with_mode(start, end, limit, batch_size, mode)
2397 .await
2398 }
2399
2400 pub async fn range_stream_with_min_sequence_number(
2401 &self,
2402 start: &Key,
2403 end: &Key,
2404 limit: usize,
2405 batch_size: usize,
2406 min_sequence_number: u64,
2407 ) -> Result<RangeStream, ClientError> {
2408 self.c
2409 .range_stream_with_min_sequence_number(
2410 start,
2411 end,
2412 limit,
2413 batch_size,
2414 min_sequence_number,
2415 )
2416 .await
2417 }
2418
2419 pub async fn range_stream_with_mode_and_min_sequence_number(
2420 &self,
2421 start: &Key,
2422 end: &Key,
2423 limit: usize,
2424 batch_size: usize,
2425 mode: RangeMode,
2426 min_sequence_number: u64,
2427 ) -> Result<RangeStream, ClientError> {
2428 self.c
2429 .range_stream_with_mode_and_min_sequence_number(
2430 start,
2431 end,
2432 limit,
2433 batch_size,
2434 mode,
2435 min_sequence_number,
2436 )
2437 .await
2438 }
2439
2440 pub async fn range_reduce(
2441 &self,
2442 start: &Key,
2443 end: &Key,
2444 request: &DomainRangeReduceRequest,
2445 ) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
2446 self.c.range_reduce(start, end, request).await
2447 }
2448
2449 pub async fn range_reduce_with_min_sequence_number(
2450 &self,
2451 start: &Key,
2452 end: &Key,
2453 request: &DomainRangeReduceRequest,
2454 min_sequence_number: u64,
2455 ) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
2456 self.c
2457 .range_reduce_with_min_sequence_number(start, end, request, min_sequence_number)
2458 .await
2459 }
2460
2461 pub async fn range_reduce_response(
2462 &self,
2463 start: &Key,
2464 end: &Key,
2465 request: &DomainRangeReduceRequest,
2466 ) -> Result<exoware_proto::query::ReduceResponse, ClientError> {
2467 self.c.range_reduce_response(start, end, request).await
2468 }
2469
2470 pub async fn range_reduce_response_with_min_sequence_number(
2471 &self,
2472 start: &Key,
2473 end: &Key,
2474 request: &DomainRangeReduceRequest,
2475 min_sequence_number: u64,
2476 ) -> Result<exoware_proto::query::ReduceResponse, ClientError> {
2477 self.c
2478 .range_reduce_response_with_min_sequence_number(
2479 start,
2480 end,
2481 request,
2482 min_sequence_number,
2483 )
2484 .await
2485 }
2486}
2487
2488impl<'a> Compact<'a> {
2489 pub async fn prune(
2490 &self,
2491 policies: &[crate::prune_policy::PrunePolicy],
2492 ) -> Result<(), ClientError> {
2493 self.c.prune(policies).await
2494 }
2495}
2496
2497impl<'a> Stream<'a> {
2498 pub async fn subscribe(
2504 &self,
2505 filter: crate::stream_filter::StreamFilter,
2506 since_sequence_number: Option<u64>,
2507 ) -> Result<StreamSubscription, ClientError> {
2508 self.c.subscribe(filter, since_sequence_number).await
2509 }
2510
2511 pub async fn get(
2514 &self,
2515 sequence_number: u64,
2516 ) -> Result<Option<Vec<(Key, Bytes)>>, ClientError> {
2517 self.c.stream_get(sequence_number).await
2518 }
2519}
2520
2521impl SerializableReadSession {
2522 pub fn fixed_sequence(&self) -> Option<u64> {
2529 self.state.fixed_sequence()
2530 }
2531
2532 pub async fn get(&self, key: &Key) -> Result<Option<Bytes>, ClientError> {
2533 let seeded_client = self.client.clone();
2534 let unseeded_client = self.client.clone();
2535 self.run_read(
2536 move |sequence| {
2537 let client = seeded_client.clone();
2538 async move { client.get_with_min_sequence_number(key, sequence).await }
2539 },
2540 move |observed_sequence| {
2541 let client = unseeded_client.clone();
2542 async move {
2543 let (response, detail) = client.send_get(key, None).await?;
2544 if let Some(detail) = detail {
2545 observed_sequence.fetch_max(detail.sequence_number, Ordering::SeqCst);
2546 }
2547 Ok(response.value)
2548 }
2549 },
2550 )
2551 .await
2552 }
2553
2554 pub async fn get_many(
2555 &self,
2556 keys: &[&Key],
2557 batch_size: u32,
2558 ) -> Result<GetManyStream, ClientError> {
2559 let keys_owned: Vec<Key> = keys.iter().map(|k| (**k).clone()).collect();
2562 let seeded_client = self.client.clone();
2563 let unseeded_client = self.client.clone();
2564 let keys_seeded = keys_owned.clone();
2565 let keys_unseeded = keys_owned;
2566 self.run_read(
2567 move |sequence| {
2568 let client = seeded_client.clone();
2569 let keys = keys_seeded.clone();
2570 async move {
2571 let refs: Vec<&Key> = keys.iter().collect();
2572 client
2573 .get_many_with_min_sequence_number(&refs, batch_size, sequence)
2574 .await
2575 }
2576 },
2577 move |observed_sequence| {
2578 let client = unseeded_client.clone();
2579 let keys = keys_unseeded.clone();
2580 async move {
2581 let refs: Vec<&Key> = keys.iter().collect();
2582 client
2583 .get_many_internal(&refs, batch_size, None, Some(observed_sequence))
2584 .await
2585 }
2586 },
2587 )
2588 .await
2589 }
2590
2591 pub async fn range(
2592 &self,
2593 start: &Key,
2594 end: &Key,
2595 limit: usize,
2596 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2597 self.range_with_mode(start, end, limit, RangeMode::Forward)
2598 .await
2599 }
2600
2601 pub async fn range_with_mode(
2602 &self,
2603 start: &Key,
2604 end: &Key,
2605 limit: usize,
2606 mode: RangeMode,
2607 ) -> Result<Vec<(Key, Bytes)>, ClientError> {
2608 let seeded_client = self.client.clone();
2609 let unseeded_client = self.client.clone();
2610 self.run_read(
2611 move |sequence| {
2612 let client = seeded_client.clone();
2613 async move {
2614 client
2615 .range_with_mode_and_min_sequence_number(start, end, limit, mode, sequence)
2616 .await
2617 }
2618 },
2619 move |observed_sequence| {
2620 let client = unseeded_client.clone();
2621 async move {
2622 let stream = client
2623 .range_stream_internal(
2624 start,
2625 end,
2626 limit,
2627 limit.max(1),
2628 mode,
2629 RangeStreamReadOptions {
2630 min_sequence_number: None,
2631 observed_sequence: Some(observed_sequence),
2632 },
2633 )
2634 .await;
2635 stream?.collect().await
2636 }
2637 },
2638 )
2639 .await
2640 }
2641
2642 pub async fn range_stream(
2643 &self,
2644 start: &Key,
2645 end: &Key,
2646 limit: usize,
2647 batch_size: usize,
2648 ) -> Result<RangeStream, ClientError> {
2649 self.range_stream_with_mode(start, end, limit, batch_size, RangeMode::Forward)
2650 .await
2651 }
2652
2653 pub async fn range_stream_with_mode(
2654 &self,
2655 start: &Key,
2656 end: &Key,
2657 limit: usize,
2658 batch_size: usize,
2659 mode: RangeMode,
2660 ) -> Result<RangeStream, ClientError> {
2661 let seeded_client = self.client.clone();
2662 let unseeded_client = self.client.clone();
2663 self.run_read(
2664 move |sequence| {
2665 let client = seeded_client.clone();
2666 async move {
2667 client
2668 .range_stream_with_mode_and_min_sequence_number(
2669 start, end, limit, batch_size, mode, sequence,
2670 )
2671 .await
2672 }
2673 },
2674 move |observed_sequence| {
2675 let client = unseeded_client.clone();
2676 async move {
2677 client
2678 .range_stream_internal(
2679 start,
2680 end,
2681 limit,
2682 batch_size,
2683 mode,
2684 RangeStreamReadOptions {
2685 min_sequence_number: None,
2686 observed_sequence: Some(observed_sequence),
2687 },
2688 )
2689 .await
2690 }
2691 },
2692 )
2693 .await
2694 }
2695
2696 pub async fn range_reduce(
2697 &self,
2698 start: &Key,
2699 end: &Key,
2700 request: &DomainRangeReduceRequest,
2701 ) -> Result<Vec<Option<KvReducedValue>>, ClientError> {
2702 let seeded_client = self.client.clone();
2703 let unseeded_client = self.client.clone();
2704 let request_seeded = request.clone();
2705 let request_unseeded = request.clone();
2706 self.run_read(
2707 move |sequence| {
2708 let client = seeded_client.clone();
2709 let request = request_seeded.clone();
2710 async move {
2711 client
2712 .range_reduce_with_min_sequence_number(start, end, &request, sequence)
2713 .await
2714 }
2715 },
2716 move |observed_sequence| {
2717 let client = unseeded_client.clone();
2718 let request = request_unseeded.clone();
2719 async move {
2720 let (response, detail) = client
2721 .range_reduce_response_internal(start, end, &request, None)
2722 .await?;
2723 if let Some(detail) = detail {
2724 observed_sequence.fetch_max(detail.sequence_number, Ordering::SeqCst);
2725 }
2726 let decoded = proto_to_domain_reduce_response(response)
2727 .map_err(ClientError::WireFormat)?;
2728 if !decoded.groups.is_empty() {
2729 return Err(ClientError::WireFormat(
2730 "grouped range reduction response returned for scalar request"
2731 .to_string(),
2732 ));
2733 }
2734 Ok(decoded
2735 .results
2736 .iter()
2737 .map(|result| result.value.clone())
2738 .collect())
2739 }
2740 },
2741 )
2742 .await
2743 }
2744
2745 pub async fn range_reduce_response(
2746 &self,
2747 start: &Key,
2748 end: &Key,
2749 request: &DomainRangeReduceRequest,
2750 ) -> Result<exoware_proto::query::ReduceResponse, ClientError> {
2751 let seeded_client = self.client.clone();
2752 let unseeded_client = self.client.clone();
2753 let request_seeded = request.clone();
2754 let request_unseeded = request.clone();
2755 self.run_read(
2756 move |sequence| {
2757 let client = seeded_client.clone();
2758 let request = request_seeded.clone();
2759 async move {
2760 client
2761 .range_reduce_response_with_min_sequence_number(
2762 start, end, &request, sequence,
2763 )
2764 .await
2765 }
2766 },
2767 move |observed_sequence| {
2768 let client = unseeded_client.clone();
2769 let request = request_unseeded.clone();
2770 async move {
2771 let (response, detail) = client
2772 .range_reduce_response_internal(start, end, &request, None)
2773 .await?;
2774 if let Some(detail) = detail {
2775 observed_sequence.fetch_max(detail.sequence_number, Ordering::SeqCst);
2776 }
2777 Ok(response)
2778 }
2779 },
2780 )
2781 .await
2782 }
2783
2784 async fn run_read<T, SeededCall, SeededFut, UnseededCall, UnseededFut>(
2785 &self,
2786 seeded_call: SeededCall,
2787 unseeded_call: UnseededCall,
2788 ) -> Result<T, ClientError>
2789 where
2790 SeededCall: Fn(u64) -> SeededFut,
2791 SeededFut: std::future::Future<Output = Result<T, ClientError>>,
2792 UnseededCall: Fn(Arc<AtomicU64>) -> UnseededFut,
2793 UnseededFut: std::future::Future<Output = Result<T, ClientError>>,
2794 {
2795 if let Some(sequence) = self.fixed_sequence() {
2796 return seeded_call(sequence).await;
2797 }
2798
2799 let gate = self.state.init_gate.lock().await;
2800
2801 if let Some(sequence) = self.fixed_sequence() {
2802 drop(gate);
2803 return seeded_call(sequence).await;
2804 }
2805
2806 let result = unseeded_call(self.state.sequence.clone()).await;
2807 drop(gate);
2808 result
2809 }
2810}
2811
2812fn client_error_from_connect(err: ConnectError) -> ClientError {
2813 ClientError::Rpc(Box::new(err))
2814}
2815
2816fn is_retryable_error(err: &ConnectError) -> bool {
2817 matches!(
2818 err.code,
2819 ErrorCode::Aborted
2820 | ErrorCode::ResourceExhausted
2821 | ErrorCode::Unavailable
2822 | ErrorCode::Unknown
2823 | ErrorCode::Internal
2828 )
2829}
2830
2831fn retry_delay_for_error(
2832 err: &ConnectError,
2833 attempt: usize,
2834 retry_config: RetryConfig,
2835) -> Duration {
2836 if let Ok(decoded) = proto_decode_connect_error(err) {
2837 if let Some(retry_info) = decoded.retry_info {
2838 if let Some(delay) = retry_info.retry_delay.as_option() {
2839 let secs = u64::try_from(delay.seconds).unwrap_or(0);
2840 let nanos = u32::try_from(delay.nanos.max(0)).unwrap_or(0);
2841 let hinted = Duration::new(secs, nanos);
2842 if !hinted.is_zero() {
2843 return hinted.min(retry_config.max_backoff);
2844 }
2845 }
2846 }
2847 }
2848 retry_backoff_delay(attempt, retry_config)
2849}
2850
2851fn retry_backoff_delay(attempt: usize, retry_config: RetryConfig) -> Duration {
2852 let exponent = (attempt.saturating_sub(1)).min(20) as u32;
2853 let factor = 1u128 << exponent;
2854 let base_ms = retry_config.initial_backoff.as_millis();
2855 let capped_ms = base_ms
2856 .saturating_mul(factor)
2857 .min(retry_config.max_backoff.as_millis());
2858 Duration::from_millis(capped_ms.min(u64::MAX as u128) as u64)
2859}
2860
2861#[cfg(test)]
2862mod tests {
2863 use super::*;
2864 use crate::kv_codec::{KvFieldKind, KvPredicate, KvPredicateCheck, KvPredicateConstraint};
2865 use exoware_proto::query::TraversalMode as ProtoTraversalMode;
2866
2867 #[test]
2868 fn hex_round_trip() {
2869 let data = vec![0x00, 0x42, 0xFF, 0xAB];
2870 let encoded = hex_encode(&data);
2871 assert_eq!(encoded, "0042ffab");
2872 let decoded = hex_decode(&encoded).unwrap();
2873 assert_eq!(decoded, data);
2874 }
2875
2876 #[test]
2877 fn client_creation() {
2878 let client = StoreClient::new("http://localhost:10000");
2879 assert_eq!(client.health_url, "http://localhost:10000");
2880 assert_eq!(client.ingest_uri.to_string(), "http://localhost:10000/");
2881 assert_eq!(client.query_uri.to_string(), "http://localhost:10000/");
2882 assert_eq!(client.stream_uri.to_string(), "http://localhost:10000/");
2883 }
2884
2885 #[test]
2886 fn builder_fails_until_all_urls_set() {
2887 assert!(matches!(
2888 StoreClient::builder().health_url("http://h").build(),
2889 Err(ClientBuildError::MissingIngestUrl)
2890 ));
2891 assert!(matches!(
2892 StoreClient::builder()
2893 .health_url("http://h")
2894 .ingest_url("http://i")
2895 .build(),
2896 Err(ClientBuildError::MissingQueryUrl)
2897 ));
2898 assert!(matches!(
2899 StoreClient::builder()
2900 .health_url("http://h")
2901 .ingest_url("http://i")
2902 .query_url("http://q")
2903 .build(),
2904 Err(ClientBuildError::MissingCompactUrl)
2905 ));
2906 assert!(matches!(
2907 StoreClient::builder()
2908 .health_url("http://h")
2909 .ingest_url("http://i")
2910 .query_url("http://q")
2911 .compact_url("http://c")
2912 .build(),
2913 Err(ClientBuildError::MissingStreamUrl)
2914 ));
2915 }
2916
2917 #[test]
2918 fn client_trims_trailing_slash() {
2919 let client = StoreClient::new("http://localhost:10000/");
2920 assert_eq!(client.health_url, "http://localhost:10000");
2921 }
2922
2923 #[test]
2924 fn create_session_starts_unseeded() {
2925 let client = PrefixedStoreClient::empty(StoreClient::new("http://localhost:10000/"));
2926 let session = client.create_session();
2927 assert_eq!(session.fixed_sequence(), None);
2928 }
2929
2930 #[test]
2931 fn range_mode_maps_to_proto_traversal() {
2932 assert_eq!(
2933 RangeMode::Forward.to_proto(),
2934 ProtoTraversalMode::TRAVERSAL_MODE_FORWARD
2935 );
2936 assert_eq!(
2937 RangeMode::Reverse.to_proto(),
2938 ProtoTraversalMode::TRAVERSAL_MODE_REVERSE
2939 );
2940 }
2941
2942 #[test]
2943 fn retry_config_standard_defaults_match_expected() {
2944 let config = RetryConfig::standard();
2945 assert_eq!(config.max_attempts, 3);
2946 assert_eq!(config.initial_backoff, Duration::from_millis(100));
2947 assert_eq!(config.max_backoff, Duration::from_millis(2_000));
2948 }
2949
2950 #[test]
2951 fn retry_config_clamps_attempts_and_backoff_bounds() {
2952 let config = RetryConfig::standard()
2953 .with_max_attempts(0)
2954 .with_initial_backoff(Duration::from_millis(250))
2955 .with_max_backoff(Duration::from_millis(50))
2956 .sanitized();
2957 assert_eq!(config.max_attempts, 1);
2958 assert_eq!(config.initial_backoff, Duration::from_millis(250));
2959 assert_eq!(config.max_backoff, Duration::from_millis(250));
2960 }
2961
2962 #[test]
2963 fn retryable_codes_include_connect_transients() {
2964 assert!(is_retryable_error(&ConnectError::aborted("retry")));
2965 assert!(is_retryable_error(&ConnectError::resource_exhausted(
2966 "retry"
2967 )));
2968 assert!(is_retryable_error(&ConnectError::unavailable("retry")));
2969 assert!(is_retryable_error(&ConnectError::internal("retry")));
2970 assert!(!is_retryable_error(&ConnectError::invalid_argument(
2971 "no retry"
2972 )));
2973 }
2974
2975 #[test]
2976 fn retry_backoff_delay_is_exponential_and_capped() {
2977 let config = RetryConfig::standard()
2978 .with_initial_backoff(Duration::from_millis(100))
2979 .with_max_backoff(Duration::from_millis(250));
2980 assert_eq!(retry_backoff_delay(1, config), Duration::from_millis(100));
2981 assert_eq!(retry_backoff_delay(2, config), Duration::from_millis(200));
2982 assert_eq!(retry_backoff_delay(3, config), Duration::from_millis(250));
2983 assert_eq!(retry_backoff_delay(4, config), Duration::from_millis(250));
2984 }
2985
2986 #[test]
2987 fn create_session_with_sequence_pins_explicit_floor() {
2988 let client = PrefixedStoreClient::empty(StoreClient::new("http://localhost:10000/"));
2989 let session = client.create_session_with_sequence(27);
2990 assert_eq!(session.fixed_sequence(), Some(27));
2991 }
2992
2993 #[test]
2994 fn store_key_prefix_round_trips_logical_keys() {
2995 let prefix = StoreKeyPrefix::new(vec![0x0A]).unwrap();
2996 let logical = Bytes::from_static(b"hello");
2997 let physical = prefix.encode_key(&logical).unwrap();
2998 assert!(prefix.matches(&physical));
2999 assert_eq!(prefix.decode_key(&physical).unwrap(), logical);
3000 }
3001
3002 #[test]
3003 fn identity_prefix_encode_decode_are_zero_copy() {
3004 let prefix = StoreKeyPrefix::identity();
3005 let logical = Bytes::from_static(b"passthrough-key");
3006 let physical = prefix.encode_key(&logical).unwrap();
3007 assert_eq!(physical.as_ptr(), logical.as_ptr());
3009 let decoded = prefix.decode_key(&physical).unwrap();
3010 assert_eq!(decoded.as_ptr(), logical.as_ptr());
3011 assert_eq!(decoded, logical);
3012 }
3013
3014 #[test]
3015 fn uniform_width_prefixes_are_pairwise_disjoint() {
3016 let all = [
3018 StoreKeyPrefix::new(vec![0]).unwrap(),
3019 StoreKeyPrefix::new(vec![1]).unwrap(),
3020 StoreKeyPrefix::new(vec![2]).unwrap(),
3021 StoreKeyPrefix::new(vec![3]).unwrap(),
3022 StoreKeyPrefix::new(vec![4]).unwrap(),
3023 StoreKeyPrefix::new(vec![5]).unwrap(),
3024 ];
3025 let logical = Bytes::from_static(b"\x00\x10whatever-block-meta-or-op-log-key");
3029 for (i, pa) in all.iter().enumerate() {
3030 let physical = pa.encode_key(&logical).unwrap();
3031 assert!(pa.matches(&physical));
3032 assert_eq!(pa.decode_key(&physical).unwrap(), logical);
3033 for (j, pb) in all.iter().enumerate() {
3034 if i == j {
3035 continue;
3036 }
3037 assert!(
3038 !pb.matches(&physical),
3039 "prefix {j} matched a key encoded under prefix {i}",
3040 );
3041 }
3042 }
3043 }
3044
3045 #[test]
3046 fn prefixed_store_client_always_carries_its_prefix() {
3047 let base = StoreClient::new("http://localhost:8090");
3048 let prefix = StoreKeyPrefix::new(vec![0]).unwrap();
3049 let client = base.prefixed(prefix.clone());
3050 assert_eq!(client.key_prefix(), &prefix);
3051 let logical = Bytes::from_static(b"row");
3053 assert_eq!(
3054 client.encode_store_key(&logical).unwrap(),
3055 prefix.encode_key(&logical).unwrap(),
3056 );
3057 }
3058
3059 #[test]
3060 fn store_key_prefix_clamps_long_logical_range_upper_bound() {
3061 let prefix = StoreKeyPrefix::new(vec![0x02]).unwrap();
3062 let logical_start = Bytes::new();
3065 let logical_end = Bytes::from(vec![0xFFu8; MAX_KEY_LEN]);
3066 assert_eq!(logical_end.len(), MAX_KEY_LEN);
3067
3068 let (physical_start, physical_end) =
3069 prefix.encode_range(&logical_start, &logical_end).unwrap();
3070 assert!(prefix.matches(&physical_start));
3071 assert!(prefix.matches(&physical_end));
3072 assert_eq!(prefix.max_logical_key_len(), MAX_KEY_LEN - 1);
3075 assert_eq!(physical_end.len(), MAX_KEY_LEN);
3076 assert_eq!(prefix.decode_key(&physical_start).unwrap(), logical_start);
3077 }
3078
3079 #[test]
3080 fn store_key_prefix_rewrites_selector_family() {
3081 let prefix = StoreKeyPrefix::new(vec![0x05]).unwrap();
3082 let logical = crate::selector::Selector {
3083 prefix: Bytes::copy_from_slice(&[0x06]),
3084 payload_regex: crate::kv_codec::Utf8::from("(?s).*"),
3085 };
3086 let physical = prefix.prefix_selector(&logical).unwrap();
3087 assert_eq!(physical.prefix.as_ref(), &[0x05, 0x06]);
3088 assert_eq!(physical.payload_regex, logical.payload_regex);
3089 }
3090
3091 #[test]
3092 fn store_key_prefix_composed_selector_reconstructs_and_strips_cleanly() {
3093 let prefix = StoreKeyPrefix::new(vec![0x05, 0x06]).unwrap();
3104 let logical = crate::selector::Selector {
3105 prefix: Bytes::copy_from_slice(&[0x07]),
3106 payload_regex: crate::kv_codec::Utf8::from("(?s).*"),
3107 };
3108
3109 let merged = prefix.prefix_selector(&logical).unwrap();
3112 assert_eq!(merged.prefix.as_ref(), &[0x05, 0x06, 0x07]);
3113
3114 let codec = Prefix::new(merged.prefix.clone()).unwrap();
3116
3117 let payload = [0xAA, 0xBB, 0xCC];
3119 let key = codec.encode(&payload).unwrap();
3120 assert_eq!(key.as_ref(), &[0x05, 0x06, 0x07, 0xAA, 0xBB, 0xCC]);
3121 assert!(codec.matches(&key));
3122 assert_eq!(
3125 codec.strip(&key).unwrap(),
3126 Bytes::copy_from_slice(&[0xAA, 0xBB, 0xCC])
3127 );
3128
3129 let (start, end) = codec.bounds();
3132 assert!(start <= key && key <= end);
3133
3134 let sibling = Prefix::new(vec![0x05, 0x06, 0x08])
3136 .unwrap()
3137 .encode(&[0x11])
3138 .unwrap();
3139 assert!(!codec.matches(&sibling));
3140 }
3141
3142 #[test]
3143 fn store_key_prefix_stream_filter_passes_payload_regex_through() {
3144 let client = StoreClient::builder()
3148 .url("http://localhost:10000")
3149 .build()
3150 .unwrap()
3151 .prefixed(StoreKeyPrefix::new(vec![0x05]).unwrap());
3152 let filter = crate::stream_filter::StreamFilter {
3153 selectors: vec![crate::selector::Selector {
3154 prefix: Bytes::copy_from_slice(&[0x06]),
3155 payload_regex: crate::kv_codec::Utf8::from("(?s)foo.*"),
3156 }],
3157 value_filters: vec![],
3158 };
3159 let physical = client.prefix_stream_filter(filter.clone()).unwrap();
3160 assert_eq!(physical.selectors[0].prefix.as_ref(), &[0x05, 0x06]);
3161 assert_eq!(
3162 physical.selectors[0].payload_regex,
3163 filter.selectors[0].payload_regex
3164 );
3165 }
3166
3167 #[test]
3168 fn prefixed_reduce_request_shifts_key_field_offsets() {
3169 let client = StoreClient::builder()
3170 .url("http://localhost:10000")
3171 .build()
3172 .unwrap()
3173 .prefixed(StoreKeyPrefix::new(vec![0x01, 0x02, 0x03]).unwrap());
3174 let request = DomainRangeReduceRequest {
3175 reducers: vec![crate::RangeReducerSpec {
3176 op: crate::RangeReduceOp::SumField,
3177 expr: Some(KvExpr::Field(KvFieldRef::Key {
3178 byte_offset: 9,
3179 kind: KvFieldKind::UInt64,
3180 })),
3181 }],
3182 group_by: vec![KvExpr::Field(KvFieldRef::ZOrderKey {
3183 bit_offset: 12,
3184 field_position: 0,
3185 field_widths: vec![8],
3186 kind: KvFieldKind::UInt64,
3187 })],
3188 filter: Some(KvPredicate {
3189 checks: vec![KvPredicateCheck {
3190 field: KvFieldRef::Value {
3191 index: 0,
3192 kind: KvFieldKind::UInt64,
3193 nullable: false,
3194 },
3195 constraint: KvPredicateConstraint::UInt64Range {
3196 min: Some(1),
3197 max: Some(9),
3198 },
3199 }],
3200 contradiction: false,
3201 }),
3202 };
3203
3204 let shifted = client.prefix_reduce_request(&request).unwrap();
3205 let Some(KvExpr::Field(KvFieldRef::Key { byte_offset, .. })) =
3206 shifted.reducers[0].expr.as_ref()
3207 else {
3208 panic!("expected key field reducer");
3209 };
3210 assert_eq!(*byte_offset, 12);
3213 let KvExpr::Field(KvFieldRef::ZOrderKey { bit_offset, .. }) = &shifted.group_by[0] else {
3214 panic!("expected z-order group field");
3215 };
3216 assert_eq!(*bit_offset, 36);
3219 }
3220
3221 #[test]
3222 fn store_write_batch_uses_each_clients_prefix() {
3223 let base = StoreClient::new("http://localhost:10000");
3224 let a = base.prefixed(StoreKeyPrefix::new(vec![1]).unwrap());
3225 let b = base.prefixed(StoreKeyPrefix::new(vec![2]).unwrap());
3226 let key_a = Bytes::from_static(b"a");
3227 let key_b = Bytes::from_static(b"b");
3228
3229 let mut batch = StoreWriteBatch::new();
3230 batch.push(&a, &key_a, b"va").unwrap();
3231 batch.push(&b, &key_b, b"vb").unwrap();
3232
3233 assert_eq!(
3234 batch.entries[0].0,
3235 a.key_prefix().encode_key(&key_a).unwrap()
3236 );
3237 assert_eq!(
3238 batch.entries[1].0,
3239 b.key_prefix().encode_key(&key_b).unwrap()
3240 );
3241 }
3242
3243 #[test]
3244 fn pushed_rows_round_trip_only_through_their_own_client() {
3245 let base = StoreClient::new("http://localhost:10000");
3246 let a = base.prefixed(StoreKeyPrefix::new(vec![1]).unwrap());
3247 let b = base.prefixed(StoreKeyPrefix::new(vec![2]).unwrap());
3248 let key = Bytes::from_static(b"shared-logical-key");
3249
3250 let mut batch = StoreWriteBatch::new();
3251 batch.push(&a, &key, b"va").unwrap();
3252 batch.push(&b, &key, b"vb").unwrap();
3253
3254 let entries = batch.entries();
3257 assert_eq!(entries.len(), 2);
3258 assert_ne!(entries[0].0, entries[1].0);
3259
3260 assert_eq!(a.decode_store_key(&entries[0].0).unwrap(), key);
3263 assert_eq!(b.decode_store_key(&entries[1].0).unwrap(), key);
3264 assert!(a.decode_store_key(&entries[1].0).is_err());
3265 assert!(b.decode_store_key(&entries[0].0).is_err());
3266 }
3267
3268 #[test]
3269 fn identity_prefix_stages_keys_verbatim() {
3270 let base = StoreClient::new("http://localhost:10000");
3271 let plain = PrefixedStoreClient::empty(base);
3272 let key = Bytes::from_static(b"raw-key");
3273
3274 let mut batch = StoreWriteBatch::new();
3275 batch.push(&plain, &key, b"v").unwrap();
3276
3277 assert_eq!(batch.entries()[0].0, key);
3278 }
3279
3280 #[test]
3281 fn push_rejects_keys_exceeding_prefixed_capacity() {
3282 let base = StoreClient::new("http://localhost:10000");
3283 let a = base.prefixed(StoreKeyPrefix::new(vec![1]).unwrap());
3284 let max = a.key_prefix().max_logical_key_len();
3285
3286 let mut batch = StoreWriteBatch::new();
3287 let at_capacity = Key::from(vec![7u8; max]);
3288 batch.push(&a, &at_capacity, b"v").unwrap();
3289
3290 let oversize = Key::from(vec![7u8; max + 1]);
3291 assert!(batch.push(&a, &oversize, b"v").is_err());
3292 assert_eq!(batch.len(), 1);
3293 }
3294
3295 fn hex_encode(data: &[u8]) -> String {
3296 hex::encode(data)
3297 }
3298
3299 fn hex_decode(s: &str) -> Option<Vec<u8>> {
3300 hex::decode(s).ok()
3301 }
3302}