Skip to main content

corium_peer/
lib.rs

1//! Peer library: remote connection to a transactor, tx-report handling,
2//! local database values, sync, and the segment cache
3//! (see `docs/architecture.md` and `docs/design/protocol.md`).
4//!
5//! A [`Connection`] subscribes to the transactor's tx-report stream and
6//! folds every report into an immutable [`Db`] value locally; queries never
7//! block on the transactor. On disconnect it reconnects and resubscribes
8//! from its basis, and the server backfills the gap from the durable log.
9
10pub mod authz;
11pub mod metrics;
12pub mod segment;
13pub mod server;
14
15use std::collections::BTreeMap;
16use std::sync::{Arc, Mutex, RwLock};
17use std::time::Duration;
18
19use corium_core::{Datom, EntityId, KeywordInterner, Schema};
20use corium_crypt::Keyring;
21use corium_db::{Db, Idents, bootstrap};
22use corium_log::TxRecord;
23use corium_protocol::auth::TokenInterceptor;
24use corium_protocol::codec::{self, CodecError};
25use corium_protocol::pb;
26use corium_protocol::pb::catalog_client::CatalogClient;
27use corium_protocol::pb::transactor_client::TransactorClient;
28use corium_query::edn::Edn;
29use thiserror::Error;
30use tokio::sync::{broadcast, watch};
31use tonic::Status;
32use tonic::service::interceptor::InterceptedService;
33use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
34
35pub use crate::segment::{CachedPeerStorage, DiscoveredPeerStorage, EncryptedPeerStorage};
36use crate::segment::{PeerStorage, SnapshotError, load_current_snapshot, open_encrypted_storage};
37pub use corium_store::{SegmentCacheConfig, SegmentCacheMetrics};
38
39type Client = TransactorClient<InterceptedService<Channel, TokenInterceptor>>;
40
41/// Peer failure.
42#[derive(Debug, Error)]
43pub enum PeerError {
44    /// Transport-level failure.
45    #[error(transparent)]
46    Transport(#[from] tonic::transport::Error),
47    /// RPC failure.
48    #[error(transparent)]
49    Rpc(#[from] Status),
50    /// Payload failed to decode.
51    #[error(transparent)]
52    Codec(#[from] CodecError),
53    /// Published storage snapshot could not be loaded.
54    #[error(transparent)]
55    Snapshot(#[from] SnapshotError),
56    /// Protocol contract violation.
57    #[error("protocol error: {0}")]
58    Protocol(String),
59    /// The connection background task has stopped.
60    #[error("connection closed")]
61    Closed,
62}
63
64/// Connection configuration.
65#[derive(Clone)]
66pub struct ConnectConfig {
67    /// Candidate transactor endpoints in preference order, e.g.
68    /// `http://127.0.0.1:4334`. With an HA pair, list the active and every
69    /// standby: the connection sticks to whichever endpoint accepts its
70    /// subscription and rotates through the rest on failure, so failover
71    /// needs no reconfiguration.
72    pub endpoints: Vec<String>,
73    /// Database name.
74    pub db: String,
75    /// Optional bearer token.
76    pub token: Option<String>,
77    /// Optional TLS configuration (`https` endpoints).
78    pub tls: Option<ClientTlsConfig>,
79    /// Minimum reconnect backoff.
80    pub reconnect_min: Duration,
81    /// Maximum reconnect backoff.
82    pub reconnect_max: Duration,
83    /// How long [`Connection::transact`] keeps retrying failures that are
84    /// provably pre-commit (standby/deposed rejections, connection
85    /// establishment) while a failover completes. Ambiguous failures — a
86    /// connection that died with the request in flight — are surfaced
87    /// immediately; the transaction may or may not have committed.
88    pub failover_timeout: Duration,
89    /// Heartbeat-silence timeout override. `None` derives three times the
90    /// server-advertised heartbeat interval; streams silent for longer are
91    /// dropped and the connection fails over.
92    pub heartbeat_timeout: Option<Duration>,
93    /// Optional direct blob/root storage used to bootstrap from the newest
94    /// published index before subscribing to the transaction-log tail.
95    storage: Option<Arc<dyn PeerStorage>>,
96    /// Keys this peer can resolve. Needed only for a storage-aware peer
97    /// reading an encrypted database: everything arriving over the wire is
98    /// already plaintext.
99    keyring: Option<Arc<dyn Keyring>>,
100    segment_cache_metrics: Option<SegmentCacheMetricsHandle>,
101}
102
103#[derive(Clone)]
104pub(crate) struct SegmentCacheMetricsHandle {
105    pub(crate) metrics: Arc<SegmentCacheMetrics>,
106    pub(crate) disk_capacity: u64,
107    pub(crate) memory_capacity: u64,
108}
109
110impl std::fmt::Debug for ConnectConfig {
111    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        formatter
113            .debug_struct("ConnectConfig")
114            .field("endpoints", &self.endpoints)
115            .field("db", &self.db)
116            .field("token", &self.token.as_ref().map(|_| "[REDACTED]"))
117            .field("tls", &self.tls.is_some())
118            .field("reconnect_min", &self.reconnect_min)
119            .field("reconnect_max", &self.reconnect_max)
120            .field("failover_timeout", &self.failover_timeout)
121            .field("heartbeat_timeout", &self.heartbeat_timeout)
122            .field("storage", &self.storage.is_some())
123            .field("keyring", &self.keyring.is_some())
124            .field("segment_cache", &self.segment_cache_metrics.is_some())
125            .finish()
126    }
127}
128
129impl ConnectConfig {
130    /// Plaintext connection with default backoff.
131    #[must_use]
132    pub fn new(endpoint: impl Into<String>, db: impl Into<String>) -> Self {
133        Self::with_failover(vec![endpoint.into()], db)
134    }
135
136    /// Plaintext connection over an ordered endpoint candidate list.
137    #[must_use]
138    pub fn with_failover(endpoints: Vec<String>, db: impl Into<String>) -> Self {
139        Self {
140            endpoints,
141            db: db.into(),
142            token: None,
143            tls: None,
144            reconnect_min: Duration::from_millis(100),
145            reconnect_max: Duration::from_secs(5),
146            failover_timeout: Duration::from_secs(30),
147            heartbeat_timeout: None,
148            storage: None,
149            keyring: None,
150            segment_cache_metrics: None,
151        }
152    }
153
154    /// Gives this peer direct read access to the transactor's blob/root
155    /// storage service.
156    ///
157    /// A storage-aware connection loads the newest published EAVT snapshot
158    /// and asks the transactor only for transactions after that snapshot.
159    #[must_use]
160    pub fn with_storage(mut self, storage: Arc<dyn PeerStorage>) -> Self {
161        self.storage = Some(storage);
162        self.segment_cache_metrics = None;
163        self
164    }
165
166    /// Gives this peer direct storage access through a bounded local SSD cache.
167    ///
168    /// # Errors
169    /// Returns an error for an invalid, unwritable, or already-owned cache directory.
170    pub fn with_storage_cache(
171        mut self,
172        storage: Arc<dyn PeerStorage>,
173        cache: &SegmentCacheConfig,
174    ) -> std::io::Result<Self> {
175        let cached = CachedPeerStorage::open(storage, cache)?;
176        self.segment_cache_metrics = Some(SegmentCacheMetricsHandle {
177            metrics: cached.metrics(),
178            disk_capacity: cache.capacity_bytes,
179            memory_capacity: cache.memory_capacity_bytes,
180        });
181        self.storage = Some(Arc::new(cached));
182        Ok(self)
183    }
184
185    /// Supplies the keys this peer can resolve.
186    ///
187    /// Only a storage-aware peer needs them, and only for a database that is
188    /// encrypted at rest: without direct storage every datom arrives over the
189    /// wire already decrypted by the transactor's peers.
190    #[must_use]
191    pub fn with_keyring(mut self, keyring: Arc<dyn Keyring>) -> Self {
192        self.keyring = Some(keyring);
193        self
194    }
195
196    /// Whether direct peer storage has been configured.
197    #[must_use]
198    pub fn has_storage(&self) -> bool {
199        self.storage.is_some()
200    }
201}
202
203/// One transaction applied to the peer's local database value.
204#[derive(Clone, Debug)]
205pub struct PeerReport {
206    /// Transaction number.
207    pub t: u64,
208    /// Commit timestamp (Unix milliseconds).
209    pub tx_instant: i64,
210    /// Datoms asserted/retracted by the transaction.
211    pub datoms: Vec<Datom>,
212    /// Database value including the transaction.
213    pub db_after: Db,
214}
215
216/// Result of a transaction submitted through a peer.
217#[derive(Clone, Debug)]
218pub struct TxResult {
219    /// Basis before the transaction.
220    pub basis_before: u64,
221    /// The transaction's `t`.
222    pub basis_t: u64,
223    /// Commit timestamp.
224    pub tx_instant: i64,
225    /// Tempid allocations.
226    pub tempids: BTreeMap<String, EntityId>,
227    /// Database value including the transaction.
228    pub db_after: Db,
229}
230
231struct PeerState {
232    schema: Schema,
233    idents: Idents,
234    interner: KeywordInterner,
235    db: Db,
236    instants: BTreeMap<u64, i64>,
237}
238
239struct Inner {
240    config: ConnectConfig,
241    state: RwLock<Option<PeerState>>,
242    basis: watch::Sender<u64>,
243    index_basis: watch::Sender<u64>,
244    reports: broadcast::Sender<PeerReport>,
245    client: Mutex<Client>,
246    /// Index into `config.endpoints` of the endpoint currently serving the
247    /// subscription (and the cached client).
248    endpoint_index: std::sync::atomic::AtomicUsize,
249    /// Server-advertised heartbeat interval (ms); 0 disables the
250    /// heartbeat-silence timeout.
251    heartbeat_ms: std::sync::atomic::AtomicU64,
252}
253
254impl Inner {
255    /// Deadline of stream silence after which the transactor is presumed
256    /// dead and the connection fails over.
257    fn heartbeat_deadline(&self) -> Option<Duration> {
258        if let Some(timeout) = self.config.heartbeat_timeout {
259            return Some(timeout);
260        }
261        let advertised = self.heartbeat_ms.load(std::sync::atomic::Ordering::Relaxed);
262        (advertised > 0).then(|| Duration::from_millis(advertised.saturating_mul(3)))
263    }
264}
265
266/// A live peer connection to a transactor-hosted database.
267pub struct Connection {
268    inner: Arc<Inner>,
269    task: tokio::task::JoinHandle<()>,
270}
271
272impl Drop for Connection {
273    fn drop(&mut self) {
274        self.task.abort();
275    }
276}
277
278async fn open_channel(config: &ConnectConfig, endpoint: &str) -> Result<Channel, PeerError> {
279    let mut endpoint = Endpoint::from_shared(endpoint.to_owned())
280        .map_err(|error| PeerError::Protocol(format!("bad endpoint: {error}")))?
281        .connect_timeout(Duration::from_secs(10));
282    if let Some(tls) = &config.tls {
283        endpoint = endpoint.tls_config(tls.clone())?;
284    }
285    Ok(endpoint.connect().await?)
286}
287
288fn make_client(channel: Channel, token: Option<String>) -> Client {
289    TransactorClient::with_interceptor(channel, TokenInterceptor::new(token))
290}
291
292impl Connection {
293    pub(crate) fn segment_cache_metrics(&self) -> Option<SegmentCacheMetricsHandle> {
294        self.inner.config.segment_cache_metrics.clone()
295    }
296
297    /// Connects and waits until the handshake and required log tail have
298    /// been applied locally. With direct storage configured, the peer first
299    /// loads the newest published index and subscribes from its basis;
300    /// otherwise it subscribes from basis zero. Candidate endpoints are tried
301    /// in order; a standby transactor rejects the subscription and the next
302    /// candidate is tried.
303    ///
304    /// # Errors
305    /// Returns [`PeerError`] when no endpoint accepts the subscription.
306    pub async fn connect(config: ConnectConfig) -> Result<Self, PeerError> {
307        if config.endpoints.is_empty() {
308            return Err(PeerError::Protocol("no endpoints configured".into()));
309        }
310        // Decryption wraps whatever storage was configured, so it sits above
311        // the segment cache and the SSD tier keeps holding ciphertext.
312        let storage = match &config.storage {
313            Some(storage) => Some(
314                open_encrypted_storage(Arc::clone(storage), &config.db, config.keyring.as_ref())
315                    .await
316                    .map_err(SnapshotError::Store)?,
317            ),
318            None => None,
319        };
320        let snapshot = match &storage {
321            Some(storage) => load_current_snapshot(storage.as_ref(), &config.db).await?,
322            None => None,
323        };
324        let start_basis = snapshot.as_ref().map_or(0, Db::basis_t);
325        // Establish the first subscription before returning so `db()` is
326        // populated and connection errors surface synchronously.
327        let mut last_error = PeerError::Closed;
328        let mut first = None;
329        for (index, endpoint) in config.endpoints.iter().enumerate() {
330            match open_channel(&config, endpoint).await {
331                Ok(channel) => {
332                    let mut client = make_client(channel, config.token.clone());
333                    match subscribe_with(&mut client, &config.db, start_basis).await {
334                        Ok(stream) => {
335                            first = Some((index, client, stream));
336                            break;
337                        }
338                        Err(error) => last_error = error,
339                    }
340                }
341                Err(error) => last_error = error,
342            }
343        }
344        let Some((index, client, mut stream)) = first else {
345            return Err(last_error);
346        };
347        let initial_state = snapshot.map(|db| PeerState {
348            schema: db.schema().clone(),
349            idents: db.idents().clone(),
350            interner: db.interner().clone(),
351            db,
352            instants: BTreeMap::new(),
353        });
354        let inner = Arc::new(Inner {
355            state: RwLock::new(initial_state),
356            basis: watch::channel(start_basis).0,
357            index_basis: watch::channel(0).0,
358            reports: broadcast::channel(1024).0,
359            client: Mutex::new(client),
360            endpoint_index: std::sync::atomic::AtomicUsize::new(index),
361            heartbeat_ms: std::sync::atomic::AtomicU64::new(0),
362            config,
363        });
364        let handshake_basis = pump_handshake(&inner, &mut stream).await?;
365        drain_until(&inner, &mut stream, handshake_basis).await?;
366        let task_inner = Arc::clone(&inner);
367        let task = tokio::spawn(async move {
368            run_loop(task_inner, Some(stream)).await;
369        });
370        Ok(Self { inner, task })
371    }
372
373    /// The connected database name.
374    #[must_use]
375    pub fn db_name(&self) -> &str {
376        &self.inner.config.db
377    }
378
379    /// Returns the current local database value without blocking on the
380    /// transactor.
381    ///
382    /// # Panics
383    /// Panics if called before the initial handshake (impossible through
384    /// [`Connection::connect`]).
385    #[must_use]
386    pub fn db(&self) -> Db {
387        self.inner
388            .state
389            .read()
390            .unwrap_or_else(std::sync::PoisonError::into_inner)
391            .as_ref()
392            .expect("connection is initialized")
393            .db
394            .clone()
395    }
396
397    /// Basis of the newest locally applied transaction.
398    #[must_use]
399    pub fn basis_t(&self) -> u64 {
400        *self.inner.basis.subscribe().borrow()
401    }
402
403    /// Basis of the newest published durable index announced by the
404    /// transactor.
405    #[must_use]
406    pub fn index_basis_t(&self) -> u64 {
407        *self.inner.index_basis.subscribe().borrow()
408    }
409
410    /// Subscribes to reports applied after this call.
411    #[must_use]
412    pub fn tx_reports(&self) -> broadcast::Receiver<PeerReport> {
413        self.inner.reports.subscribe()
414    }
415
416    /// Transaction instants recorded locally for `[start, end)`, paired
417    /// with their datoms (the peer-side `tx-range`).
418    #[must_use]
419    pub fn tx_range(&self, start: u64, end: Option<u64>) -> Vec<TxRecord> {
420        let guard = self
421            .inner
422            .state
423            .read()
424            .unwrap_or_else(std::sync::PoisonError::into_inner);
425        let Some(state) = guard.as_ref() else {
426            return Vec::new();
427        };
428        state
429            .db
430            .tx_range(start, end)
431            .into_iter()
432            .map(|(t, datoms)| TxRecord {
433                t,
434                tx_instant: state.instants.get(&t).copied().unwrap_or_default(),
435                datoms,
436            })
437            .collect()
438    }
439
440    /// Submits a transaction (EDN transaction forms) and waits until it is
441    /// applied locally, so a following [`Connection::db`] observes it.
442    ///
443    /// Failures that are provably pre-commit — a standby or deposed
444    /// transactor rejecting the request, or a connection that could not be
445    /// established — are retried until [`ConnectConfig::failover_timeout`],
446    /// riding out an HA takeover. A connection that dies with the request
447    /// in flight is ambiguous (the transaction may or may not have
448    /// committed) and surfaces as an error, exactly like a transactor
449    /// crash between durability and reply; callers decide whether to check
450    /// and resubmit.
451    ///
452    /// # Errors
453    /// Returns [`PeerError`] for rejected transactions or transport failure.
454    pub async fn transact(&self, forms: Vec<Edn>) -> Result<TxResult, PeerError> {
455        self.transact_at(forms, None).await
456    }
457
458    /// Submits a transaction conditionally against `expected_basis_t`.
459    ///
460    /// # Errors
461    /// Returns [`PeerError`] for a stale basis, rejected transaction, or
462    /// transport failure.
463    pub async fn transact_at(
464        &self,
465        forms: Vec<Edn>,
466        expected_basis_t: Option<u64>,
467    ) -> Result<TxResult, PeerError> {
468        let tx_data = codec::encode_edn(&Edn::Vector(forms));
469        let deadline = tokio::time::Instant::now() + self.inner.config.failover_timeout;
470        let response = loop {
471            match self
472                .transact_raw_at(tx_data.clone(), expected_basis_t)
473                .await
474            {
475                Ok(response) => break response,
476                Err(error) if retry_is_safe(&error) && tokio::time::Instant::now() < deadline => {
477                    tokio::time::sleep(self.inner.config.reconnect_min).await;
478                }
479                Err(error) => return Err(error),
480            }
481        };
482        let tempids = decode_tempids(&response.tempids)?;
483        let db_after = self.sync_to(response.basis_t).await?;
484        Ok(TxResult {
485            basis_before: response.basis_before,
486            basis_t: response.basis_t,
487            tx_instant: response.tx_instant,
488            tempids,
489            db_after,
490        })
491    }
492
493    /// Submits already-encoded transaction data, returning the raw wire
494    /// response (used by the peer server's transact proxy).
495    ///
496    /// # Errors
497    /// Returns [`PeerError`] for rejected transactions or transport failure.
498    pub async fn transact_raw(&self, tx_data: Vec<u8>) -> Result<pb::TransactResponse, PeerError> {
499        self.transact_raw_at(tx_data, None).await
500    }
501
502    /// Submits already-encoded transaction data with an optional basis fence.
503    ///
504    /// # Errors
505    /// Returns [`PeerError`] for a stale basis, rejected transaction, or
506    /// transport failure.
507    pub async fn transact_raw_at(
508        &self,
509        tx_data: Vec<u8>,
510        expected_basis_t: Option<u64>,
511    ) -> Result<pb::TransactResponse, PeerError> {
512        let request = pb::TransactRequest {
513            db: self.inner.config.db.clone(),
514            protocol_version: corium_protocol::PROTOCOL_VERSION,
515            tx_data,
516            expected_basis_t,
517        };
518        let mut client = self
519            .inner
520            .client
521            .lock()
522            .unwrap_or_else(std::sync::PoisonError::into_inner)
523            .clone();
524        Ok(client.transact(request).await?.into_inner())
525    }
526
527    /// Waits until the local basis reaches the transactor's current basis.
528    ///
529    /// # Errors
530    /// Returns [`PeerError`] on transport failure.
531    pub async fn sync(&self) -> Result<Db, PeerError> {
532        let mut client = self
533            .inner
534            .client
535            .lock()
536            .unwrap_or_else(std::sync::PoisonError::into_inner)
537            .clone();
538        let response = client
539            .sync(pb::SyncRequest {
540                db: self.inner.config.db.clone(),
541                t: 0,
542            })
543            .await?
544            .into_inner();
545        self.sync_to(response.basis_t).await
546    }
547
548    /// Waits until the local basis reaches `t`, returning the database value.
549    ///
550    /// # Errors
551    /// Returns [`PeerError::Closed`] if the connection task stops.
552    pub async fn sync_to(&self, t: u64) -> Result<Db, PeerError> {
553        let mut basis = self.inner.basis.subscribe();
554        loop {
555            if *basis.borrow() >= t {
556                return Ok(self.db());
557            }
558            basis.changed().await.map_err(|_| PeerError::Closed)?;
559        }
560    }
561
562    /// Transactor-side status for the connected database.
563    ///
564    /// # Errors
565    /// Returns [`PeerError`] on transport failure.
566    pub async fn status(&self) -> Result<pb::StatusResponse, PeerError> {
567        let mut client = self
568            .inner
569            .client
570            .lock()
571            .unwrap_or_else(std::sync::PoisonError::into_inner)
572            .clone();
573        Ok(client
574            .status(pb::StatusRequest {
575                db: self.inner.config.db.clone(),
576            })
577            .await?
578            .into_inner())
579    }
580
581    /// Opens an independent upstream subscription (used by the peer server
582    /// to relay tx-report streams to thin clients).
583    ///
584    /// # Errors
585    /// Returns [`PeerError`] on transport failure.
586    pub async fn subscribe_raw(
587        &self,
588        from_basis_t: u64,
589    ) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
590        let mut client = self
591            .inner
592            .client
593            .lock()
594            .unwrap_or_else(std::sync::PoisonError::into_inner)
595            .clone();
596        Ok(client
597            .subscribe(pb::SubscribeRequest {
598                db: self.inner.config.db.clone(),
599                protocol_version: corium_protocol::PROTOCOL_VERSION,
600                from_basis_t,
601            })
602            .await?
603            .into_inner())
604    }
605}
606
607/// Whether a transact failure is provably pre-commit and therefore safe to
608/// retry without risking a duplicate transaction.
609fn retry_is_safe(error: &PeerError) -> bool {
610    match error {
611        // The channel could not even be built; no request was sent.
612        PeerError::Transport(_) => true,
613        PeerError::Rpc(status) => match status.code() {
614            // A standby (not lease holder) or freshly deposed transactor
615            // refuses before reaching the commit point.
616            tonic::Code::FailedPrecondition => {
617                let message = status.message();
618                message.contains("standby") || message.contains("deposed")
619            }
620            // Unavailable with a connect-phase source means the request was
621            // never sent. Anything else (a connection that died mid-call)
622            // is ambiguous and must surface.
623            tonic::Code::Unavailable => status.message().contains("connect"),
624            _ => false,
625        },
626        _ => false,
627    }
628}
629
630fn decode_tempids(bytes: &[u8]) -> Result<BTreeMap<String, EntityId>, PeerError> {
631    let Edn::Map(pairs) = codec::decode_edn(bytes)? else {
632        return Err(PeerError::Protocol("tempids must be a map".into()));
633    };
634    let mut tempids = BTreeMap::new();
635    for (key, value) in pairs {
636        let (Edn::Str(name), Edn::Long(raw)) = (key, value) else {
637            return Err(PeerError::Protocol("bad tempid entry".into()));
638        };
639        let raw = u64::try_from(raw).map_err(|_| PeerError::Protocol("bad entity id".into()))?;
640        tempids.insert(name, EntityId::from_raw(raw));
641    }
642    Ok(tempids)
643}
644
645async fn subscribe_with(
646    client: &mut Client,
647    db: &str,
648    from_basis_t: u64,
649) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
650    Ok(client
651        .subscribe(pb::SubscribeRequest {
652            db: db.to_owned(),
653            protocol_version: corium_protocol::PROTOCOL_VERSION,
654            from_basis_t,
655        })
656        .await?
657        .into_inner())
658}
659
660/// Consumes the stream's handshake, installing schema/naming, and returns
661/// the server basis at subscription time.
662async fn pump_handshake(
663    inner: &Arc<Inner>,
664    stream: &mut tonic::Streaming<pb::SubscribeItem>,
665) -> Result<u64, PeerError> {
666    let first = stream
667        .message()
668        .await?
669        .and_then(|item| item.item)
670        .ok_or_else(|| PeerError::Protocol("subscription ended before handshake".into()))?;
671    let pb::subscribe_item::Item::Handshake(handshake) = first else {
672        return Err(PeerError::Protocol(
673            "subscription must begin with a handshake".into(),
674        ));
675    };
676    let local_basis = *inner.basis.subscribe().borrow();
677    if handshake.basis_t < local_basis {
678        return Err(PeerError::Protocol(format!(
679            "published snapshot basis {local_basis} is newer than transactor basis {}",
680            handshake.basis_t
681        )));
682    }
683    let (schema, idents) = codec::decode_schema(&handshake.schema)?;
684    inner.heartbeat_ms.store(
685        handshake.heartbeat_interval_ms,
686        std::sync::atomic::Ordering::Relaxed,
687    );
688    {
689        let mut guard = inner
690            .state
691            .write()
692            .unwrap_or_else(std::sync::PoisonError::into_inner);
693        if let Some(state) = guard.as_mut() {
694            // Reconnect: schema/idents are fixed after bootstrap; keep the
695            // locally accumulated database value and naming.
696            state.schema = schema;
697            state.idents = idents;
698        } else {
699            let interner = KeywordInterner::default();
700            let db = Db::new(schema.clone()).with_naming(idents.clone(), interner.clone());
701            *guard = Some(PeerState {
702                schema,
703                idents,
704                interner,
705                db,
706                instants: BTreeMap::new(),
707            });
708        }
709    }
710    let _ = inner.index_basis.send_replace(handshake.index_basis_t);
711    Ok(handshake.basis_t)
712}
713
714/// Applies stream items until the local basis reaches `target`.
715async fn drain_until(
716    inner: &Arc<Inner>,
717    stream: &mut tonic::Streaming<pb::SubscribeItem>,
718    target: u64,
719) -> Result<(), PeerError> {
720    while *inner.basis.subscribe().borrow() < target {
721        let Some(item) = stream.message().await? else {
722            return Err(PeerError::Protocol(
723                "subscription ended during backfill".into(),
724            ));
725        };
726        if let Some(item) = item.item {
727            apply_item(inner, item)?;
728        }
729    }
730    Ok(())
731}
732
733fn apply_item(inner: &Arc<Inner>, item: pb::subscribe_item::Item) -> Result<(), PeerError> {
734    match item {
735        pb::subscribe_item::Item::Report(report) => {
736            let peer_report = {
737                let mut guard = inner
738                    .state
739                    .write()
740                    .unwrap_or_else(std::sync::PoisonError::into_inner);
741                let state = guard
742                    .as_mut()
743                    .ok_or_else(|| PeerError::Protocol("report before handshake".into()))?;
744                if report.t <= state.db.basis_t() {
745                    return Ok(());
746                }
747                let before = state.interner.len();
748                let datoms = codec::decode_datoms(&report.datoms, &mut state.interner)?;
749                let tx_instant =
750                    bootstrap::asserted_instant(report.t, &datoms).unwrap_or(report.tx_instant);
751                if state.interner.len() > before {
752                    state.db = state
753                        .db
754                        .clone()
755                        .with_naming(state.idents.clone(), state.interner.clone());
756                }
757                state.db = state.db.with_transaction_at(report.t, tx_instant, &datoms);
758                state.instants.insert(report.t, tx_instant);
759                PeerReport {
760                    t: report.t,
761                    tx_instant,
762                    datoms,
763                    db_after: state.db.clone(),
764                }
765            };
766            let _ = inner.basis.send_replace(peer_report.t);
767            let _ = inner.reports.send(peer_report);
768        }
769        pb::subscribe_item::Item::IndexBasis(index) => {
770            let _ = inner.index_basis.send_replace(index.index_basis_t);
771        }
772        pb::subscribe_item::Item::Heartbeat(_) | pb::subscribe_item::Item::Handshake(_) => {}
773    }
774    Ok(())
775}
776
777/// Long-running consume/reconnect loop: on stream end, error, or heartbeat
778/// silence past the timeout, rebuilds the channel with exponential backoff
779/// and resubscribes from the local basis; the server backfills the gap.
780/// Reconnection rotates through the candidate endpoints, so when the
781/// active transactor dies the loop lands on whichever standby takes over
782/// the lease (a standby rejects the subscription until then).
783async fn run_loop(inner: Arc<Inner>, initial: Option<tonic::Streaming<pb::SubscribeItem>>) {
784    let mut stream = initial;
785    let mut backoff = inner.config.reconnect_min;
786    let mut candidate = inner
787        .endpoint_index
788        .load(std::sync::atomic::Ordering::Relaxed);
789    loop {
790        if let Some(active) = stream.as_mut() {
791            let next = match inner.heartbeat_deadline() {
792                Some(deadline) => match tokio::time::timeout(deadline, active.message()).await {
793                    Ok(next) => next,
794                    // Heartbeat silence: the transactor is presumed dead
795                    // even though the transport has not noticed (partition,
796                    // stalled process); drop the stream and fail over.
797                    Err(_elapsed) => Ok(None),
798                },
799                None => active.message().await,
800            };
801            match next {
802                Ok(Some(item)) => {
803                    backoff = inner.config.reconnect_min;
804                    if let Some(item) = item.item
805                        && apply_item(&inner, item).is_err()
806                    {
807                        stream = None;
808                    }
809                    continue;
810                }
811                Ok(None) | Err(_) => {
812                    stream = None;
813                }
814            }
815        }
816        tokio::time::sleep(backoff).await;
817        backoff = (backoff * 2).min(inner.config.reconnect_max);
818        let endpoints = &inner.config.endpoints;
819        let endpoint = &endpoints[candidate % endpoints.len()];
820        let Ok(channel) = open_channel(&inner.config, endpoint).await else {
821            candidate = (candidate + 1) % endpoints.len();
822            continue;
823        };
824        let mut client = make_client(channel, inner.config.token.clone());
825        let from = *inner.basis.subscribe().borrow();
826        match subscribe_with(&mut client, &inner.config.db, from).await {
827            Ok(mut fresh) => {
828                if pump_handshake(&inner, &mut fresh).await.is_ok() {
829                    // Sticky success: transact/status/sync now go here too.
830                    *inner
831                        .client
832                        .lock()
833                        .unwrap_or_else(std::sync::PoisonError::into_inner) = client;
834                    inner.endpoint_index.store(
835                        candidate % endpoints.len(),
836                        std::sync::atomic::Ordering::Relaxed,
837                    );
838                    stream = Some(fresh);
839                } else {
840                    candidate = (candidate + 1) % endpoints.len();
841                }
842            }
843            Err(_) => {
844                candidate = (candidate + 1) % endpoints.len();
845            }
846        }
847    }
848}
849
850/// Catalog (admin) client for a transactor endpoint.
851pub struct Admin {
852    client: CatalogClient<InterceptedService<Channel, TokenInterceptor>>,
853}
854
855impl Admin {
856    /// Connects to a transactor's catalog service.
857    ///
858    /// # Errors
859    /// Returns [`PeerError`] when the endpoint is unreachable.
860    pub async fn connect(
861        endpoint: &str,
862        token: Option<String>,
863        tls: Option<ClientTlsConfig>,
864    ) -> Result<Self, PeerError> {
865        let config = ConnectConfig {
866            tls,
867            token: token.clone(),
868            ..ConnectConfig::new(endpoint, "")
869        };
870        let channel = open_channel(&config, endpoint).await?;
871        Ok(Self {
872            client: CatalogClient::with_interceptor(channel, TokenInterceptor::new(token)),
873        })
874    }
875
876    /// Creates a database with EDN schema forms; `false` when it existed.
877    ///
878    /// # Errors
879    /// Returns [`PeerError`] for invalid schema or transport failure.
880    pub async fn create_database(&mut self, db: &str, schema: &[Edn]) -> Result<bool, PeerError> {
881        self.create_database_with_key(db, schema, None).await
882    }
883
884    /// Creates a database, optionally encrypted at rest under `storage_key`.
885    ///
886    /// `storage_key` is a key-encryption key URI the *transactor* resolves —
887    /// it is where the database's data keys are wrapped, and no key material
888    /// crosses this connection. It is fixed at creation: a database made
889    /// without one can only become encrypted through a backup and restore.
890    ///
891    /// # Errors
892    /// Returns [`PeerError`] for invalid schema, a key the transactor cannot
893    /// resolve, or transport failure.
894    pub async fn create_database_with_key(
895        &mut self,
896        db: &str,
897        schema: &[Edn],
898        storage_key: Option<&str>,
899    ) -> Result<bool, PeerError> {
900        let response = self
901            .client
902            .create_database(pb::CreateDatabaseRequest {
903                db: db.to_owned(),
904                schema: codec::encode_edn(&Edn::Vector(schema.to_vec())),
905                storage_key: storage_key.unwrap_or_default().to_owned(),
906            })
907            .await?;
908        Ok(response.into_inner().created)
909    }
910
911    /// Reports a database's storage-encryption keys and their epochs.
912    ///
913    /// # Errors
914    /// Returns [`PeerError`] on transport failure or an unknown database.
915    pub async fn key_status(&mut self, db: &str) -> Result<pb::KeyStatusResponse, PeerError> {
916        Ok(self
917            .client
918            .key_status(pb::KeyStatusRequest { db: db.to_owned() })
919            .await?
920            .into_inner())
921    }
922
923    /// Opens a new storage-key epoch that new writes use, rewriting nothing.
924    ///
925    /// # Errors
926    /// Returns [`PeerError`] on transport failure, an unknown or unencrypted
927    /// database, or a key the transactor cannot wrap under.
928    pub async fn rotate_storage_key(&mut self, db: &str) -> Result<u32, PeerError> {
929        Ok(self
930            .client
931            .rotate_storage_key(pb::RotateStorageKeyRequest { db: db.to_owned() })
932            .await?
933            .into_inner()
934            .epoch)
935    }
936
937    /// Re-wraps every data key under `kek`. No stored data is touched.
938    ///
939    /// # Errors
940    /// Returns [`PeerError`] on transport failure, an unknown or unencrypted
941    /// database, or a KEK the transactor cannot resolve.
942    pub async fn rewrap_keys(&mut self, db: &str, kek: &str) -> Result<(), PeerError> {
943        self.client
944            .rewrap_keys(pb::RewrapKeysRequest {
945                db: db.to_owned(),
946                kek: kek.to_owned(),
947            })
948            .await?;
949        Ok(())
950    }
951
952    /// Forks `db` into a new database `target` duplicating it at
953    /// transaction basis `as_of_t` (`None` forks at the current basis).
954    /// Returns the fork's basis, or `None` when the target already existed.
955    ///
956    /// # Errors
957    /// Returns [`PeerError`] on transport failure, an unknown source, an
958    /// invalid target name, or a basis ahead of the source's.
959    pub async fn fork_database(
960        &mut self,
961        db: &str,
962        target: &str,
963        as_of_t: Option<u64>,
964    ) -> Result<Option<u64>, PeerError> {
965        let response = self
966            .client
967            .fork_database(pb::ForkDatabaseRequest {
968                db: db.to_owned(),
969                target: target.to_owned(),
970                as_of_t: as_of_t.unwrap_or(0),
971            })
972            .await?
973            .into_inner();
974        Ok(response.created.then_some(response.basis_t))
975    }
976
977    /// Deletes a database; `false` when it did not exist.
978    ///
979    /// # Errors
980    /// Returns [`PeerError`] on transport failure.
981    pub async fn delete_database(&mut self, db: &str) -> Result<bool, PeerError> {
982        let response = self
983            .client
984            .delete_database(pb::DeleteDatabaseRequest { db: db.to_owned() })
985            .await?;
986        Ok(response.into_inner().deleted)
987    }
988
989    /// Lists hosted databases.
990    ///
991    /// # Errors
992    /// Returns [`PeerError`] on transport failure.
993    pub async fn list_databases(&mut self) -> Result<Vec<String>, PeerError> {
994        let response = self
995            .client
996            .list_databases(pb::ListDatabasesRequest {})
997            .await?;
998        Ok(response.into_inner().dbs)
999    }
1000
1001    /// Sweeps blobs unreachable from every live database root.
1002    ///
1003    /// # Errors
1004    /// Returns [`PeerError`] on transport failure.
1005    pub async fn gc_deleted_databases(&mut self) -> Result<u64, PeerError> {
1006        self.gc_deleted_databases_with_retention(None).await
1007    }
1008
1009    /// Sweeps unreachable blobs with an optional minimum retention window.
1010    ///
1011    /// # Errors
1012    /// Returns [`PeerError`] on transport failure.
1013    pub async fn gc_deleted_databases_with_retention(
1014        &mut self,
1015        retention: Option<Duration>,
1016    ) -> Result<u64, PeerError> {
1017        let response = self
1018            .client
1019            .gc_deleted_databases(pb::GcDeletedDatabasesRequest {
1020                retention_millis: retention.map(duration_millis),
1021            })
1022            .await?;
1023        Ok(response.into_inner().swept_blobs)
1024    }
1025
1026    /// Asks the transactor to publish `db`'s covering indexes now,
1027    /// bypassing its pacing policy; returns the resulting index basis.
1028    ///
1029    /// # Errors
1030    /// Returns [`PeerError`] on transport failure or an unknown database.
1031    pub async fn request_index(&mut self, db: &str) -> Result<u64, PeerError> {
1032        let response = self
1033            .client
1034            .request_index(pb::RequestIndexRequest { db: db.to_owned() })
1035            .await?;
1036        Ok(response.into_inner().index_basis_t)
1037    }
1038
1039    /// Overrides `db`'s index-publication pacing at runtime; `None` fields
1040    /// are left unchanged, so an all-`None` update reads the current
1041    /// policy. Returns the policy now in effect.
1042    ///
1043    /// # Errors
1044    /// Returns [`PeerError`] on transport failure or an unknown database.
1045    pub async fn set_index_policy(
1046        &mut self,
1047        db: &str,
1048        update: IndexPolicySettings,
1049    ) -> Result<IndexPolicySettings, PeerError> {
1050        let response = self
1051            .client
1052            .set_index_policy(pb::SetIndexPolicyRequest {
1053                db: db.to_owned(),
1054                interval_ms: update.interval_ms,
1055                backoff: update.backoff,
1056                tail_threshold: update.tail_threshold,
1057                tail_deadline_ms: update.tail_deadline_ms,
1058            })
1059            .await?
1060            .into_inner();
1061        Ok(IndexPolicySettings {
1062            interval_ms: Some(response.interval_ms),
1063            backoff: Some(response.backoff),
1064            tail_threshold: Some(response.tail_threshold),
1065            tail_deadline_ms: Some(response.tail_deadline_ms),
1066        })
1067    }
1068
1069    /// Fixes the current transaction basis and returns an independent
1070    /// connection to the transactor's underlying storage service, for a
1071    /// backup client or a storage-aware peer bootstrapping from storage.
1072    ///
1073    /// # Errors
1074    /// Returns [`PeerError`] for an unknown database or transport failure.
1075    pub async fn get_storage_info(
1076        &mut self,
1077        db: &str,
1078    ) -> Result<pb::GetStorageInfoResponse, PeerError> {
1079        Ok(self
1080            .client
1081            .get_storage_info(pb::GetStorageInfoRequest { db: db.to_owned() })
1082            .await?
1083            .into_inner())
1084    }
1085}
1086
1087/// Index-publication pacing fields for [`Admin::set_index_policy`]: `None`
1088/// in a request leaves the field unchanged; responses set every field to
1089/// the effective policy.
1090#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1091pub struct IndexPolicySettings {
1092    /// Base interval between publications (ms).
1093    pub interval_ms: Option<u64>,
1094    /// Minimum wait before the next publication, as a multiple of the
1095    /// previous publication's duration (0 disables the backoff).
1096    pub backoff: Option<u32>,
1097    /// Pending-datom count below which a due publication is deferred
1098    /// (0 publishes any pending work).
1099    pub tail_threshold: Option<u64>,
1100    /// Longest a below-threshold tail defers publication (ms).
1101    pub tail_deadline_ms: Option<u64>,
1102}
1103
1104fn duration_millis(duration: Duration) -> u64 {
1105    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1106}
1107
1108#[cfg(test)]
1109mod tests {
1110    use super::*;
1111
1112    #[test]
1113    fn gc_retention_wire_value_preserves_zero_and_subseconds() {
1114        assert_eq!(duration_millis(Duration::ZERO), 0);
1115        assert_eq!(duration_millis(Duration::from_millis(500)), 500);
1116    }
1117}