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 tokio::sync::OnceCell;
54use tonic::transport::Channel;
55use tsoracle_core::{Epoch, LOGICAL_MAX, Timestamp};
56use tsoracle_proto::v1::tso_service_client::TsoServiceClient;
57
58pub(crate) const MAX_TIMESTAMPS_PER_RPC: u32 = LOGICAL_MAX + 1;
63
64use crate::channel_pool::ChannelPool;
65
66type ControlPlaneClient = (
67 String,
68 Duration,
69 crate::budget::PairBudget,
70 TsoServiceClient<Channel>,
71 Arc<OnceCell<Channel>>,
72);
73
74pub struct ClientBuilder {
75 endpoints: Vec<String>,
76 flush_interval: Duration,
77 connector: Option<Arc<crate::transport::ChannelConnector>>,
78 tls_required: bool,
79 retry_policy: RetryPolicy,
80}
81
82impl ClientBuilder {
83 pub fn endpoints(endpoints: Vec<String>) -> Self {
84 ClientBuilder {
85 endpoints,
86 flush_interval: Duration::from_millis(1),
87 connector: None,
88 tls_required: false,
89 retry_policy: RetryPolicy::default(),
90 }
91 }
92
93 pub fn batch_flush_interval(mut self, flush_interval: Duration) -> Self {
94 self.flush_interval = flush_interval;
95 self
96 }
97
98 pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
110 self.retry_policy = policy;
111 self
112 }
113
114 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
130 pub fn tls_config(mut self, cfg: tonic::transport::ClientTlsConfig) -> Self {
131 self.connector = Some(crate::transport::tls_connector(
132 cfg,
133 self.retry_policy.clone(),
134 ));
135 self.tls_required = true;
136 self
137 }
138
139 pub fn channel_connector<F, Fut>(mut self, connector: F) -> Self
150 where
151 F: Fn(&str) -> Fut + Send + Sync + 'static,
152 Fut: std::future::Future<Output = Result<tonic::transport::Channel, crate::BoxError>>
153 + Send
154 + 'static,
155 {
156 let wrapped: Arc<crate::transport::ChannelConnector> = Arc::new(move |endpoint: &str| {
157 let fut = connector(endpoint);
158 Box::pin(async move { fut.await.map_err(ClientError::Connector) })
159 });
160 self.connector = Some(wrapped);
161 self.tls_required = false;
162 self
163 }
164
165 pub async fn build(self) -> Result<Client, ClientError> {
166 if self.endpoints.is_empty() {
167 return Err(ClientError::NoReachableEndpoints);
168 }
169 let pool = Arc::new(ChannelPool::new(
170 self.endpoints,
171 self.connector,
172 self.tls_required,
173 self.retry_policy,
174 ));
175 let pool_for_rpc = pool.clone();
176 let driver = driver::Driver::spawn(
177 move |count| {
178 let pool = pool_for_rpc.clone();
179 Box::pin(async move { retry::issue_rpc(&pool, count).await })
180 },
181 self.flush_interval,
182 );
183 Ok(Client { pool, driver })
184 }
185}
186
187pub struct Client {
188 pool: Arc<ChannelPool>,
189 driver: driver::Driver,
190}
191
192impl Client {
193 pub async fn connect(endpoints: Vec<String>) -> Result<Self, ClientError> {
194 ClientBuilder::endpoints(endpoints).build().await
195 }
196
197 pub fn cached_leader(&self) -> Option<String> {
207 self.pool.cached_leader()
208 }
209
210 pub async fn get_ts(&self) -> Result<Timestamp, ClientError> {
211 Ok(self.driver.request(1).await?[0])
212 }
213
214 pub async fn get_ts_batch(&self, count: u32) -> Result<Vec<Timestamp>, ClientError> {
215 if count == 0 || count > MAX_TIMESTAMPS_PER_RPC {
216 return Err(ClientError::InvalidCount(count));
217 }
218 self.driver.request(count).await
219 }
220
221 pub async fn get_seq(&self, key: &str, count: u32) -> Result<SeqBlock, ClientError> {
242 if key.is_empty() || key.len() > tsoracle_core::MAX_SEQ_KEY_LEN {
243 return Err(ClientError::InvalidSeqKey);
244 }
245 if count == 0 {
246 return Err(ClientError::InvalidCount(0));
247 }
248 retry::issue_seq_rpc(&self.pool, key, count).await
249 }
250
251 pub async fn get_seq_batch(
265 &self,
266 entries: &[(&str, u32)],
267 ) -> Result<Vec<SeqBlock>, ClientError> {
268 if entries.is_empty() {
269 return Err(ClientError::InvalidCount(0));
270 }
271 let mut seen = std::collections::HashSet::with_capacity(entries.len());
272 for (key, count) in entries {
273 if key.is_empty() || key.len() > tsoracle_core::MAX_SEQ_KEY_LEN {
274 return Err(ClientError::InvalidSeqKey);
275 }
276 if *count == 0 {
277 return Err(ClientError::InvalidCount(0));
278 }
279 if !seen.insert(*key) {
280 return Err(ClientError::InvalidSeqKey);
282 }
283 }
284 retry::issue_seq_batch_rpc(&self.pool, entries).await
285 }
286
287 async fn control_plane_client(&self) -> Result<ControlPlaneClient, ClientError> {
288 let endpoint = self
289 .pool
290 .cached_leader()
291 .or_else(|| self.pool.iter_round_robin().into_iter().next())
292 .ok_or(ClientError::NoReachableEndpoints)?;
293 let budget = self.pool.retry_policy().per_attempt_deadline;
294 let pair = crate::budget::PairBudget::start(budget);
295 let (svc, cell) =
296 match tokio::time::timeout(budget, self.pool.client_with_cell(&endpoint)).await {
297 Ok(Ok(leased)) => leased,
298 Ok(Err(err)) => return Err(err),
299 Err(_) => {
300 return Err(ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
301 "connect exceeded per_attempt_deadline of {budget:?}"
302 ))));
303 }
304 };
305 Ok((endpoint, budget, pair, svc, cell))
306 }
307
308 pub async fn get_current_max_safe(&self) -> Result<MaxSafe, ClientError> {
324 let endpoint = self
325 .pool
326 .cached_leader()
327 .or_else(|| self.pool.iter_round_robin().into_iter().next())
328 .ok_or(ClientError::NoReachableEndpoints)?;
329 let budget = self.pool.retry_policy().per_attempt_deadline;
330 let pair = crate::budget::PairBudget::start(budget);
334 let (mut svc, cell) =
335 match tokio::time::timeout(budget, self.pool.client_with_cell(&endpoint)).await {
336 Ok(Ok(leased)) => leased,
337 Ok(Err(err)) => return Err(err),
340 Err(_) => {
341 return Err(ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
342 "connect exceeded per_attempt_deadline of {budget:?}"
343 ))));
344 }
345 };
346 let rpc_budget = pair.remaining();
347 let rpc = svc.get_current_max_safe(tsoracle_proto::v1::GetCurrentMaxSafeRequest {});
348 let err = match tokio::time::timeout(rpc_budget, rpc).await {
349 Ok(Ok(response)) => {
350 let inner = response.into_inner();
351 return Ok(MaxSafe {
352 max_safe_physical_ms: inner.max_safe_physical_ms,
353 epoch: Epoch::from_wire(inner.epoch_hi, inner.epoch_lo),
354 });
355 }
356 Ok(Err(status)) => ClientError::Rpc(status),
357 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
361 "rpc exceeded its share of per_attempt_deadline \
362 ({rpc_budget:?} of {budget:?})"
363 ))),
364 };
365 if crate::retry_policy::is_transport_failure(&err) {
366 self.pool.evict_if_current(&endpoint, &cell);
367 }
368 Err(err)
369 }
370
371 pub async fn acquire_lease(
378 &self,
379 holder: &[u8],
380 holder_epoch: u64,
381 ttl: Duration,
382 ) -> Result<Lease, ClientError> {
383 let ttl_ms = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX);
384 let (endpoint, budget, pair, mut svc, cell) = self.control_plane_client().await?;
385 let rpc_budget = pair.remaining();
386 let rpc = svc.acquire_lease(tsoracle_proto::v1::AcquireLeaseRequest {
387 holder: holder.to_vec(),
388 holder_epoch,
389 ttl_ms,
390 });
391 let err = match tokio::time::timeout(rpc_budget, rpc).await {
392 Ok(Ok(response)) => {
393 let inner = response.into_inner();
394 return Ok(Lease {
395 lease_id: inner.lease_id,
396 ts_upper_bound: inner.ts_upper_bound,
397 expires_at_ms: inner.expires_at_ms,
398 epoch: required_epoch(inner.epoch, "lease response missing epoch")?,
399 });
400 }
401 Ok(Err(status)) => ClientError::Rpc(status),
402 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
403 "rpc exceeded its share of per_attempt_deadline \
404 ({rpc_budget:?} of {budget:?})"
405 ))),
406 };
407 if crate::retry_policy::is_transport_failure(&err) {
408 self.pool.evict_if_current(&endpoint, &cell);
409 }
410 Err(err)
411 }
412
413 pub async fn renew_lease(&self, lease_id: u64) -> Result<LeaseRenewal, ClientError> {
417 let (endpoint, budget, pair, mut svc, cell) = self.control_plane_client().await?;
418 let rpc_budget = pair.remaining();
419 let rpc = svc.renew_lease(tsoracle_proto::v1::RenewLeaseRequest { lease_id });
420 let err = match tokio::time::timeout(rpc_budget, rpc).await {
421 Ok(Ok(response)) => {
422 let inner = response.into_inner();
423 return Ok(LeaseRenewal {
424 ts_upper_bound: inner.ts_upper_bound,
425 expires_at_ms: inner.expires_at_ms,
426 epoch: required_epoch(inner.epoch, "lease renewal missing epoch")?,
427 });
428 }
429 Ok(Err(status)) => ClientError::Rpc(status),
430 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
431 "rpc exceeded its share of per_attempt_deadline \
432 ({rpc_budget:?} of {budget:?})"
433 ))),
434 };
435 if crate::retry_policy::is_transport_failure(&err) {
436 self.pool.evict_if_current(&endpoint, &cell);
437 }
438 Err(err)
439 }
440
441 pub async fn release_lease(&self, lease_id: u64) -> Result<(), ClientError> {
444 let (endpoint, budget, pair, mut svc, cell) = self.control_plane_client().await?;
445 let rpc_budget = pair.remaining();
446 let rpc = svc.release_lease(tsoracle_proto::v1::ReleaseLeaseRequest { lease_id });
447 let err = match tokio::time::timeout(rpc_budget, rpc).await {
448 Ok(Ok(_response)) => return Ok(()),
449 Ok(Err(status)) => ClientError::Rpc(status),
450 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
451 "rpc exceeded its share of per_attempt_deadline \
452 ({rpc_budget:?} of {budget:?})"
453 ))),
454 };
455 if crate::retry_policy::is_transport_failure(&err) {
456 self.pool.evict_if_current(&endpoint, &cell);
457 }
458 Err(err)
459 }
460
461 pub async fn get_safe_frontier(&self) -> Result<SafeFrontier, ClientError> {
463 let (endpoint, budget, pair, mut svc, cell) = self.control_plane_client().await?;
464 let rpc_budget = pair.remaining();
465 let rpc = svc.get_safe_frontier(tsoracle_proto::v1::GetSafeFrontierRequest {});
466 let err = match tokio::time::timeout(rpc_budget, rpc).await {
467 Ok(Ok(response)) => {
468 let inner = response.into_inner();
469 return Ok(SafeFrontier {
470 frontier_physical_ms: inner.frontier_physical_ms,
471 epoch: Epoch::from_wire(inner.epoch_hi, inner.epoch_lo),
472 });
473 }
474 Ok(Err(status)) => ClientError::Rpc(status),
475 Err(_) => ClientError::Rpc(tonic::Status::deadline_exceeded(format!(
476 "rpc exceeded its share of per_attempt_deadline \
477 ({rpc_budget:?} of {budget:?})"
478 ))),
479 };
480 if crate::retry_policy::is_transport_failure(&err) {
481 self.pool.evict_if_current(&endpoint, &cell);
482 }
483 Err(err)
484 }
485}
486
487#[derive(Copy, Clone, Debug, PartialEq, Eq)]
490pub struct MaxSafe {
491 pub max_safe_physical_ms: u64,
494 pub epoch: Epoch,
496}
497
498#[derive(Copy, Clone, Debug, PartialEq, Eq)]
500pub struct Lease {
501 pub lease_id: u64,
502 pub ts_upper_bound: u64,
503 pub expires_at_ms: u64,
504 pub epoch: Epoch,
505}
506
507#[derive(Copy, Clone, Debug, PartialEq, Eq)]
509pub struct LeaseRenewal {
510 pub ts_upper_bound: u64,
511 pub expires_at_ms: u64,
512 pub epoch: Epoch,
513}
514
515#[derive(Copy, Clone, Debug, PartialEq, Eq)]
517pub struct SafeFrontier {
518 pub frontier_physical_ms: u64,
519 pub epoch: Epoch,
520}
521
522fn required_epoch(
523 epoch: Option<tsoracle_proto::v1::EpochWire>,
524 message: &'static str,
525) -> Result<Epoch, ClientError> {
526 let epoch = epoch.ok_or_else(|| ClientError::Rpc(tonic::Status::internal(message)))?;
527 Ok(Epoch::from_wire(epoch.hi, epoch.lo))
528}
529
530#[cfg(test)]
531mod tests {
532 use super::*;
533
534 #[tokio::test]
535 async fn cached_leader_is_none_before_any_rpc() {
536 let client = Client::connect(vec!["http://127.0.0.1:1".into()])
542 .await
543 .expect("build with a non-empty endpoint list must succeed");
544 assert_eq!(client.cached_leader(), None);
545 }
546
547 #[tokio::test]
548 async fn build_rejects_empty_endpoint_list() {
549 match ClientBuilder::endpoints(Vec::new()).build().await {
553 Err(ClientError::NoReachableEndpoints) => {}
554 Err(other) => panic!("expected NoReachableEndpoints, got {other:?}"),
555 Ok(_) => panic!("expected Err, got Ok(Client)"),
556 }
557 }
558
559 #[tokio::test]
560 async fn channel_connector_error_surfaces_as_connector_variant() {
561 let builder = ClientBuilder::endpoints(vec!["a:1".into()]).channel_connector(
562 |_endpoint: &str| async move {
563 Err::<tonic::transport::Channel, crate::BoxError>(
564 std::io::Error::other("boom").into(),
565 )
566 },
567 );
568 let client = builder.build().await.expect("build must not fail");
569 let result = client.get_ts().await;
570 match result {
571 Err(ClientError::Connector(inner)) => {
572 assert!(inner.to_string().contains("boom"));
573 }
574 other => panic!("expected ClientError::Connector, got {other:?}"),
575 }
576 }
577
578 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
583 async fn marker_connector_failure() -> Result<tonic::transport::Channel, crate::BoxError> {
584 Err("MARKER".into())
585 }
586
587 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
588 #[tokio::test]
589 async fn tls_config_then_channel_connector_last_wins() {
590 let builder = ClientBuilder::endpoints(vec!["a:1".into()])
594 .tls_config(tonic::transport::ClientTlsConfig::new())
595 .channel_connector(|_endpoint: &str| marker_connector_failure());
596 let client = builder.build().await.expect("build must not fail");
597 match client.get_ts().await {
598 Err(ClientError::Connector(inner)) => {
599 assert!(inner.to_string().contains("MARKER"));
600 }
601 other => panic!("expected Connector(MARKER), got {other:?}"),
602 }
603 }
604
605 #[cfg(any(feature = "tls-rustls", feature = "tls-native"))]
606 #[tokio::test]
607 async fn channel_connector_then_tls_config_last_wins() {
608 let builder = ClientBuilder::endpoints(vec!["a:1".into()])
612 .channel_connector(|_endpoint: &str| marker_connector_failure())
613 .tls_config(tonic::transport::ClientTlsConfig::new());
614 let client = builder.build().await.expect("build must not fail");
615 let result = client.get_ts().await;
616 if let Err(ClientError::Connector(inner)) = &result
617 && inner.to_string().contains("MARKER")
618 {
619 panic!("tls_config set last must overwrite the prior channel_connector");
620 }
621 }
622
623 #[tokio::test]
624 async fn batch_flush_interval_overrides_default() {
625 let custom = Duration::from_millis(25);
630 let builder = ClientBuilder::endpoints(vec!["http://127.0.0.1:1".into()])
631 .batch_flush_interval(custom);
632 assert_eq!(builder.flush_interval, custom);
633 }
634
635 #[tokio::test]
636 async fn retry_policy_override_propagates_to_builder() {
637 let policy = RetryPolicy {
642 max_attempts: 7,
643 per_attempt_deadline: Duration::from_millis(11),
644 overall_deadline: Duration::from_millis(13),
645 base_backoff: Duration::from_millis(17),
646 leader_ttl: Duration::from_millis(19),
647 };
648 let builder = ClientBuilder::endpoints(vec!["http://127.0.0.1:1".into()])
649 .retry_policy(policy.clone());
650 assert_eq!(builder.retry_policy.max_attempts, policy.max_attempts);
651 assert_eq!(
652 builder.retry_policy.per_attempt_deadline,
653 policy.per_attempt_deadline
654 );
655 assert_eq!(
656 builder.retry_policy.overall_deadline,
657 policy.overall_deadline
658 );
659 assert_eq!(builder.retry_policy.base_backoff, policy.base_backoff);
660 assert_eq!(builder.retry_policy.leader_ttl, policy.leader_ttl);
661 }
662
663 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
676 async fn get_current_max_safe_returns_within_per_attempt_deadline_when_connector_hangs() {
677 let policy = RetryPolicy {
678 max_attempts: 1,
679 per_attempt_deadline: Duration::from_millis(100),
680 overall_deadline: Duration::from_millis(300),
681 base_backoff: Duration::ZERO,
682 leader_ttl: Duration::from_secs(30),
683 };
684 let client = ClientBuilder::endpoints(vec!["hang:1".into()])
685 .channel_connector(|_endpoint: &str| async {
686 std::future::pending::<Result<tonic::transport::Channel, crate::BoxError>>().await
692 })
693 .retry_policy(policy)
694 .build()
695 .await
696 .expect("builder must accept the policy");
697 let outer_safety = Duration::from_secs(5);
704 let start = std::time::Instant::now();
705 let result = match tokio::time::timeout(outer_safety, client.get_current_max_safe()).await {
706 Ok(r) => r,
707 Err(_) => panic!(
708 "get_current_max_safe failed to honor its own per_attempt_deadline; \
709 the {outer_safety:?} outer safety net had to fire — this is the \
710 security finding's exact symptom (channel acquisition or RPC was \
711 never bounded)",
712 ),
713 };
714 let elapsed = start.elapsed();
715 assert!(
716 result.is_err(),
717 "a hanging connector must surface as Err, got {result:?}",
718 );
719 assert!(
720 elapsed < Duration::from_secs(2),
721 "deadline must short-circuit; took {elapsed:?} (per_attempt_deadline was 100ms)",
722 );
723 }
724
725 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
726 async fn get_ts_returns_within_overall_deadline_when_all_endpoints_unreachable() {
727 let policy = RetryPolicy {
734 max_attempts: 3,
735 per_attempt_deadline: Duration::from_millis(100),
736 overall_deadline: Duration::from_millis(300),
737 base_backoff: Duration::ZERO,
738 leader_ttl: Duration::from_secs(30),
739 };
740 let client = ClientBuilder::endpoints(vec![
741 "http://127.0.0.1:1".into(),
742 "http://127.0.0.1:2".into(),
743 "http://127.0.0.1:3".into(),
744 ])
745 .retry_policy(policy)
746 .build()
747 .await
748 .expect("builder must accept the policy");
749 let start = std::time::Instant::now();
750 let result = client.get_ts().await;
751 let elapsed = start.elapsed();
752 assert!(result.is_err(), "no listener can reply: {result:?}");
753 assert!(
754 elapsed < Duration::from_secs(2),
755 "deadline must short-circuit; took {elapsed:?}"
756 );
757 }
758
759 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
765 async fn get_seq_returns_gapless_blocks_and_rejects_invalid_inputs() {
766 use std::sync::Arc;
767 use std::sync::atomic::{AtomicU64, Ordering};
768 const EPOCH: u128 = 42;
769
770 let counter = Arc::new(AtomicU64::new(0));
771 let server_counter = counter.clone();
772 let addr = crate::test_support::FakeTso::new()
775 .on_get_seq(move |req| {
776 let counter = server_counter.clone();
777 async move {
778 if req.key.is_empty() {
779 return Err(tonic::Status::invalid_argument("invalid sequence key"));
780 }
781 if req.count == 0 {
782 return Err(tonic::Status::invalid_argument(
783 "count must be between 1 and the maximum",
784 ));
785 }
786 let start = counter.fetch_add(u64::from(req.count), Ordering::SeqCst);
787 let (hi, lo) = tsoracle_core::Epoch(EPOCH).to_wire();
788 Ok(tsoracle_proto::v1::GetSeqResponse {
789 key: req.key,
790 start,
791 count: req.count,
792 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
793 })
794 }
795 })
796 .spawn()
797 .await;
798
799 let endpoint = format!("http://{addr}");
800 let client = ClientBuilder::endpoints(vec![endpoint])
801 .retry_policy(RetryPolicy {
802 max_attempts: 3,
803 per_attempt_deadline: Duration::from_secs(2),
804 overall_deadline: Duration::from_secs(5),
805 base_backoff: Duration::ZERO,
806 leader_ttl: Duration::from_secs(30),
807 })
808 .build()
809 .await
810 .expect("client must build");
811
812 let block1 = client
814 .get_seq("orders", 5)
815 .await
816 .expect("first get_seq must succeed");
817 assert_eq!(block1.start, 0, "first block must start at 0");
818 assert_eq!(block1.count, 5, "first block must have count 5");
819 assert_eq!(block1.epoch, EPOCH, "epoch must match the server's epoch");
820
821 let block2 = client
823 .get_seq("orders", 3)
824 .await
825 .expect("second get_seq must succeed");
826 assert_eq!(block2.start, 5, "second block must start at 5 (gapless)");
827 assert_eq!(block2.count, 3);
828 assert_eq!(block2.epoch, EPOCH);
829
830 match client.get_seq("", 1).await {
832 Err(ClientError::InvalidSeqKey) => {}
833 other => panic!("empty key must return InvalidSeqKey, got {other:?}"),
834 }
835
836 match client.get_seq("orders", 0).await {
838 Err(ClientError::InvalidCount(0)) => {}
839 other => panic!("count=0 must return InvalidCount, got {other:?}"),
840 }
841 }
842
843 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
849 async fn get_seq_forwards_large_count_to_server() {
850 let addr = crate::test_support::FakeTso::new()
851 .on_get_seq(|req| async move {
852 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
853 Ok(tsoracle_proto::v1::GetSeqResponse {
854 key: req.key,
855 start: 0,
856 count: req.count,
857 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
858 })
859 })
860 .spawn()
861 .await;
862
863 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
864 .retry_policy(RetryPolicy {
865 max_attempts: 3,
866 per_attempt_deadline: Duration::from_secs(2),
867 overall_deadline: Duration::from_secs(5),
868 base_backoff: Duration::ZERO,
869 leader_ttl: Duration::from_secs(30),
870 })
871 .build()
872 .await
873 .unwrap();
874
875 let big = tsoracle_core::DEFAULT_MAX_SEQ_COUNT + 1;
878 let block = client
879 .get_seq("orders", big)
880 .await
881 .expect("a large count must be forwarded to the server, not locally rejected");
882 assert_eq!(block.count, big);
883 }
884
885 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
894 async fn get_seq_over_cap_count_surfaces_rpc_invalid_argument() {
895 use std::sync::Arc;
896 use std::sync::atomic::{AtomicU64, Ordering};
897 let calls = Arc::new(AtomicU64::new(0));
898 let server_calls = calls.clone();
899 let addr = crate::test_support::FakeTso::new()
903 .on_get_seq(move |_req| {
904 let calls = server_calls.clone();
905 async move {
906 calls.fetch_add(1, Ordering::SeqCst);
907 Err(tonic::Status::invalid_argument(
908 "count must be between 1 and the maximum",
909 ))
910 }
911 })
912 .spawn()
913 .await;
914
915 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
916 .retry_policy(RetryPolicy {
917 max_attempts: 3,
918 per_attempt_deadline: Duration::from_secs(2),
919 overall_deadline: Duration::from_secs(5),
920 base_backoff: Duration::ZERO,
921 leader_ttl: Duration::from_secs(30),
922 })
923 .build()
924 .await
925 .unwrap();
926
927 let big = tsoracle_core::DEFAULT_MAX_SEQ_COUNT + 1;
930 match client.get_seq("orders", big).await {
931 Err(ClientError::Rpc(status)) => {
932 assert_eq!(status.code(), tonic::Code::InvalidArgument);
933 assert!(
934 status.message().contains("maximum"),
935 "must surface the server's over-cap message, not a synthetic \
936 or mislabelled error; got {:?}",
937 status.message()
938 );
939 }
940 other => panic!(
941 "an over-cap count with a valid key must surface as \
942 ClientError::Rpc(InvalidArgument), not InvalidSeqKey; got {other:?}"
943 ),
944 }
945 let n = calls.load(Ordering::SeqCst);
947 assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
948 }
949
950 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
956 async fn get_seq_bare_failed_precondition_surfaces_definitive_error() {
957 use std::sync::Arc;
958 use std::sync::atomic::{AtomicU64, Ordering};
959 let calls = Arc::new(AtomicU64::new(0));
960 let server_calls = calls.clone();
961 let addr = crate::test_support::FakeTso::new()
964 .on_get_seq(move |_req| {
965 let calls = server_calls.clone();
966 async move {
967 calls.fetch_add(1, Ordering::SeqCst);
968 Err(tonic::Status::failed_precondition("dense counter overflow"))
969 }
970 })
971 .spawn()
972 .await;
973
974 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
975 .retry_policy(RetryPolicy {
976 max_attempts: 3,
977 per_attempt_deadline: Duration::from_secs(2),
978 overall_deadline: Duration::from_secs(5),
979 base_backoff: Duration::ZERO,
980 leader_ttl: Duration::from_secs(30),
981 })
982 .build()
983 .await
984 .unwrap();
985
986 match client.get_seq("orders", 1).await {
987 Err(ClientError::Rpc(status)) => {
988 assert_eq!(status.code(), tonic::Code::FailedPrecondition);
989 assert!(
990 status.message().contains("overflow"),
991 "must surface the server's overflow message, not a synthetic \
992 'no leader yet'; got {:?}",
993 status.message()
994 );
995 }
996 other => panic!("expected the overflow FailedPrecondition, got {other:?}"),
997 }
998 let n = calls.load(Ordering::SeqCst);
1000 assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
1001 }
1002
1003 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1009 async fn get_seq_ambiguous_post_send_failure_is_uncertain_without_retry() {
1010 use std::sync::Arc;
1011 use std::sync::atomic::{AtomicU64, Ordering};
1012 let calls = Arc::new(AtomicU64::new(0));
1013 let server_calls = calls.clone();
1014 let addr = crate::test_support::FakeTso::new()
1017 .on_get_seq(move |_req| {
1018 let calls = server_calls.clone();
1019 async move {
1020 calls.fetch_add(1, Ordering::SeqCst);
1021 Err(tonic::Status::unavailable("connection reset mid-call"))
1022 }
1023 })
1024 .spawn()
1025 .await;
1026
1027 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1028 .retry_policy(RetryPolicy {
1029 max_attempts: 3,
1030 per_attempt_deadline: Duration::from_secs(2),
1031 overall_deadline: Duration::from_secs(5),
1032 base_backoff: Duration::ZERO,
1033 leader_ttl: Duration::from_secs(30),
1034 })
1035 .build()
1036 .await
1037 .unwrap();
1038
1039 match client.get_seq("orders", 1).await {
1040 Err(ClientError::SeqUncertain) => {}
1041 other => panic!("post-send Unavailable must be SeqUncertain, got {other:?}"),
1042 }
1043 assert_eq!(
1044 calls.load(Ordering::SeqCst),
1045 1,
1046 "non-idempotent get_seq must NOT retry an ambiguous post-send failure"
1047 );
1048 }
1049
1050 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1057 async fn get_seq_post_send_internal_is_uncertain_without_retry() {
1058 use std::sync::Arc;
1059 use std::sync::atomic::{AtomicU64, Ordering};
1060 let calls = Arc::new(AtomicU64::new(0));
1061 let server_calls = calls.clone();
1062 let addr = crate::test_support::FakeTso::new()
1065 .on_get_seq(move |_req| {
1066 let calls = server_calls.clone();
1067 async move {
1068 calls.fetch_add(1, Ordering::SeqCst);
1069 Err(tonic::Status::internal("permanent dense driver fault"))
1070 }
1071 })
1072 .spawn()
1073 .await;
1074
1075 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1076 .retry_policy(RetryPolicy {
1077 max_attempts: 3,
1078 per_attempt_deadline: Duration::from_secs(2),
1079 overall_deadline: Duration::from_secs(5),
1080 base_backoff: Duration::ZERO,
1081 leader_ttl: Duration::from_secs(30),
1082 })
1083 .build()
1084 .await
1085 .unwrap();
1086
1087 match client.get_seq("orders", 1).await {
1088 Err(ClientError::SeqUncertain) => {}
1089 other => panic!("post-send INTERNAL must be SeqUncertain, got {other:?}"),
1090 }
1091 assert_eq!(
1092 calls.load(Ordering::SeqCst),
1093 1,
1094 "post-send INTERNAL must NOT be retried (possible committed advance)"
1095 );
1096 }
1097
1098 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1104 async fn get_seq_response_count_mismatch_is_uncertain() {
1105 use std::sync::Arc;
1106 use std::sync::atomic::{AtomicU64, Ordering};
1107 let calls = Arc::new(AtomicU64::new(0));
1108 let server_calls = calls.clone();
1109 let addr = crate::test_support::FakeTso::new()
1112 .on_get_seq(move |req| {
1113 let calls = server_calls.clone();
1114 async move {
1115 calls.fetch_add(1, Ordering::SeqCst);
1116 let (hi, lo) = tsoracle_core::Epoch(7).to_wire();
1117 Ok(tsoracle_proto::v1::GetSeqResponse {
1118 key: req.key,
1119 start: 0,
1120 count: req.count + 1,
1121 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1122 })
1123 }
1124 })
1125 .spawn()
1126 .await;
1127
1128 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1129 .retry_policy(RetryPolicy {
1130 max_attempts: 3,
1131 per_attempt_deadline: Duration::from_secs(2),
1132 overall_deadline: Duration::from_secs(5),
1133 base_backoff: Duration::ZERO,
1134 leader_ttl: Duration::from_secs(30),
1135 })
1136 .build()
1137 .await
1138 .unwrap();
1139
1140 match client.get_seq("orders", 5).await {
1141 Err(ClientError::SeqUncertain) => {}
1142 other => panic!("count-mismatch success must be SeqUncertain, got {other:?}"),
1143 }
1144 assert_eq!(
1145 calls.load(Ordering::SeqCst),
1146 1,
1147 "a malformed success must not trigger a (double-spending) retry"
1148 );
1149 }
1150
1151 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1154 async fn get_seq_response_key_mismatch_is_uncertain() {
1155 use std::sync::Arc;
1156 use std::sync::atomic::{AtomicU64, Ordering};
1157 let calls = Arc::new(AtomicU64::new(0));
1158 let server_calls = calls.clone();
1159 let addr = crate::test_support::FakeTso::new()
1162 .on_get_seq(move |req| {
1163 let calls = server_calls.clone();
1164 async move {
1165 calls.fetch_add(1, Ordering::SeqCst);
1166 let (hi, lo) = tsoracle_core::Epoch(7).to_wire();
1167 Ok(tsoracle_proto::v1::GetSeqResponse {
1168 key: format!("{}-tampered", req.key),
1169 start: 0,
1170 count: req.count,
1171 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1172 })
1173 }
1174 })
1175 .spawn()
1176 .await;
1177
1178 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1179 .retry_policy(RetryPolicy {
1180 max_attempts: 3,
1181 per_attempt_deadline: Duration::from_secs(2),
1182 overall_deadline: Duration::from_secs(5),
1183 base_backoff: Duration::ZERO,
1184 leader_ttl: Duration::from_secs(30),
1185 })
1186 .build()
1187 .await
1188 .unwrap();
1189
1190 match client.get_seq("orders", 5).await {
1191 Err(ClientError::SeqUncertain) => {}
1192 other => panic!("key-mismatch success must be SeqUncertain, got {other:?}"),
1193 }
1194 assert_eq!(
1195 calls.load(Ordering::SeqCst),
1196 1,
1197 "malformed success: one call"
1198 );
1199 }
1200
1201 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1205 async fn get_seq_batch_returns_blocks_in_request_order() {
1206 use std::sync::Arc;
1207 use std::sync::atomic::{AtomicU64, Ordering};
1208 const EPOCH: u128 = 77;
1209
1210 let counter = Arc::new(AtomicU64::new(0));
1213 let server_counter = counter.clone();
1214 let addr = crate::test_support::FakeTso::new()
1215 .on_get_seq_batch(move |req| {
1216 let counter = server_counter.clone();
1217 async move {
1218 let mut grants = Vec::with_capacity(req.entries.len());
1219 for entry in &req.entries {
1220 let start = counter.fetch_add(u64::from(entry.count), Ordering::SeqCst);
1221 grants.push(tsoracle_proto::v1::SeqGrantEntry {
1222 key: entry.key.clone(),
1223 start,
1224 count: entry.count,
1225 });
1226 }
1227 let (hi, lo) = tsoracle_core::Epoch(EPOCH).to_wire();
1228 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1229 grants,
1230 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1231 })
1232 }
1233 })
1234 .spawn()
1235 .await;
1236
1237 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1238 .retry_policy(RetryPolicy {
1239 max_attempts: 3,
1240 per_attempt_deadline: Duration::from_secs(2),
1241 overall_deadline: Duration::from_secs(5),
1242 base_backoff: Duration::ZERO,
1243 leader_ttl: Duration::from_secs(30),
1244 })
1245 .build()
1246 .await
1247 .expect("client must build");
1248
1249 let blocks = client
1250 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1251 .await
1252 .expect("get_seq_batch must succeed");
1253
1254 assert_eq!(blocks.len(), 2, "one block per entry");
1255 assert_eq!(blocks[0].start, 0);
1257 assert_eq!(blocks[0].count, 5);
1258 assert_eq!(blocks[0].epoch, EPOCH);
1259 assert_eq!(blocks[1].start, 5);
1261 assert_eq!(blocks[1].count, 3);
1262 assert_eq!(blocks[1].epoch, EPOCH);
1263 }
1264
1265 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1267 async fn get_seq_batch_rejects_duplicate_keys_without_hitting_server() {
1268 use std::sync::Arc;
1269 use std::sync::atomic::{AtomicU64, Ordering};
1270
1271 let calls = Arc::new(AtomicU64::new(0));
1272 let server_calls = calls.clone();
1273 let addr = crate::test_support::FakeTso::new()
1275 .on_get_seq_batch(move |_req| {
1276 let calls = server_calls.clone();
1277 async move {
1278 calls.fetch_add(1, Ordering::SeqCst);
1279 Err(tonic::Status::internal("should not be called"))
1280 }
1281 })
1282 .spawn()
1283 .await;
1284
1285 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1286 .retry_policy(RetryPolicy {
1287 max_attempts: 3,
1288 per_attempt_deadline: Duration::from_secs(2),
1289 overall_deadline: Duration::from_secs(5),
1290 base_backoff: Duration::ZERO,
1291 leader_ttl: Duration::from_secs(30),
1292 })
1293 .build()
1294 .await
1295 .expect("client must build");
1296
1297 match client.get_seq_batch(&[("orders", 1), ("orders", 2)]).await {
1299 Err(ClientError::InvalidSeqKey) => {}
1300 other => panic!("duplicate key must return InvalidSeqKey, got {other:?}"),
1301 }
1302 assert_eq!(
1303 calls.load(Ordering::SeqCst),
1304 0,
1305 "duplicate-key rejection must not reach the server"
1306 );
1307 }
1308
1309 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1313 async fn get_seq_batch_malformed_success_is_uncertain() {
1314 use std::sync::Arc;
1315 use std::sync::atomic::{AtomicU64, Ordering};
1316
1317 {
1319 let calls = Arc::new(AtomicU64::new(0));
1320 let server_calls = calls.clone();
1321 let addr = crate::test_support::FakeTso::new()
1322 .on_get_seq_batch(move |_req| {
1323 let calls = server_calls.clone();
1324 async move {
1325 calls.fetch_add(1, Ordering::SeqCst);
1326 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1327 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1329 grants: vec![tsoracle_proto::v1::SeqGrantEntry {
1330 key: "orders".into(),
1331 start: 0,
1332 count: 5,
1333 }],
1334 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1335 })
1336 }
1337 })
1338 .spawn()
1339 .await;
1340
1341 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1342 .retry_policy(RetryPolicy {
1343 max_attempts: 3,
1344 per_attempt_deadline: Duration::from_secs(2),
1345 overall_deadline: Duration::from_secs(5),
1346 base_backoff: Duration::ZERO,
1347 leader_ttl: Duration::from_secs(30),
1348 })
1349 .build()
1350 .await
1351 .expect("client must build");
1352
1353 match client
1354 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1355 .await
1356 {
1357 Err(ClientError::SeqUncertain) => {}
1358 other => panic!("wrong grants.len must surface as SeqUncertain, got {other:?}"),
1359 }
1360 assert_eq!(
1361 calls.load(Ordering::SeqCst),
1362 1,
1363 "malformed success must not trigger a retry"
1364 );
1365 }
1366
1367 {
1369 let calls = Arc::new(AtomicU64::new(0));
1370 let server_calls = calls.clone();
1371 let addr = crate::test_support::FakeTso::new()
1372 .on_get_seq_batch(move |req| {
1373 let calls = server_calls.clone();
1374 async move {
1375 calls.fetch_add(1, Ordering::SeqCst);
1376 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1377 let grants = req
1379 .entries
1380 .iter()
1381 .enumerate()
1382 .map(|(i, e)| tsoracle_proto::v1::SeqGrantEntry {
1383 key: e.key.clone(),
1384 start: 0,
1385 count: if i == 0 { e.count + 1 } else { e.count },
1387 })
1388 .collect();
1389 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1390 grants,
1391 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1392 })
1393 }
1394 })
1395 .spawn()
1396 .await;
1397
1398 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1399 .retry_policy(RetryPolicy {
1400 max_attempts: 3,
1401 per_attempt_deadline: Duration::from_secs(2),
1402 overall_deadline: Duration::from_secs(5),
1403 base_backoff: Duration::ZERO,
1404 leader_ttl: Duration::from_secs(30),
1405 })
1406 .build()
1407 .await
1408 .expect("client must build");
1409
1410 match client
1411 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1412 .await
1413 {
1414 Err(ClientError::SeqUncertain) => {}
1415 other => panic!("count-mismatch grant must surface as SeqUncertain, got {other:?}"),
1416 }
1417 assert_eq!(
1418 calls.load(Ordering::SeqCst),
1419 1,
1420 "malformed success must not trigger a retry"
1421 );
1422 }
1423
1424 {
1426 let calls = Arc::new(AtomicU64::new(0));
1427 let server_calls = calls.clone();
1428 let addr = crate::test_support::FakeTso::new()
1429 .on_get_seq_batch(move |req| {
1430 let calls = server_calls.clone();
1431 async move {
1432 calls.fetch_add(1, Ordering::SeqCst);
1433 let grants = req
1434 .entries
1435 .iter()
1436 .map(|e| tsoracle_proto::v1::SeqGrantEntry {
1437 key: e.key.clone(),
1438 start: 0,
1439 count: e.count,
1440 })
1441 .collect();
1442 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1444 grants,
1445 epoch: None,
1446 })
1447 }
1448 })
1449 .spawn()
1450 .await;
1451
1452 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1453 .retry_policy(RetryPolicy {
1454 max_attempts: 3,
1455 per_attempt_deadline: Duration::from_secs(2),
1456 overall_deadline: Duration::from_secs(5),
1457 base_backoff: Duration::ZERO,
1458 leader_ttl: Duration::from_secs(30),
1459 })
1460 .build()
1461 .await
1462 .expect("client must build");
1463
1464 match client
1465 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1466 .await
1467 {
1468 Err(ClientError::SeqUncertain) => {}
1469 other => panic!("missing epoch must surface as SeqUncertain, got {other:?}"),
1470 }
1471 assert_eq!(
1472 calls.load(Ordering::SeqCst),
1473 1,
1474 "malformed success must not trigger a retry"
1475 );
1476 }
1477
1478 {
1481 let calls = Arc::new(AtomicU64::new(0));
1482 let server_calls = calls.clone();
1483 let addr = crate::test_support::FakeTso::new()
1484 .on_get_seq_batch(move |req| {
1485 let calls = server_calls.clone();
1486 async move {
1487 calls.fetch_add(1, Ordering::SeqCst);
1488 let (hi, lo) = tsoracle_core::Epoch(1).to_wire();
1489 let grants = req
1491 .entries
1492 .iter()
1493 .enumerate()
1494 .map(|(i, e)| tsoracle_proto::v1::SeqGrantEntry {
1495 key: if i == 0 {
1496 format!("{}-tampered", e.key)
1497 } else {
1498 e.key.clone()
1499 },
1500 start: 0,
1501 count: e.count,
1502 })
1503 .collect();
1504 Ok(tsoracle_proto::v1::GetSeqBatchResponse {
1505 grants,
1506 epoch: Some(tsoracle_proto::v1::EpochWire { hi, lo }),
1507 })
1508 }
1509 })
1510 .spawn()
1511 .await;
1512
1513 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1514 .retry_policy(RetryPolicy {
1515 max_attempts: 3,
1516 per_attempt_deadline: Duration::from_secs(2),
1517 overall_deadline: Duration::from_secs(5),
1518 base_backoff: Duration::ZERO,
1519 leader_ttl: Duration::from_secs(30),
1520 })
1521 .build()
1522 .await
1523 .expect("client must build");
1524
1525 match client
1526 .get_seq_batch(&[("orders", 5), ("invoices", 3)])
1527 .await
1528 {
1529 Err(ClientError::SeqUncertain) => {}
1530 other => panic!("key-mismatch grant must surface as SeqUncertain, got {other:?}"),
1531 }
1532 assert_eq!(
1533 calls.load(Ordering::SeqCst),
1534 1,
1535 "malformed success must not trigger a retry"
1536 );
1537 }
1538 }
1539
1540 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1543 async fn get_seq_batch_post_send_error_is_uncertain_without_retry() {
1544 use std::sync::Arc;
1545 use std::sync::atomic::{AtomicU64, Ordering};
1546
1547 let calls = Arc::new(AtomicU64::new(0));
1548 let server_calls = calls.clone();
1549 let addr = crate::test_support::FakeTso::new()
1550 .on_get_seq_batch(move |_req| {
1551 let calls = server_calls.clone();
1552 async move {
1553 calls.fetch_add(1, Ordering::SeqCst);
1554 Err(tonic::Status::internal("permanent dense driver fault"))
1555 }
1556 })
1557 .spawn()
1558 .await;
1559
1560 let client = ClientBuilder::endpoints(vec![format!("http://{addr}")])
1561 .retry_policy(RetryPolicy {
1562 max_attempts: 3,
1563 per_attempt_deadline: Duration::from_secs(2),
1564 overall_deadline: Duration::from_secs(5),
1565 base_backoff: Duration::ZERO,
1566 leader_ttl: Duration::from_secs(30),
1567 })
1568 .build()
1569 .await
1570 .expect("client must build");
1571
1572 match client
1573 .get_seq_batch(&[("orders", 1), ("invoices", 2)])
1574 .await
1575 {
1576 Err(ClientError::SeqUncertain) => {}
1577 other => panic!("post-send INTERNAL must be SeqUncertain, got {other:?}"),
1578 }
1579 assert_eq!(
1580 calls.load(Ordering::SeqCst),
1581 1,
1582 "post-send INTERNAL must NOT be retried (possible committed advance)"
1583 );
1584 }
1585}