Skip to main content

tsoracle_client/
lib.rs

1//
2//  ░▀█▀░█▀▀░█▀█░█▀▄░█▀█░█▀▀░█░░░█▀▀
3//  ░░█░░▀▀█░█░█░█▀▄░█▀█░█░░░█░░░█▀▀
4//  ░░▀░░▀▀▀░▀▀▀░▀░▀░▀░▀░▀▀▀░▀▀▀░▀▀▀
5//
6//  tsoracle — Distributed Timestamp Oracle
7//  https://www.tsoracle.rs
8//
9//  Copyright (c) 2026 Prisma Risk
10//
11//  Licensed under the Apache License, Version 2.0 (the "License");
12//  you may not use this file except in compliance with the License.
13//  You may obtain a copy of the License at
14//
15//      https://www.apache.org/licenses/LICENSE-2.0
16//
17//  Unless required by applicable law or agreed to in writing, software
18//  distributed under the License is distributed on an "AS IS" BASIS,
19//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20//  See the License for the specific language governing permissions and
21//  limitations under the License.
22//
23
24#![doc = include_str!("../README.md")]
25// Panic policy (see CONTRIBUTING.md). `cfg_attr(not(test), ...)` skips the lint
26// for the lib's own unit tests; integration tests are separate compilation units.
27#![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
55/// The server's per-call cap on requested timestamps, fixed by the 18-bit
56/// logical width. Callers asking for more than this can't be served by any
57/// single RPC; the client rejects them up-front rather than burning a queue
58/// slot and round-trip to learn the same thing from the server.
59pub(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    /// Override the default [`RetryPolicy`].
88    ///
89    /// The policy controls per-attempt deadlines, the overall deadline
90    /// across all candidate endpoints, the cap on attempts, and the
91    /// jittered backoff base. The per-attempt deadline is also pushed
92    /// down to `tonic::transport::Endpoint::connect_timeout` and
93    /// `Endpoint::timeout` for the built-in default and TLS transport
94    /// paths so a blackholed peer fails fast at the transport layer.
95    /// User-supplied [`Self::channel_connector`] closures own their
96    /// own Endpoint config; the policy still bounds the retry loop's
97    /// outer `tokio::time::timeout` around them.
98    pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
99        self.retry_policy = policy;
100        self
101    }
102
103    /// Configure the client to dial bare endpoints with TLS. Bare `host:port`
104    /// becomes `https://host:port`; explicit `http://...` endpoints supplied
105    /// in [`Self::endpoints`] remain plaintext; explicit `https://...`
106    /// endpoints use the provided TLS config.
107    ///
108    /// Wire-supplied `http://...` leader-hint trailers are NOT honored under
109    /// `tls_config` — they are dropped to prevent a contacted peer from
110    /// downgrading the transport. Operator-supplied configuration still wins;
111    /// untrusted wire input does not.
112    ///
113    /// Setting both [`Self::channel_connector`] and `tls_config` is allowed;
114    /// the last call wins (standard builder semantics). Calling
115    /// `channel_connector` after `tls_config` also clears the
116    /// reject-plaintext-hint policy, since the caller-owned connector owns
117    /// its own scheme policy.
118    #[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    /// Replace the default plaintext channel construction with a caller-owned
129    /// closure. The closure is invoked on first use of each endpoint —
130    /// configured endpoints and leader-hint redirects alike. Errors returned
131    /// from the closure surface as [`ClientError::Connector`].
132    ///
133    /// See module docs for the interaction with [`Self::tls_config`]
134    /// (last-wins) and the scheme matrix. A caller-owned connector replaces
135    /// the built-in TLS plumbing entirely, including the
136    /// reject-plaintext-leader-hint policy — the closure is responsible for
137    /// whatever scheme policy it wants to enforce.
138    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    /// The endpoint the client currently believes is the leader, or `None`
187    /// if no leader has been observed yet or the cached entry has aged past
188    /// the configured `leader_ttl`.
189    ///
190    /// Read-only diagnostic surface for ops dashboards and integration tests
191    /// asserting that a client has converged to the expected leader. It
192    /// reflects the cache as last updated by a completed RPC — it neither
193    /// triggers nor waits on any network round-trip, and the TTL check is
194    /// lazy (an expired entry reads as `None`).
195    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    /// Request a contiguous block of `count` dense ordinals for the given
211    /// sequence `key`.
212    ///
213    /// **Non-idempotent:** if this call returns `ClientError::SeqUncertain` it
214    /// means a commit may or may not have occurred on the server — do NOT
215    /// silently retry, as that would risk spending a second block. The caller
216    /// must resolve the ambiguity (e.g. by reading back the current counter
217    /// with a coordinator-level read) before re-issuing.
218    ///
219    /// All other errors are pre-commit-certain (the server rejected the request
220    /// before any durable advance) and are safe to retry.
221    ///
222    /// The client only pre-rejects universally-invalid inputs — an empty/oversized
223    /// `key` and a zero `count` (a block always covers at least one ordinal).
224    /// The upper bound on `count` is the server's configured
225    /// `ServerBuilder::max_seq_count`, which is deployment-specific, so the client
226    /// does not second-guess it: an over-cap `count` is forwarded and the server
227    /// rejects it with `INVALID_ARGUMENT` (surfaced as `ClientError::Rpc`). This
228    /// keeps a client built against one cap correct against a server configured
229    /// with another.
230    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    /// Atomically request contiguous dense blocks for several DISTINCT keys in
241    /// one round-trip. Returns one [`SeqBlock`] per entry, in request order.
242    ///
243    /// **Non-idempotent (whole batch):** on `ClientError::SeqUncertain` the
244    /// entire batch may or may not have committed — do NOT retry; reconcile
245    /// first (the atomic server contract means it is all-or-nothing, so there is
246    /// never a partial mix to untangle). All other errors are pre-commit-certain.
247    ///
248    /// Client-side pre-checks reject only universally-invalid input: an empty
249    /// batch, a duplicate key, an empty/oversized key, or a zero count. The
250    /// per-entry count cap and the batch-size cap are server-side and are
251    /// forwarded for the server to reject, so a client built against one
252    /// configuration stays correct against a server with another.
253    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                // Distinct-key contract — mirror the server's pre-commit reject.
270                return Err(ClientError::InvalidSeqKey);
271            }
272        }
273        retry::issue_seq_batch_rpc(&self.pool, entries).await
274    }
275
276    /// Read the leader's current safe-point in physical-millisecond units.
277    ///
278    /// Targets the cached leader if known, otherwise the first configured
279    /// endpoint. Followers return `MaxSafe::max_safe_physical_ms == 0` rather
280    /// than erroring, matching the proto contract; pollers needing freshness
281    /// should target the leader endpoint.
282    ///
283    /// Single-shot by design — the proto contract is "followers return 0"
284    /// rather than NOT_LEADER, so there is no hint to chase and no worklist
285    /// to drain; a caller polling for freshness retries the next tick rather
286    /// than the next endpoint. The one `(connect, RPC)` pair is bounded by
287    /// [`RetryPolicy::per_attempt_deadline`] (shared across both phases via
288    /// `PairBudget`, exactly like one `get_ts` attempt), and a transport-class
289    /// failure evicts the cached channel so a half-open / black-holing
290    /// connection is dropped before the next poll lands on it.
291    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        // One budget shared across connect + RPC: a slow connect eats into
299        // the RPC's time rather than each phase getting a fresh full budget,
300        // so a single call never runs longer than ~`per_attempt_deadline`.
301        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                // `client_with_cell` already evicts its own cell on a failed
306                // dial, so there is nothing to evict here.
307                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            // A timed-out RPC surfaces as `DeadlineExceeded` (transport-class
326            // per `is_transport_failure`), so the eviction tail below drops the
327            // possibly-half-open channel — matching the `get_ts` attempt path.
328            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/// The leader's view of the durable safe-point, returned by
341/// [`Client::get_current_max_safe`].
342#[derive(Copy, Clone, Debug, PartialEq, Eq)]
343pub struct MaxSafe {
344    /// Safe-point in physical-millisecond units; `0` is the cold-start sentinel
345    /// (also the value any follower returns).
346    pub max_safe_physical_ms: u64,
347    /// Leader epoch that issued this view.
348    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        // A freshly built client has issued no RPC, so the channel pool's
358        // leader cache is empty and `cached_leader()` reports `None`. This
359        // pins the "nothing observed yet" branch of the diagnostic accessor
360        // without needing a server; the post-RPC `Some(addr)` case is covered
361        // end-to-end in `tsoracle-tests`.
362        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        // Validation prevents a Client whose `pool` has no endpoints to try;
371        // every RPC would fail-fast with `NoReachableEndpoints` and burn no
372        // network roundtrips at all, so reject up-front instead.
373        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    // Marker payload used by both last-wins tests. The free function holds
400    // the body in one source location so coverage credit flows from the test
401    // where the closure DOES run; the test that asserts "this never runs"
402    // just calls the same helper.
403    #[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        // channel_connector is set LAST, so its path runs on get_ts. The
412        // marker error surfaces as `ClientError::Connector`, proving the
413        // builder did not silently keep the prior tls_config.
414        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        // tls_config is set LAST, so the connector path is replaced and its
430        // marker error must NOT surface. The tls_config path produces a
431        // transport-level failure (or NoReachableEndpoints / Rpc) instead.
432        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        // The builder's `batch_flush_interval` knob feeds the driver's
447        // coalescing window; without a test it could silently revert to the
448        // default and no-one would notice from black-box behavior. We
449        // confirm the override path by reaching into the builder fields.
450        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        // The builder field is what `build` hands to the pool and retry
459        // loop. If `retry_policy()` ever silently stops storing the
460        // override, the loop reverts to defaults; this test pins the
461        // override path against that.
462        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    /// Acceptance criterion for the `get_current_max_safe` deadline fix
485    /// (security finding 8c3ea943): a single-shot call against a
486    /// black-holed endpoint must surface `DeadlineExceeded` within the
487    /// configured `per_attempt_deadline`, not park indefinitely on a
488    /// connector whose future never resolves. Exercised through the
489    /// public `channel_connector` surface — `std::future::pending()`
490    /// guarantees the connect future never completes, so any code path
491    /// without an outer `tokio::time::timeout` hangs until an external
492    /// test timeout kills the runtime. The wall-clock bound is generous
493    /// enough to absorb CI scheduler jitter but well below the OS-level
494    /// TCP timeout, mirroring the structure of the `get_ts`
495    /// equivalent below.
496    #[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                // The future never resolves; only an outer timeout can end
508                // the await. This is the exact attack shape the finding
509                // describes: a user-supplied connector (or a black-holed
510                // peer behind one) for which the caller relied on
511                // `RetryPolicy` to bound the wait.
512                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        // Outer safety timeout: when the fix is missing, `get_current_max_safe`
519        // awaits a never-resolving connector indefinitely, which would hang the
520        // whole test runner rather than fail this case. The 5s ceiling converts
521        // that hang into a clean panic with a diagnostic message. Under a
522        // correctly-deadlined call this outer guard never fires — the inner
523        // future returns in ~100ms, well below the elapsed-time assertion.
524        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        // End-to-end test of the issue's acceptance criterion: with no
549        // listener bound at the configured endpoints, a `get_ts` call
550        // must return well before the OS-default TCP timeout (~75 s on
551        // Linux). The bound here is generous enough to absorb CI
552        // scheduler jitter — the assertion is "fast", not "exactly the
553        // configured deadline".
554        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    /// Integration test: `get_seq` against a real loopback gRPC server that
581    /// simulates a file-driver leader. The server uses a fake `TsoService`
582    /// with a real `get_seq` implementation (incrementing an atomic counter).
583    /// Asserts: consecutive blocks are gapless, the epoch is propagated, and
584    /// invalid inputs are rejected up-front.
585    #[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        // A leader serving gapless blocks from a single atomic counter, rejecting
594        // empty keys / zero counts as InvalidArgument.
595        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        // First call: block [0, 5).
634        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        // Second call: block [5, 8) — gapless continuation.
643        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        // Empty key is rejected up-front by the client, not the server.
652        match client.get_seq("", 1).await {
653            Err(ClientError::InvalidSeqKey) => {}
654            other => panic!("empty key must return InvalidSeqKey, got {other:?}"),
655        }
656
657        // Zero count is rejected up-front by the client.
658        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    /// The upper bound on `count` is the server's configured `max_seq_count`, not
665    /// a fixed client-side constant, so the client must forward a large `count`
666    /// rather than short-circuit it with a local `InvalidCount`. Here a count far
667    /// above the old default reaches a cap-less echo server and is served —
668    /// proving the client dropped its hard-coded ceiling check.
669    #[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        // Far above the old hard-coded default — would have been a local
697        // InvalidCount before; now it is forwarded and the echo server serves it.
698        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    /// The rejection companion to `get_seq_forwards_large_count_to_server`: when
707    /// the server's configured `ServerBuilder::max_seq_count` is smaller than the
708    /// requested `count`, it rejects with `INVALID_ARGUMENT`
709    /// (`CoreError::SeqCountTooLarge`). Because the client no longer pre-checks the
710    /// upper bound, that rejection must surface through `Client::get_seq` as
711    /// `ClientError::Rpc` preserving the server's message — exactly as the `get_seq`
712    /// doc contract promises — and must NOT be mislabelled `InvalidSeqKey`: a valid
713    /// key was supplied; only the count was over-cap.
714    #[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        // Mirrors the server's SeqCountTooLarge → INVALID_ARGUMENT mapping for a
721        // count above the configured max_seq_count. Counts calls to prove the
722        // definitive-error path stays bounded (no unbounded ride-out).
723        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        // A valid key with an over-cap count: the key is fine, so InvalidSeqKey
749        // would be the wrong label — the docs promise ClientError::Rpc.
750        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        // Definitive-error path: bounded attempts, never an unbounded ride-out.
767        let n = calls.load(Ordering::SeqCst);
768        assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
769    }
770
771    /// Regression guard for the dense classification bug: a bare
772    /// `FailedPrecondition` from the server (NO leader-hint trailer) — e.g.
773    /// `SeqOverflow` — must surface immediately as that definitive error,
774    /// preserving the server's message, NOT be misread as a NOT_LEADER
775    /// "no leader yet" and ridden out across passes.
776    #[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        // Bare FailedPrecondition with NO leader-hint trailer — mirrors the
783        // server's SeqOverflow mapping. Counts calls to prove bounded attempts.
784        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        // Definitive-error path: bounded attempts, never an unbounded ride-out.
820        let n = calls.load(Ordering::SeqCst);
821        assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
822    }
823
824    /// The non-idempotency contract: a post-send transport failure
825    /// (`Unavailable` after the request reached the server) is ambiguous — the
826    /// advance may have committed — so `get_seq` returns `SeqUncertain` and must
827    /// NOT transparently retry (a retry could silently spend a second block).
828    /// Asserts the server is invoked exactly once.
829    #[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        // A post-send transport failure (Unavailable after the request reached
836        // the server) — ambiguous; counts calls to prove no retry.
837        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    /// The non-idempotency contract extended to application faults: a post-send
872    /// `INTERNAL` (the server's mapping for a permanent dense driver fault — the
873    /// file driver emits this for a directory-fsync failure that happens *after*
874    /// the durable rename) is ambiguous, not pre-commit-certain. `get_seq` must
875    /// return `SeqUncertain` and invoke the server exactly once, never surfacing
876    /// it as a definitive (caller-retry-safe) error.
877    #[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        // A post-send INTERNAL (permanent dense driver fault) — ambiguous; counts
884        // calls to prove no retry.
885        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    /// A successful `GetSeqResponse` whose `count` does not match the request is
920    /// a protocol violation: the server committed an advance, but the block it
921    /// describes is not the one the caller asked for. The client must not hand
922    /// back a malformed block — it surfaces `SeqUncertain` (a commit occurred,
923    /// reconcile) and invokes the server exactly once.
924    #[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        // Honour the key but return a count the caller did NOT request — a
931        // protocol violation; counts calls to prove no (double-spending) retry.
932        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    /// A successful `GetSeqResponse` echoing a different `key` than requested is
973    /// likewise a protocol violation → `SeqUncertain`, server invoked once.
974    #[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        // Echo a DIFFERENT key than requested — a protocol violation; counts
981        // calls to prove the malformed success is not retried.
982        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    /// `get_seq_batch` returns one [`SeqBlock`] per entry, in request order,
1023    /// with each block's start/count/epoch matching the server's echoed grant.
1024    /// Uses a fake server that echoes grants positionally from a shared counter.
1025    #[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        // Simple echo server: for each entry, allocate `count` ordinals from a
1032        // shared counter and echo them back in request order.
1033        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        // First entry: "orders" with count 5, starting at 0.
1077        assert_eq!(blocks[0].start, 0);
1078        assert_eq!(blocks[0].count, 5);
1079        assert_eq!(blocks[0].epoch, EPOCH);
1080        // Second entry: "invoices" with count 3, gapless after orders.
1081        assert_eq!(blocks[1].start, 5);
1082        assert_eq!(blocks[1].count, 3);
1083        assert_eq!(blocks[1].epoch, EPOCH);
1084    }
1085
1086    /// Duplicate keys are rejected client-side without hitting the server.
1087    #[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        // Server should never be reached; count calls to prove it.
1095        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        // Duplicate key "orders" — must be caught client-side.
1119        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    /// A malformed success response (wrong grants length, mismatched key/count,
1131    /// or missing epoch) surfaces as `SeqUncertain`. The server is called
1132    /// exactly once — the ambiguity is not retried.
1133    #[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        // Case 1: grants.len() != entries.len() (one fewer grant than requested).
1139        {
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                        // Return only one grant for a two-entry request.
1149                        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        // Case 2: grant.count mismatches the requested count.
1189        {
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                        // Echo the right key but wrong count for the first entry.
1199                        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                                // Tamper: inflate count for the first entry.
1207                                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        // Case 3: missing epoch on an otherwise-valid response.
1246        {
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                        // No epoch — protocol violation.
1264                        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        // Case 4: grant.key mismatches the requested key (the other branch of
1300        // the per-entry `||` guard, distinct from the count-mismatch Case 2).
1301        {
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                        // Echo the right count but tamper the first grant's key.
1311                        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    /// A post-send server error (e.g. `INTERNAL`) from `GetSeqBatch` surfaces as
1362    /// `SeqUncertain` and the server is invoked exactly once — never retried.
1363    #[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}