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 tokio::sync::OnceCell;
54use tonic::transport::Channel;
55use tsoracle_core::{Epoch, LOGICAL_MAX, Timestamp};
56use tsoracle_proto::v1::tso_service_client::TsoServiceClient;
57
58/// The server's per-call cap on requested timestamps, fixed by the 18-bit
59/// logical width. Callers asking for more than this can't be served by any
60/// single RPC; the client rejects them up-front rather than burning a queue
61/// slot and round-trip to learn the same thing from the server.
62pub(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    /// Override the default [`RetryPolicy`].
99    ///
100    /// The policy controls per-attempt deadlines, the overall deadline
101    /// across all candidate endpoints, the cap on attempts, and the
102    /// jittered backoff base. The per-attempt deadline is also pushed
103    /// down to `tonic::transport::Endpoint::connect_timeout` and
104    /// `Endpoint::timeout` for the built-in default and TLS transport
105    /// paths so a blackholed peer fails fast at the transport layer.
106    /// User-supplied [`Self::channel_connector`] closures own their
107    /// own Endpoint config; the policy still bounds the retry loop's
108    /// outer `tokio::time::timeout` around them.
109    pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
110        self.retry_policy = policy;
111        self
112    }
113
114    /// Configure the client to dial bare endpoints with TLS. Bare `host:port`
115    /// becomes `https://host:port`; explicit `http://...` endpoints supplied
116    /// in [`Self::endpoints`] remain plaintext; explicit `https://...`
117    /// endpoints use the provided TLS config.
118    ///
119    /// Wire-supplied `http://...` leader-hint trailers are NOT honored under
120    /// `tls_config` — they are dropped to prevent a contacted peer from
121    /// downgrading the transport. Operator-supplied configuration still wins;
122    /// untrusted wire input does not.
123    ///
124    /// Setting both [`Self::channel_connector`] and `tls_config` is allowed;
125    /// the last call wins (standard builder semantics). Calling
126    /// `channel_connector` after `tls_config` also clears the
127    /// reject-plaintext-hint policy, since the caller-owned connector owns
128    /// its own scheme policy.
129    #[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    /// Replace the default plaintext channel construction with a caller-owned
140    /// closure. The closure is invoked on first use of each endpoint —
141    /// configured endpoints and leader-hint redirects alike. Errors returned
142    /// from the closure surface as [`ClientError::Connector`].
143    ///
144    /// See module docs for the interaction with [`Self::tls_config`]
145    /// (last-wins) and the scheme matrix. A caller-owned connector replaces
146    /// the built-in TLS plumbing entirely, including the
147    /// reject-plaintext-leader-hint policy — the closure is responsible for
148    /// whatever scheme policy it wants to enforce.
149    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    /// The endpoint the client currently believes is the leader, or `None`
198    /// if no leader has been observed yet or the cached entry has aged past
199    /// the configured `leader_ttl`.
200    ///
201    /// Read-only diagnostic surface for ops dashboards and integration tests
202    /// asserting that a client has converged to the expected leader. It
203    /// reflects the cache as last updated by a completed RPC — it neither
204    /// triggers nor waits on any network round-trip, and the TTL check is
205    /// lazy (an expired entry reads as `None`).
206    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    /// Request a contiguous block of `count` dense ordinals for the given
222    /// sequence `key`.
223    ///
224    /// **Non-idempotent:** if this call returns `ClientError::SeqUncertain` it
225    /// means a commit may or may not have occurred on the server — do NOT
226    /// silently retry, as that would risk spending a second block. The caller
227    /// must resolve the ambiguity (e.g. by reading back the current counter
228    /// with a coordinator-level read) before re-issuing.
229    ///
230    /// All other errors are pre-commit-certain (the server rejected the request
231    /// before any durable advance) and are safe to retry.
232    ///
233    /// The client only pre-rejects universally-invalid inputs — an empty/oversized
234    /// `key` and a zero `count` (a block always covers at least one ordinal).
235    /// The upper bound on `count` is the server's configured
236    /// `ServerBuilder::max_seq_count`, which is deployment-specific, so the client
237    /// does not second-guess it: an over-cap `count` is forwarded and the server
238    /// rejects it with `INVALID_ARGUMENT` (surfaced as `ClientError::Rpc`). This
239    /// keeps a client built against one cap correct against a server configured
240    /// with another.
241    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    /// Atomically request contiguous dense blocks for several DISTINCT keys in
252    /// one round-trip. Returns one [`SeqBlock`] per entry, in request order.
253    ///
254    /// **Non-idempotent (whole batch):** on `ClientError::SeqUncertain` the
255    /// entire batch may or may not have committed — do NOT retry; reconcile
256    /// first (the atomic server contract means it is all-or-nothing, so there is
257    /// never a partial mix to untangle). All other errors are pre-commit-certain.
258    ///
259    /// Client-side pre-checks reject only universally-invalid input: an empty
260    /// batch, a duplicate key, an empty/oversized key, or a zero count. The
261    /// per-entry count cap and the batch-size cap are server-side and are
262    /// forwarded for the server to reject, so a client built against one
263    /// configuration stays correct against a server with another.
264    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                // Distinct-key contract — mirror the server's pre-commit reject.
281                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    /// Read the leader's current safe-point in physical-millisecond units.
309    ///
310    /// Targets the cached leader if known, otherwise the first configured
311    /// endpoint. Followers return `MaxSafe::max_safe_physical_ms == 0` rather
312    /// than erroring, matching the proto contract; pollers needing freshness
313    /// should target the leader endpoint.
314    ///
315    /// Single-shot by design — the proto contract is "followers return 0"
316    /// rather than NOT_LEADER, so there is no hint to chase and no worklist
317    /// to drain; a caller polling for freshness retries the next tick rather
318    /// than the next endpoint. The one `(connect, RPC)` pair is bounded by
319    /// [`RetryPolicy::per_attempt_deadline`] (shared across both phases via
320    /// `PairBudget`, exactly like one `get_ts` attempt), and a transport-class
321    /// failure evicts the cached channel so a half-open / black-holing
322    /// connection is dropped before the next poll lands on it.
323    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        // One budget shared across connect + RPC: a slow connect eats into
331        // the RPC's time rather than each phase getting a fresh full budget,
332        // so a single call never runs longer than ~`per_attempt_deadline`.
333        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                // `client_with_cell` already evicts its own cell on a failed
338                // dial, so there is nothing to evict here.
339                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            // A timed-out RPC surfaces as `DeadlineExceeded` (transport-class
358            // per `is_transport_failure`), so the eviction tail below drops the
359            // possibly-half-open channel — matching the `get_ts` attempt path.
360            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    /// Acquire a stamping lease for an opaque holder group.
372    ///
373    /// Single-attempt control-plane call: callers own retry and re-route policy.
374    /// Retrying an ambiguous acquire is safe for the same `(holder,
375    /// holder_epoch)` because the server treats that pair idempotently while
376    /// the lease is live.
377    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    /// Renew a live lease, re-arming its acquire-time TTL.
414    ///
415    /// Single-attempt control-plane call; callers own retry and re-route policy.
416    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    /// Release a lease. The server treats unknown or already-released leases as
442    /// success, making graceful surrender idempotent.
443    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    /// Read the lease-aware safe frontier in physical-millisecond units.
462    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/// The leader's view of the durable safe-point, returned by
488/// [`Client::get_current_max_safe`].
489#[derive(Copy, Clone, Debug, PartialEq, Eq)]
490pub struct MaxSafe {
491    /// Safe-point in physical-millisecond units; `0` is the cold-start sentinel
492    /// (also the value any follower returns).
493    pub max_safe_physical_ms: u64,
494    /// Leader epoch that issued this view.
495    pub epoch: Epoch,
496}
497
498/// A granted stamping lease.
499#[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/// A renewed stamping lease.
508#[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/// Lease-aware safe frontier.
516#[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        // A freshly built client has issued no RPC, so the channel pool's
537        // leader cache is empty and `cached_leader()` reports `None`. This
538        // pins the "nothing observed yet" branch of the diagnostic accessor
539        // without needing a server; the post-RPC `Some(addr)` case is covered
540        // end-to-end in `tsoracle-tests`.
541        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        // Validation prevents a Client whose `pool` has no endpoints to try;
550        // every RPC would fail-fast with `NoReachableEndpoints` and burn no
551        // network roundtrips at all, so reject up-front instead.
552        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    // Marker payload used by both last-wins tests. The free function holds
579    // the body in one source location so coverage credit flows from the test
580    // where the closure DOES run; the test that asserts "this never runs"
581    // just calls the same helper.
582    #[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        // channel_connector is set LAST, so its path runs on get_ts. The
591        // marker error surfaces as `ClientError::Connector`, proving the
592        // builder did not silently keep the prior tls_config.
593        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        // tls_config is set LAST, so the connector path is replaced and its
609        // marker error must NOT surface. The tls_config path produces a
610        // transport-level failure (or NoReachableEndpoints / Rpc) instead.
611        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        // The builder's `batch_flush_interval` knob feeds the driver's
626        // coalescing window; without a test it could silently revert to the
627        // default and no-one would notice from black-box behavior. We
628        // confirm the override path by reaching into the builder fields.
629        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        // The builder field is what `build` hands to the pool and retry
638        // loop. If `retry_policy()` ever silently stops storing the
639        // override, the loop reverts to defaults; this test pins the
640        // override path against that.
641        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    /// Acceptance criterion for the `get_current_max_safe` deadline fix
664    /// (security finding 8c3ea943): a single-shot call against a
665    /// black-holed endpoint must surface `DeadlineExceeded` within the
666    /// configured `per_attempt_deadline`, not park indefinitely on a
667    /// connector whose future never resolves. Exercised through the
668    /// public `channel_connector` surface — `std::future::pending()`
669    /// guarantees the connect future never completes, so any code path
670    /// without an outer `tokio::time::timeout` hangs until an external
671    /// test timeout kills the runtime. The wall-clock bound is generous
672    /// enough to absorb CI scheduler jitter but well below the OS-level
673    /// TCP timeout, mirroring the structure of the `get_ts`
674    /// equivalent below.
675    #[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                // The future never resolves; only an outer timeout can end
687                // the await. This is the exact attack shape the finding
688                // describes: a user-supplied connector (or a black-holed
689                // peer behind one) for which the caller relied on
690                // `RetryPolicy` to bound the wait.
691                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        // Outer safety timeout: when the fix is missing, `get_current_max_safe`
698        // awaits a never-resolving connector indefinitely, which would hang the
699        // whole test runner rather than fail this case. The 5s ceiling converts
700        // that hang into a clean panic with a diagnostic message. Under a
701        // correctly-deadlined call this outer guard never fires — the inner
702        // future returns in ~100ms, well below the elapsed-time assertion.
703        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        // End-to-end test of the issue's acceptance criterion: with no
728        // listener bound at the configured endpoints, a `get_ts` call
729        // must return well before the OS-default TCP timeout (~75 s on
730        // Linux). The bound here is generous enough to absorb CI
731        // scheduler jitter — the assertion is "fast", not "exactly the
732        // configured deadline".
733        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    /// Integration test: `get_seq` against a real loopback gRPC server that
760    /// simulates a file-driver leader. The server uses a fake `TsoService`
761    /// with a real `get_seq` implementation (incrementing an atomic counter).
762    /// Asserts: consecutive blocks are gapless, the epoch is propagated, and
763    /// invalid inputs are rejected up-front.
764    #[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        // A leader serving gapless blocks from a single atomic counter, rejecting
773        // empty keys / zero counts as InvalidArgument.
774        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        // First call: block [0, 5).
813        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        // Second call: block [5, 8) — gapless continuation.
822        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        // Empty key is rejected up-front by the client, not the server.
831        match client.get_seq("", 1).await {
832            Err(ClientError::InvalidSeqKey) => {}
833            other => panic!("empty key must return InvalidSeqKey, got {other:?}"),
834        }
835
836        // Zero count is rejected up-front by the client.
837        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    /// The upper bound on `count` is the server's configured `max_seq_count`, not
844    /// a fixed client-side constant, so the client must forward a large `count`
845    /// rather than short-circuit it with a local `InvalidCount`. Here a count far
846    /// above the old default reaches a cap-less echo server and is served —
847    /// proving the client dropped its hard-coded ceiling check.
848    #[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        // Far above the old hard-coded default — would have been a local
876        // InvalidCount before; now it is forwarded and the echo server serves it.
877        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    /// The rejection companion to `get_seq_forwards_large_count_to_server`: when
886    /// the server's configured `ServerBuilder::max_seq_count` is smaller than the
887    /// requested `count`, it rejects with `INVALID_ARGUMENT`
888    /// (`CoreError::SeqCountTooLarge`). Because the client no longer pre-checks the
889    /// upper bound, that rejection must surface through `Client::get_seq` as
890    /// `ClientError::Rpc` preserving the server's message — exactly as the `get_seq`
891    /// doc contract promises — and must NOT be mislabelled `InvalidSeqKey`: a valid
892    /// key was supplied; only the count was over-cap.
893    #[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        // Mirrors the server's SeqCountTooLarge → INVALID_ARGUMENT mapping for a
900        // count above the configured max_seq_count. Counts calls to prove the
901        // definitive-error path stays bounded (no unbounded ride-out).
902        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        // A valid key with an over-cap count: the key is fine, so InvalidSeqKey
928        // would be the wrong label — the docs promise ClientError::Rpc.
929        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        // Definitive-error path: bounded attempts, never an unbounded ride-out.
946        let n = calls.load(Ordering::SeqCst);
947        assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
948    }
949
950    /// Regression guard for the dense classification bug: a bare
951    /// `FailedPrecondition` from the server (NO leader-hint trailer) — e.g.
952    /// `SeqOverflow` — must surface immediately as that definitive error,
953    /// preserving the server's message, NOT be misread as a NOT_LEADER
954    /// "no leader yet" and ridden out across passes.
955    #[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        // Bare FailedPrecondition with NO leader-hint trailer — mirrors the
962        // server's SeqOverflow mapping. Counts calls to prove bounded attempts.
963        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        // Definitive-error path: bounded attempts, never an unbounded ride-out.
999        let n = calls.load(Ordering::SeqCst);
1000        assert!((1..=3).contains(&n), "expected bounded attempts, got {n}");
1001    }
1002
1003    /// The non-idempotency contract: a post-send transport failure
1004    /// (`Unavailable` after the request reached the server) is ambiguous — the
1005    /// advance may have committed — so `get_seq` returns `SeqUncertain` and must
1006    /// NOT transparently retry (a retry could silently spend a second block).
1007    /// Asserts the server is invoked exactly once.
1008    #[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        // A post-send transport failure (Unavailable after the request reached
1015        // the server) — ambiguous; counts calls to prove no retry.
1016        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    /// The non-idempotency contract extended to application faults: a post-send
1051    /// `INTERNAL` (the server's mapping for a permanent dense driver fault — the
1052    /// file driver emits this for a directory-fsync failure that happens *after*
1053    /// the durable rename) is ambiguous, not pre-commit-certain. `get_seq` must
1054    /// return `SeqUncertain` and invoke the server exactly once, never surfacing
1055    /// it as a definitive (caller-retry-safe) error.
1056    #[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        // A post-send INTERNAL (permanent dense driver fault) — ambiguous; counts
1063        // calls to prove no retry.
1064        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    /// A successful `GetSeqResponse` whose `count` does not match the request is
1099    /// a protocol violation: the server committed an advance, but the block it
1100    /// describes is not the one the caller asked for. The client must not hand
1101    /// back a malformed block — it surfaces `SeqUncertain` (a commit occurred,
1102    /// reconcile) and invokes the server exactly once.
1103    #[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        // Honour the key but return a count the caller did NOT request — a
1110        // protocol violation; counts calls to prove no (double-spending) retry.
1111        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    /// A successful `GetSeqResponse` echoing a different `key` than requested is
1152    /// likewise a protocol violation → `SeqUncertain`, server invoked once.
1153    #[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        // Echo a DIFFERENT key than requested — a protocol violation; counts
1160        // calls to prove the malformed success is not retried.
1161        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    /// `get_seq_batch` returns one [`SeqBlock`] per entry, in request order,
1202    /// with each block's start/count/epoch matching the server's echoed grant.
1203    /// Uses a fake server that echoes grants positionally from a shared counter.
1204    #[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        // Simple echo server: for each entry, allocate `count` ordinals from a
1211        // shared counter and echo them back in request order.
1212        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        // First entry: "orders" with count 5, starting at 0.
1256        assert_eq!(blocks[0].start, 0);
1257        assert_eq!(blocks[0].count, 5);
1258        assert_eq!(blocks[0].epoch, EPOCH);
1259        // Second entry: "invoices" with count 3, gapless after orders.
1260        assert_eq!(blocks[1].start, 5);
1261        assert_eq!(blocks[1].count, 3);
1262        assert_eq!(blocks[1].epoch, EPOCH);
1263    }
1264
1265    /// Duplicate keys are rejected client-side without hitting the server.
1266    #[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        // Server should never be reached; count calls to prove it.
1274        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        // Duplicate key "orders" — must be caught client-side.
1298        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    /// A malformed success response (wrong grants length, mismatched key/count,
1310    /// or missing epoch) surfaces as `SeqUncertain`. The server is called
1311    /// exactly once — the ambiguity is not retried.
1312    #[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        // Case 1: grants.len() != entries.len() (one fewer grant than requested).
1318        {
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                        // Return only one grant for a two-entry request.
1328                        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        // Case 2: grant.count mismatches the requested count.
1368        {
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                        // Echo the right key but wrong count for the first entry.
1378                        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                                // Tamper: inflate count for the first entry.
1386                                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        // Case 3: missing epoch on an otherwise-valid response.
1425        {
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                        // No epoch — protocol violation.
1443                        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        // Case 4: grant.key mismatches the requested key (the other branch of
1479        // the per-entry `||` guard, distinct from the count-mismatch Case 2).
1480        {
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                        // Echo the right count but tamper the first grant's key.
1490                        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    /// A post-send server error (e.g. `INTERNAL`) from `GetSeqBatch` surfaces as
1541    /// `SeqUncertain` and the server is invoked exactly once — never retried.
1542    #[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}