1#![doc = include_str!("../README.md")]
25#![cfg_attr(not(test), warn(clippy::unwrap_used, clippy::expect_used))]
28
29mod attempt;
30mod budget;
31mod channel_pool;
32mod driver;
33mod driver_supervisor;
34mod error;
35mod leader_hint;
36mod response;
37mod retry;
38mod retry_policy;
39mod seq_attempt;
40mod transport;
41mod worklist;
42
43#[cfg(test)]
44mod test_support;
45
46pub use error::ClientError;
47pub use retry_policy::RetryPolicy;
48pub use seq_attempt::SeqBlock;
49pub use transport::BoxError;
50
51use std::sync::Arc;
52use std::time::Duration;
53use tsoracle_core::{Epoch, LOGICAL_MAX, Timestamp};
54
55pub(crate) const MAX_TIMESTAMPS_PER_RPC: u32 = LOGICAL_MAX + 1;
60
61use crate::channel_pool::ChannelPool;
62
63pub struct ClientBuilder {
64 endpoints: Vec<String>,
65 flush_interval: Duration,
66 connector: Option<Arc<crate::transport::ChannelConnector>>,
67 tls_required: bool,
68 retry_policy: RetryPolicy,
69}
70
71impl ClientBuilder {
72 pub fn endpoints(endpoints: Vec<String>) -> Self {
73 ClientBuilder {
74 endpoints,
75 flush_interval: Duration::from_millis(1),
76 connector: None,
77 tls_required: false,
78 retry_policy: RetryPolicy::default(),
79 }
80 }
81
82 pub fn batch_flush_interval(mut self, flush_interval: Duration) -> Self {
83 self.flush_interval = flush_interval;
84 self
85 }
86
87 pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
99 self.retry_policy = policy;
100 self
101 }
102
103 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
119 pub fn tls_config(mut self, cfg: tonic::transport::ClientTlsConfig) -> Self {
120 self.connector = Some(crate::transport::tls_connector(
121 cfg,
122 self.retry_policy.clone(),
123 ));
124 self.tls_required = true;
125 self
126 }
127
128 pub fn channel_connector<F, Fut>(mut self, connector: F) -> Self
139 where
140 F: Fn(&str) -> Fut + Send + Sync + 'static,
141 Fut: std::future::Future<Output = Result<tonic::transport::Channel, crate::BoxError>>
142 + Send
143 + 'static,
144 {
145 let wrapped: Arc<crate::transport::ChannelConnector> = Arc::new(move |endpoint: &str| {
146 let fut = connector(endpoint);
147 Box::pin(async move { fut.await.map_err(ClientError::Connector) })
148 });
149 self.connector = Some(wrapped);
150 self.tls_required = false;
151 self
152 }
153
154 pub async fn build(self) -> Result<Client, ClientError> {
155 if self.endpoints.is_empty() {
156 return Err(ClientError::NoReachableEndpoints);
157 }
158 let pool = Arc::new(ChannelPool::new(
159 self.endpoints,
160 self.connector,
161 self.tls_required,
162 self.retry_policy,
163 ));
164 let pool_for_rpc = pool.clone();
165 let driver = driver::Driver::spawn(
166 move |count| {
167 let pool = pool_for_rpc.clone();
168 Box::pin(async move { retry::issue_rpc(&pool, count).await })
169 },
170 self.flush_interval,
171 );
172 Ok(Client { pool, driver })
173 }
174}
175
176pub struct Client {
177 pool: Arc<ChannelPool>,
178 driver: driver::Driver,
179}
180
181impl Client {
182 pub async fn connect(endpoints: Vec<String>) -> Result<Self, ClientError> {
183 ClientBuilder::endpoints(endpoints).build().await
184 }
185
186 pub fn cached_leader(&self) -> Option<String> {
196 self.pool.cached_leader()
197 }
198
199 pub async fn get_ts(&self) -> Result<Timestamp, ClientError> {
200 Ok(self.driver.request(1).await?[0])
201 }
202
203 pub async fn get_ts_batch(&self, count: u32) -> Result<Vec<Timestamp>, ClientError> {
204 if count == 0 || count > MAX_TIMESTAMPS_PER_RPC {
205 return Err(ClientError::InvalidCount(count));
206 }
207 self.driver.request(count).await
208 }
209
210 pub async fn get_seq(&self, key: &str, count: u32) -> Result<SeqBlock, ClientError> {
231 if key.is_empty() || key.len() > tsoracle_core::MAX_SEQ_KEY_LEN {
232 return Err(ClientError::InvalidSeqKey);
233 }
234 if count == 0 {
235 return Err(ClientError::InvalidCount(0));
236 }
237 retry::issue_seq_rpc(&self.pool, key, count).await
238 }
239
240 pub async fn get_seq_batch(
254 &self,
255 entries: &[(&str, u32)],
256 ) -> Result<Vec<SeqBlock>, ClientError> {
257 if entries.is_empty() {
258 return Err(ClientError::InvalidCount(0));
259 }
260 let mut seen = std::collections::HashSet::with_capacity(entries.len());
261 for (key, count) in entries {
262 if key.is_empty() || key.len() > tsoracle_core::MAX_SEQ_KEY_LEN {
263 return Err(ClientError::InvalidSeqKey);
264 }
265 if *count == 0 {
266 return Err(ClientError::InvalidCount(0));
267 }
268 if !seen.insert(*key) {
269 return Err(ClientError::InvalidSeqKey);
271 }
272 }
273 retry::issue_seq_batch_rpc(&self.pool, entries).await
274 }
275
276 pub async fn get_current_max_safe(&self) -> Result<MaxSafe, ClientError> {
292 let endpoint = self
293 .pool
294 .cached_leader()
295 .or_else(|| self.pool.iter_round_robin().into_iter().next())
296 .ok_or(ClientError::NoReachableEndpoints)?;
297 let budget = self.pool.retry_policy().per_attempt_deadline;
298 let pair = crate::budget::PairBudget::start(budget);
302 let (mut svc, cell) =
303 match tokio::time::timeout(budget, self.pool.client_with_cell(&endpoint)).await {
304 Ok(Ok(leased)) => leased,
305 Ok(Err(err)) => return Err(err),
308 Err(_) => {
309 return Err(ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
310 "connect exceeded per_attempt_deadline of {budget:?}"
311 ))));
312 }
313 };
314 let rpc_budget = pair.remaining();
315 let rpc = svc.get_current_max_safe(tsoracle_proto::v1::GetCurrentMaxSafeRequest {});
316 let err = match tokio::time::timeout(rpc_budget, rpc).await {
317 Ok(Ok(response)) => {
318 let inner = response.into_inner();
319 return Ok(MaxSafe {
320 max_safe_physical_ms: inner.max_safe_physical_ms,
321 epoch: Epoch::from_wire(inner.epoch_hi, inner.epoch_lo),
322 });
323 }
324 Ok(Err(status)) => ClientError::Rpc(status),
325 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
329 "rpc exceeded its share of per_attempt_deadline \
330 ({rpc_budget:?} of {budget:?})"
331 ))),
332 };
333 if crate::retry_policy::is_transport_failure(&err) {
334 self.pool.evict_if_current(&endpoint, &cell);
335 }
336 Err(err)
337 }
338}
339
340#[derive(Copy, Clone, Debug, PartialEq, Eq)]
343pub struct MaxSafe {
344 pub max_safe_physical_ms: u64,
347 pub epoch: Epoch,
349}
350
351#[cfg(test)]
352mod tests {
353 use super::*;
354
355 #[tokio::test]
356 async fn cached_leader_is_none_before_any_rpc() {
357 let client = Client::connect(vec!["http://127.0.0.1:1".into()])
363 .await
364 .expect("build with a non-empty endpoint list must succeed");
365 assert_eq!(client.cached_leader(), None);
366 }
367
368 #[tokio::test]
369 async fn build_rejects_empty_endpoint_list() {
370 match ClientBuilder::endpoints(Vec::new()).build().await {
374 Err(ClientError::NoReachableEndpoints) => {}
375 Err(other) => panic!("expected NoReachableEndpoints, got {other:?}"),
376 Ok(_) => panic!("expected Err, got Ok(Client)"),
377 }
378 }
379
380 #[tokio::test]
381 async fn channel_connector_error_surfaces_as_connector_variant() {
382 let builder = ClientBuilder::endpoints(vec!["a:1".into()]).channel_connector(
383 |_endpoint: &str| async move {
384 Err::<tonic::transport::Channel, crate::BoxError>(
385 std::io::Error::other("boom").into(),
386 )
387 },
388 );
389 let client = builder.build().await.expect("build must not fail");
390 let result = client.get_ts().await;
391 match result {
392 Err(ClientError::Connector(inner)) => {
393 assert!(inner.to_string().contains("boom"));
394 }
395 other => panic!("expected ClientError::Connector, got {other:?}"),
396 }
397 }
398
399 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
404 async fn marker_connector_failure() -> Result<tonic::transport::Channel, crate::BoxError> {
405 Err("MARKER".into())
406 }
407
408 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
409 #[tokio::test]
410 async fn tls_config_then_channel_connector_last_wins() {
411 let builder = ClientBuilder::endpoints(vec!["a:1".into()])
415 .tls_config(tonic::transport::ClientTlsConfig::new())
416 .channel_connector(|_endpoint: &str| marker_connector_failure());
417 let client = builder.build().await.expect("build must not fail");
418 match client.get_ts().await {
419 Err(ClientError::Connector(inner)) => {
420 assert!(inner.to_string().contains("MARKER"));
421 }
422 other => panic!("expected Connector(MARKER), got {other:?}"),
423 }
424 }
425
426 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
427 #[tokio::test]
428 async fn channel_connector_then_tls_config_last_wins() {
429 let builder = ClientBuilder::endpoints(vec!["a:1".into()])
433 .channel_connector(|_endpoint: &str| marker_connector_failure())
434 .tls_config(tonic::transport::ClientTlsConfig::new());
435 let client = builder.build().await.expect("build must not fail");
436 let result = client.get_ts().await;
437 if let Err(ClientError::Connector(inner)) = &result
438 && inner.to_string().contains("MARKER")
439 {
440 panic!("tls_config set last must overwrite the prior channel_connector");
441 }
442 }
443
444 #[tokio::test]
445 async fn batch_flush_interval_overrides_default() {
446 let custom = Duration::from_millis(25);
451 let builder = ClientBuilder::endpoints(vec!["http://127.0.0.1:1".into()])
452 .batch_flush_interval(custom);
453 assert_eq!(builder.flush_interval, custom);
454 }
455
456 #[tokio::test]
457 async fn retry_policy_override_propagates_to_builder() {
458 let policy = RetryPolicy {
463 max_attempts: 7,
464 per_attempt_deadline: Duration::from_millis(11),
465 overall_deadline: Duration::from_millis(13),
466 base_backoff: Duration::from_millis(17),
467 leader_ttl: Duration::from_millis(19),
468 };
469 let builder = ClientBuilder::endpoints(vec!["http://127.0.0.1:1".into()])
470 .retry_policy(policy.clone());
471 assert_eq!(builder.retry_policy.max_attempts, policy.max_attempts);
472 assert_eq!(
473 builder.retry_policy.per_attempt_deadline,
474 policy.per_attempt_deadline
475 );
476 assert_eq!(
477 builder.retry_policy.overall_deadline,
478 policy.overall_deadline
479 );
480 assert_eq!(builder.retry_policy.base_backoff, policy.base_backoff);
481 assert_eq!(builder.retry_policy.leader_ttl, policy.leader_ttl);
482 }
483
484 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
497 async fn get_current_max_safe_returns_within_per_attempt_deadline_when_connector_hangs() {
498 let policy = RetryPolicy {
499 max_attempts: 1,
500 per_attempt_deadline: Duration::from_millis(100),
501 overall_deadline: Duration::from_millis(300),
502 base_backoff: Duration::ZERO,
503 leader_ttl: Duration::from_secs(30),
504 };
505 let client = ClientBuilder::endpoints(vec!["hang:1".into()])
506 .channel_connector(|_endpoint: &str| async {
507 std::future::pending::<Result<tonic::transport::Channel, crate::BoxError>>().await
513 })
514 .retry_policy(policy)
515 .build()
516 .await
517 .expect("builder must accept the policy");
518 let outer_safety = Duration::from_secs(5);
525 let start = std::time::Instant::now();
526 let result = match tokio::time::timeout(outer_safety, client.get_current_max_safe()).await {
527 Ok(r) => r,
528 Err(_) => panic!(
529 "get_current_max_safe failed to honor its own per_attempt_deadline; \
530 the {outer_safety:?} outer safety net had to fire — this is the \
531 security finding's exact symptom (channel acquisition or RPC was \
532 never bounded)",
533 ),
534 };
535 let elapsed = start.elapsed();
536 assert!(
537 result.is_err(),
538 "a hanging connector must surface as Err, got {result:?}",
539 );
540 assert!(
541 elapsed < Duration::from_secs(2),
542 "deadline must short-circuit; took {elapsed:?} (per_attempt_deadline was 100ms)",
543 );
544 }
545
546 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
547 async fn get_ts_returns_within_overall_deadline_when_all_endpoints_unreachable() {
548 let policy = RetryPolicy {
555 max_attempts: 3,
556 per_attempt_deadline: Duration::from_millis(100),
557 overall_deadline: Duration::from_millis(300),
558 base_backoff: Duration::ZERO,
559 leader_ttl: Duration::from_secs(30),
560 };
561 let client = ClientBuilder::endpoints(vec![
562 "http://127.0.0.1:1".into(),
563 "http://127.0.0.1:2".into(),
564 "http://127.0.0.1:3".into(),
565 ])
566 .retry_policy(policy)
567 .build()
568 .await
569 .expect("builder must accept the policy");
570 let start = std::time::Instant::now();
571 let result = client.get_ts().await;
572 let elapsed = start.elapsed();
573 assert!(result.is_err(), "no listener can reply: {result:?}");
574 assert!(
575 elapsed < Duration::from_secs(2),
576 "deadline must short-circuit; took {elapsed:?}"
577 );
578 }
579
580 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
586 async fn get_seq_returns_gapless_blocks_and_rejects_invalid_inputs() {
587 use std::sync::Arc;
588 use std::sync::atomic::{AtomicU64, Ordering};
589 const EPOCH: u128 = 42;
590
591 let counter = Arc::new(AtomicU64::new(0));
592 let server_counter = counter.clone();
593 let addr = crate::test_support::FakeTso::new()
596 .on_get_seq(move |req| {
597 let counter = server_counter.clone();
598 async move {
599 if req.key.is_empty() {
600 return Err(tonic::Status::invalid_argument("invalid sequence key"));
601 }
602 if req.count == 0 {
603 return Err(tonic::Status::invalid_argument(
604 "count must be between 1 and the maximum",
605 ));
606 }
607 let start = counter.fetch_add(u64::from(req.count), Ordering::SeqCst);
608 let (hi, lo) = tsoracle_core::Epoch(EPOCH).to_wire();
609 Ok(tsoracle_proto::v1::GetSeqResponse {
610 key: req.key,
611 start,
612 count: req.count,
613 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
614 })
615 }
616 })
617 .spawn()
618 .await;
619
620 let endpoint = format!("http://{addr}");
621 let client = ClientBuilder::endpoints(vec![endpoint])
622 .retry_policy(RetryPolicy {
623 max_attempts: 3,
624 per_attempt_deadline: Duration::from_secs(2),
625 overall_deadline: Duration::from_secs(5),
626 base_backoff: Duration::ZERO,
627 leader_ttl: Duration::from_secs(30),
628 })
629 .build()
630 .await
631 .expect("client must build");
632
633 let block1 = client
635 .get_seq("orders", 5)
636 .await
637 .expect("first get_seq must succeed");
638 assert_eq!(block1.start, 0, "first block must start at 0");
639 assert_eq!(block1.count, 5, "first block must have count 5");
640 assert_eq!(block1.epoch, EPOCH, "epoch must match the server's epoch");
641
642 let block2 = client
644 .get_seq("orders", 3)
645 .await
646 .expect("second get_seq must succeed");
647 assert_eq!(block2.start, 5, "second block must start at 5 (gapless)");
648 assert_eq!(block2.count, 3);
649 assert_eq!(block2.epoch, EPOCH);
650
651 match client.get_seq("", 1).await {
653 Err(ClientError::InvalidSeqKey) => {}
654 other => panic!("empty key must return InvalidSeqKey, got {other:?}"),
655 }
656
657 match client.get_seq("orders", 0).await {
659 Err(ClientError::InvalidCount(0)) => {}
660 other => panic!("count=0 must return InvalidCount, got {other:?}"),
661 }
662 }
663
664 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
670 async fn get_seq_forwards_large_count_to_server() {
671 let addr = crate::test_support::FakeTso::new()
672 .on_get_seq(|req| async move {
673 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
674 Ok(tsoracle_proto::v1::GetSeqResponse {
675 key: req.key,
676 start: 0,
677 count: req.count,
678 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
679 })
680 })
681 .spawn()
682 .await;
683
684 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
685 .retry_policy(RetryPolicy {
686 max_attempts: 3,
687 per_attempt_deadline: Duration::from_secs(2),
688 overall_deadline: Duration::from_secs(5),
689 base_backoff: Duration::ZERO,
690 leader_ttl: Duration::from_secs(30),
691 })
692 .build()
693 .await
694 .unwrap();
695
696 let big = tsoracle_core::DEFAULT_MAX_SEQ_COUNT + 1;
699 let block = client
700 .get_seq("orders", big)
701 .await
702 .expect("a large count must be forwarded to the server, not locally rejected");
703 assert_eq!(block.count, big);
704 }
705
706 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
715 async fn get_seq_over_cap_count_surfaces_rpc_invalid_argument() {
716 use std::sync::Arc;
717 use std::sync::atomic::{AtomicU64, Ordering};
718 let calls = Arc::new(AtomicU64::new(0));
719 let server_calls = calls.clone();
720 let addr = crate::test_support::FakeTso::new()
724 .on_get_seq(move |_req| {
725 let calls = server_calls.clone();
726 async move {
727 calls.fetch_add(1, Ordering::SeqCst);
728 Err(tonic::Status::invalid_argument(
729 "count must be between 1 and the maximum",
730 ))
731 }
732 })
733 .spawn()
734 .await;
735
736 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
737 .retry_policy(RetryPolicy {
738 max_attempts: 3,
739 per_attempt_deadline: Duration::from_secs(2),
740 overall_deadline: Duration::from_secs(5),
741 base_backoff: Duration::ZERO,
742 leader_ttl: Duration::from_secs(30),
743 })
744 .build()
745 .await
746 .unwrap();
747
748 let big = tsoracle_core::DEFAULT_MAX_SEQ_COUNT + 1;
751 match client.get_seq("orders", big).await {
752 Err(ClientError::Rpc(status)) => {
753 assert_eq!(status.code(), tonic::Code::InvalidArgument);
754 assert!(
755 status.message().contains("maximum"),
756 "must surface the server's over-cap message, not a synthetic \
757 or mislabelled error; got {:?}",
758 status.message()
759 );
760 }
761 other => panic!(
762 "an over-cap count with a valid key must surface as \
763 ClientError::Rpc(InvalidArgument), not InvalidSeqKey; got {other:?}"
764 ),
765 }
766 let n = calls.load(Ordering::SeqCst);
768 assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
769 }
770
771 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
777 async fn get_seq_bare_failed_precondition_surfaces_definitive_error() {
778 use std::sync::Arc;
779 use std::sync::atomic::{AtomicU64, Ordering};
780 let calls = Arc::new(AtomicU64::new(0));
781 let server_calls = calls.clone();
782 let addr = crate::test_support::FakeTso::new()
785 .on_get_seq(move |_req| {
786 let calls = server_calls.clone();
787 async move {
788 calls.fetch_add(1, Ordering::SeqCst);
789 Err(tonic::Status::failed_precondition("dense counter overflow"))
790 }
791 })
792 .spawn()
793 .await;
794
795 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
796 .retry_policy(RetryPolicy {
797 max_attempts: 3,
798 per_attempt_deadline: Duration::from_secs(2),
799 overall_deadline: Duration::from_secs(5),
800 base_backoff: Duration::ZERO,
801 leader_ttl: Duration::from_secs(30),
802 })
803 .build()
804 .await
805 .unwrap();
806
807 match client.get_seq("orders", 1).await {
808 Err(ClientError::Rpc(status)) => {
809 assert_eq!(status.code(), tonic::Code::FailedPrecondition);
810 assert!(
811 status.message().contains("overflow"),
812 "must surface the server's overflow message, not a synthetic \
813 'no leader yet'; got {:?}",
814 status.message()
815 );
816 }
817 other => panic!("expected the overflow FailedPrecondition, got {other:?}"),
818 }
819 let n = calls.load(Ordering::SeqCst);
821 assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
822 }
823
824 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
830 async fn get_seq_ambiguous_post_send_failure_is_uncertain_without_retry() {
831 use std::sync::Arc;
832 use std::sync::atomic::{AtomicU64, Ordering};
833 let calls = Arc::new(AtomicU64::new(0));
834 let server_calls = calls.clone();
835 let addr = crate::test_support::FakeTso::new()
838 .on_get_seq(move |_req| {
839 let calls = server_calls.clone();
840 async move {
841 calls.fetch_add(1, Ordering::SeqCst);
842 Err(tonic::Status::unavailable("connection reset mid-call"))
843 }
844 })
845 .spawn()
846 .await;
847
848 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
849 .retry_policy(RetryPolicy {
850 max_attempts: 3,
851 per_attempt_deadline: Duration::from_secs(2),
852 overall_deadline: Duration::from_secs(5),
853 base_backoff: Duration::ZERO,
854 leader_ttl: Duration::from_secs(30),
855 })
856 .build()
857 .await
858 .unwrap();
859
860 match client.get_seq("orders", 1).await {
861 Err(ClientError::SeqUncertain) => {}
862 other => panic!("post-send Unavailable must be SeqUncertain, got {other:?}"),
863 }
864 assert_eq!(
865 calls.load(Ordering::SeqCst),
866 1,
867 "non-idempotent get_seq must NOT retry an ambiguous post-send failure"
868 );
869 }
870
871 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
878 async fn get_seq_post_send_internal_is_uncertain_without_retry() {
879 use std::sync::Arc;
880 use std::sync::atomic::{AtomicU64, Ordering};
881 let calls = Arc::new(AtomicU64::new(0));
882 let server_calls = calls.clone();
883 let addr = crate::test_support::FakeTso::new()
886 .on_get_seq(move |_req| {
887 let calls = server_calls.clone();
888 async move {
889 calls.fetch_add(1, Ordering::SeqCst);
890 Err(tonic::Status::internal("permanent dense driver fault"))
891 }
892 })
893 .spawn()
894 .await;
895
896 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
897 .retry_policy(RetryPolicy {
898 max_attempts: 3,
899 per_attempt_deadline: Duration::from_secs(2),
900 overall_deadline: Duration::from_secs(5),
901 base_backoff: Duration::ZERO,
902 leader_ttl: Duration::from_secs(30),
903 })
904 .build()
905 .await
906 .unwrap();
907
908 match client.get_seq("orders", 1).await {
909 Err(ClientError::SeqUncertain) => {}
910 other => panic!("post-send INTERNAL must be SeqUncertain, got {other:?}"),
911 }
912 assert_eq!(
913 calls.load(Ordering::SeqCst),
914 1,
915 "post-send INTERNAL must NOT be retried (possible committed advance)"
916 );
917 }
918
919 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
925 async fn get_seq_response_count_mismatch_is_uncertain() {
926 use std::sync::Arc;
927 use std::sync::atomic::{AtomicU64, Ordering};
928 let calls = Arc::new(AtomicU64::new(0));
929 let server_calls = calls.clone();
930 let addr = crate::test_support::FakeTso::new()
933 .on_get_seq(move |req| {
934 let calls = server_calls.clone();
935 async move {
936 calls.fetch_add(1, Ordering::SeqCst);
937 let (hi, lo) = tsoracle_core::Epoch(7).to_wire();
938 Ok(tsoracle_proto::v1::GetSeqResponse {
939 key: req.key,
940 start: 0,
941 count: req.count + 1,
942 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
943 })
944 }
945 })
946 .spawn()
947 .await;
948
949 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
950 .retry_policy(RetryPolicy {
951 max_attempts: 3,
952 per_attempt_deadline: Duration::from_secs(2),
953 overall_deadline: Duration::from_secs(5),
954 base_backoff: Duration::ZERO,
955 leader_ttl: Duration::from_secs(30),
956 })
957 .build()
958 .await
959 .unwrap();
960
961 match client.get_seq("orders", 5).await {
962 Err(ClientError::SeqUncertain) => {}
963 other => panic!("count-mismatch success must be SeqUncertain, got {other:?}"),
964 }
965 assert_eq!(
966 calls.load(Ordering::SeqCst),
967 1,
968 "a malformed success must not trigger a (double-spending) retry"
969 );
970 }
971
972 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
975 async fn get_seq_response_key_mismatch_is_uncertain() {
976 use std::sync::Arc;
977 use std::sync::atomic::{AtomicU64, Ordering};
978 let calls = Arc::new(AtomicU64::new(0));
979 let server_calls = calls.clone();
980 let addr = crate::test_support::FakeTso::new()
983 .on_get_seq(move |req| {
984 let calls = server_calls.clone();
985 async move {
986 calls.fetch_add(1, Ordering::SeqCst);
987 let (hi, lo) = tsoracle_core::Epoch(7).to_wire();
988 Ok(tsoracle_proto::v1::GetSeqResponse {
989 key: format!("{}-tampered", req.key),
990 start: 0,
991 count: req.count,
992 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
993 })
994 }
995 })
996 .spawn()
997 .await;
998
999 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1000 .retry_policy(RetryPolicy {
1001 max_attempts: 3,
1002 per_attempt_deadline: Duration::from_secs(2),
1003 overall_deadline: Duration::from_secs(5),
1004 base_backoff: Duration::ZERO,
1005 leader_ttl: Duration::from_secs(30),
1006 })
1007 .build()
1008 .await
1009 .unwrap();
1010
1011 match client.get_seq("orders", 5).await {
1012 Err(ClientError::SeqUncertain) => {}
1013 other => panic!("key-mismatch success must be SeqUncertain, got {other:?}"),
1014 }
1015 assert_eq!(
1016 calls.load(Ordering::SeqCst),
1017 1,
1018 "malformed success: one call"
1019 );
1020 }
1021
1022 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1026 async fn get_seq_batch_returns_blocks_in_request_order() {
1027 use std::sync::Arc;
1028 use std::sync::atomic::{AtomicU64, Ordering};
1029 const EPOCH: u128 = 77;
1030
1031 let counter = Arc::new(AtomicU64::new(0));
1034 let server_counter = counter.clone();
1035 let addr = crate::test_support::FakeTso::new()
1036 .on_get_seq_batch(move |req| {
1037 let counter = server_counter.clone();
1038 async move {
1039 let mut grants = Vec::with_capacity(req.entries.len());
1040 for entry in &req.entries {
1041 let start = counter.fetch_add(u64::from(entry.count), Ordering::SeqCst);
1042 grants.push(tsoracle_proto::v1::SeqGrantEntry {
1043 key: entry.key.clone(),
1044 start,
1045 count: entry.count,
1046 });
1047 }
1048 let (hi, lo) = tsoracle_core::Epoch(EPOCH).to_wire();
1049 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1050 grants,
1051 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1052 })
1053 }
1054 })
1055 .spawn()
1056 .await;
1057
1058 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1059 .retry_policy(RetryPolicy {
1060 max_attempts: 3,
1061 per_attempt_deadline: Duration::from_secs(2),
1062 overall_deadline: Duration::from_secs(5),
1063 base_backoff: Duration::ZERO,
1064 leader_ttl: Duration::from_secs(30),
1065 })
1066 .build()
1067 .await
1068 .expect("client must build");
1069
1070 let blocks = client
1071 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1072 .await
1073 .expect("get_seq_batch must succeed");
1074
1075 assert_eq!(blocks.len(), 2, "one block per entry");
1076 assert_eq!(blocks[0].start, 0);
1078 assert_eq!(blocks[0].count, 5);
1079 assert_eq!(blocks[0].epoch, EPOCH);
1080 assert_eq!(blocks[1].start, 5);
1082 assert_eq!(blocks[1].count, 3);
1083 assert_eq!(blocks[1].epoch, EPOCH);
1084 }
1085
1086 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1088 async fn get_seq_batch_rejects_duplicate_keys_without_hitting_server() {
1089 use std::sync::Arc;
1090 use std::sync::atomic::{AtomicU64, Ordering};
1091
1092 let calls = Arc::new(AtomicU64::new(0));
1093 let server_calls = calls.clone();
1094 let addr = crate::test_support::FakeTso::new()
1096 .on_get_seq_batch(move |_req| {
1097 let calls = server_calls.clone();
1098 async move {
1099 calls.fetch_add(1, Ordering::SeqCst);
1100 Err(tonic::Status::internal("should not be called"))
1101 }
1102 })
1103 .spawn()
1104 .await;
1105
1106 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1107 .retry_policy(RetryPolicy {
1108 max_attempts: 3,
1109 per_attempt_deadline: Duration::from_secs(2),
1110 overall_deadline: Duration::from_secs(5),
1111 base_backoff: Duration::ZERO,
1112 leader_ttl: Duration::from_secs(30),
1113 })
1114 .build()
1115 .await
1116 .expect("client must build");
1117
1118 match client.get_seq_batch(&[("orders", 1), ("orders", 2)]).await {
1120 Err(ClientError::InvalidSeqKey) => {}
1121 other => panic!("duplicate key must return InvalidSeqKey, got {other:?}"),
1122 }
1123 assert_eq!(
1124 calls.load(Ordering::SeqCst),
1125 0,
1126 "duplicate-key rejection must not reach the server"
1127 );
1128 }
1129
1130 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1134 async fn get_seq_batch_malformed_success_is_uncertain() {
1135 use std::sync::Arc;
1136 use std::sync::atomic::{AtomicU64, Ordering};
1137
1138 {
1140 let calls = Arc::new(AtomicU64::new(0));
1141 let server_calls = calls.clone();
1142 let addr = crate::test_support::FakeTso::new()
1143 .on_get_seq_batch(move |_req| {
1144 let calls = server_calls.clone();
1145 async move {
1146 calls.fetch_add(1, Ordering::SeqCst);
1147 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1148 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1150 grants: vec![tsoracle_proto::v1::SeqGrantEntry {
1151 key: "orders".into(),
1152 start: 0,
1153 count: 5,
1154 }],
1155 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1156 })
1157 }
1158 })
1159 .spawn()
1160 .await;
1161
1162 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1163 .retry_policy(RetryPolicy {
1164 max_attempts: 3,
1165 per_attempt_deadline: Duration::from_secs(2),
1166 overall_deadline: Duration::from_secs(5),
1167 base_backoff: Duration::ZERO,
1168 leader_ttl: Duration::from_secs(30),
1169 })
1170 .build()
1171 .await
1172 .expect("client must build");
1173
1174 match client
1175 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1176 .await
1177 {
1178 Err(ClientError::SeqUncertain) => {}
1179 other => panic!("wrong grants.len must surface as SeqUncertain, got {other:?}"),
1180 }
1181 assert_eq!(
1182 calls.load(Ordering::SeqCst),
1183 1,
1184 "malformed success must not trigger a retry"
1185 );
1186 }
1187
1188 {
1190 let calls = Arc::new(AtomicU64::new(0));
1191 let server_calls = calls.clone();
1192 let addr = crate::test_support::FakeTso::new()
1193 .on_get_seq_batch(move |req| {
1194 let calls = server_calls.clone();
1195 async move {
1196 calls.fetch_add(1, Ordering::SeqCst);
1197 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1198 let grants = req
1200 .entries
1201 .iter()
1202 .enumerate()
1203 .map(|(i, e)| tsoracle_proto::v1::SeqGrantEntry {
1204 key: e.key.clone(),
1205 start: 0,
1206 count: if i == 0 { e.count + 1 } else { e.count },
1208 })
1209 .collect();
1210 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1211 grants,
1212 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1213 })
1214 }
1215 })
1216 .spawn()
1217 .await;
1218
1219 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1220 .retry_policy(RetryPolicy {
1221 max_attempts: 3,
1222 per_attempt_deadline: Duration::from_secs(2),
1223 overall_deadline: Duration::from_secs(5),
1224 base_backoff: Duration::ZERO,
1225 leader_ttl: Duration::from_secs(30),
1226 })
1227 .build()
1228 .await
1229 .expect("client must build");
1230
1231 match client
1232 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1233 .await
1234 {
1235 Err(ClientError::SeqUncertain) => {}
1236 other => panic!("count-mismatch grant must surface as SeqUncertain, got {other:?}"),
1237 }
1238 assert_eq!(
1239 calls.load(Ordering::SeqCst),
1240 1,
1241 "malformed success must not trigger a retry"
1242 );
1243 }
1244
1245 {
1247 let calls = Arc::new(AtomicU64::new(0));
1248 let server_calls = calls.clone();
1249 let addr = crate::test_support::FakeTso::new()
1250 .on_get_seq_batch(move |req| {
1251 let calls = server_calls.clone();
1252 async move {
1253 calls.fetch_add(1, Ordering::SeqCst);
1254 let grants = req
1255 .entries
1256 .iter()
1257 .map(|e| tsoracle_proto::v1::SeqGrantEntry {
1258 key: e.key.clone(),
1259 start: 0,
1260 count: e.count,
1261 })
1262 .collect();
1263 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1265 grants,
1266 epoch: None,
1267 })
1268 }
1269 })
1270 .spawn()
1271 .await;
1272
1273 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1274 .retry_policy(RetryPolicy {
1275 max_attempts: 3,
1276 per_attempt_deadline: Duration::from_secs(2),
1277 overall_deadline: Duration::from_secs(5),
1278 base_backoff: Duration::ZERO,
1279 leader_ttl: Duration::from_secs(30),
1280 })
1281 .build()
1282 .await
1283 .expect("client must build");
1284
1285 match client
1286 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1287 .await
1288 {
1289 Err(ClientError::SeqUncertain) => {}
1290 other => panic!("missing epoch must surface as SeqUncertain, got {other:?}"),
1291 }
1292 assert_eq!(
1293 calls.load(Ordering::SeqCst),
1294 1,
1295 "malformed success must not trigger a retry"
1296 );
1297 }
1298
1299 {
1302 let calls = Arc::new(AtomicU64::new(0));
1303 let server_calls = calls.clone();
1304 let addr = crate::test_support::FakeTso::new()
1305 .on_get_seq_batch(move |req| {
1306 let calls = server_calls.clone();
1307 async move {
1308 calls.fetch_add(1, Ordering::SeqCst);
1309 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1310 let grants = req
1312 .entries
1313 .iter()
1314 .enumerate()
1315 .map(|(i, e)| tsoracle_proto::v1::SeqGrantEntry {
1316 key: if i == 0 {
1317 format!("{}-tampered", e.key)
1318 } else {
1319 e.key.clone()
1320 },
1321 start: 0,
1322 count: e.count,
1323 })
1324 .collect();
1325 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1326 grants,
1327 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1328 })
1329 }
1330 })
1331 .spawn()
1332 .await;
1333
1334 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1335 .retry_policy(RetryPolicy {
1336 max_attempts: 3,
1337 per_attempt_deadline: Duration::from_secs(2),
1338 overall_deadline: Duration::from_secs(5),
1339 base_backoff: Duration::ZERO,
1340 leader_ttl: Duration::from_secs(30),
1341 })
1342 .build()
1343 .await
1344 .expect("client must build");
1345
1346 match client
1347 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1348 .await
1349 {
1350 Err(ClientError::SeqUncertain) => {}
1351 other => panic!("key-mismatch grant must surface as SeqUncertain, got {other:?}"),
1352 }
1353 assert_eq!(
1354 calls.load(Ordering::SeqCst),
1355 1,
1356 "malformed success must not trigger a retry"
1357 );
1358 }
1359 }
1360
1361 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1364 async fn get_seq_batch_post_send_error_is_uncertain_without_retry() {
1365 use std::sync::Arc;
1366 use std::sync::atomic::{AtomicU64, Ordering};
1367
1368 let calls = Arc::new(AtomicU64::new(0));
1369 let server_calls = calls.clone();
1370 let addr = crate::test_support::FakeTso::new()
1371 .on_get_seq_batch(move |_req| {
1372 let calls = server_calls.clone();
1373 async move {
1374 calls.fetch_add(1, Ordering::SeqCst);
1375 Err(tonic::Status::internal("permanent dense driver fault"))
1376 }
1377 })
1378 .spawn()
1379 .await;
1380
1381 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1382 .retry_policy(RetryPolicy {
1383 max_attempts: 3,
1384 per_attempt_deadline: Duration::from_secs(2),
1385 overall_deadline: Duration::from_secs(5),
1386 base_backoff: Duration::ZERO,
1387 leader_ttl: Duration::from_secs(30),
1388 })
1389 .build()
1390 .await
1391 .expect("client must build");
1392
1393 match client
1394 .get_seq_batch(&[("orders", 1), ("invoices", 2)])
1395 .await
1396 {
1397 Err(ClientError::SeqUncertain) => {}
1398 other => panic!("post-send INTERNAL must be SeqUncertain, got {other:?}"),
1399 }
1400 assert_eq!(
1401 calls.load(Ordering::SeqCst),
1402 1,
1403 "post-send INTERNAL must NOT be retried (possible committed advance)"
1404 );
1405 }
1406}