1pub mod metrics;
11pub mod segment;
12pub mod server;
13
14use std::collections::BTreeMap;
15use std::sync::{Arc, Mutex, RwLock};
16use std::time::Duration;
17
18use corium_core::{Datom, EntityId, KeywordInterner, Schema};
19use corium_db::{Db, Idents};
20use corium_log::TxRecord;
21use corium_protocol::auth::TokenInterceptor;
22use corium_protocol::codec::{self, CodecError};
23use corium_protocol::pb;
24use corium_protocol::pb::catalog_client::CatalogClient;
25use corium_protocol::pb::transactor_client::TransactorClient;
26use corium_query::edn::Edn;
27use thiserror::Error;
28use tokio::sync::{broadcast, watch};
29use tonic::Status;
30use tonic::service::interceptor::InterceptedService;
31use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
32
33use crate::segment::{PeerStorage, SnapshotError, load_current_snapshot};
34
35type Client = TransactorClient<InterceptedService<Channel, TokenInterceptor>>;
36
37#[derive(Debug, Error)]
39pub enum PeerError {
40 #[error(transparent)]
42 Transport(#[from] tonic::transport::Error),
43 #[error(transparent)]
45 Rpc(#[from] Status),
46 #[error(transparent)]
48 Codec(#[from] CodecError),
49 #[error(transparent)]
51 Snapshot(#[from] SnapshotError),
52 #[error("protocol error: {0}")]
54 Protocol(String),
55 #[error("connection closed")]
57 Closed,
58}
59
60#[derive(Clone)]
62pub struct ConnectConfig {
63 pub endpoints: Vec<String>,
69 pub db: String,
71 pub token: Option<String>,
73 pub tls: Option<ClientTlsConfig>,
75 pub reconnect_min: Duration,
77 pub reconnect_max: Duration,
79 pub failover_timeout: Duration,
85 pub heartbeat_timeout: Option<Duration>,
89 storage: Option<Arc<dyn PeerStorage>>,
92}
93
94impl std::fmt::Debug for ConnectConfig {
95 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96 formatter
97 .debug_struct("ConnectConfig")
98 .field("endpoints", &self.endpoints)
99 .field("db", &self.db)
100 .field("token", &self.token.as_ref().map(|_| "[REDACTED]"))
101 .field("tls", &self.tls.is_some())
102 .field("reconnect_min", &self.reconnect_min)
103 .field("reconnect_max", &self.reconnect_max)
104 .field("failover_timeout", &self.failover_timeout)
105 .field("heartbeat_timeout", &self.heartbeat_timeout)
106 .field("storage", &self.storage.is_some())
107 .finish()
108 }
109}
110
111impl ConnectConfig {
112 #[must_use]
114 pub fn new(endpoint: impl Into<String>, db: impl Into<String>) -> Self {
115 Self::with_failover(vec![endpoint.into()], db)
116 }
117
118 #[must_use]
120 pub fn with_failover(endpoints: Vec<String>, db: impl Into<String>) -> Self {
121 Self {
122 endpoints,
123 db: db.into(),
124 token: None,
125 tls: None,
126 reconnect_min: Duration::from_millis(100),
127 reconnect_max: Duration::from_secs(5),
128 failover_timeout: Duration::from_secs(30),
129 heartbeat_timeout: None,
130 storage: None,
131 }
132 }
133
134 #[must_use]
140 pub fn with_storage(mut self, storage: Arc<dyn PeerStorage>) -> Self {
141 self.storage = Some(storage);
142 self
143 }
144
145 #[must_use]
147 pub fn has_storage(&self) -> bool {
148 self.storage.is_some()
149 }
150}
151
152#[derive(Clone, Debug)]
154pub struct PeerReport {
155 pub t: u64,
157 pub tx_instant: i64,
159 pub datoms: Vec<Datom>,
161 pub db_after: Db,
163}
164
165#[derive(Clone, Debug)]
167pub struct TxResult {
168 pub basis_before: u64,
170 pub basis_t: u64,
172 pub tx_instant: i64,
174 pub tempids: BTreeMap<String, EntityId>,
176 pub db_after: Db,
178}
179
180struct PeerState {
181 schema: Schema,
182 idents: Idents,
183 interner: KeywordInterner,
184 db: Db,
185 instants: BTreeMap<u64, i64>,
186}
187
188struct Inner {
189 config: ConnectConfig,
190 state: RwLock<Option<PeerState>>,
191 basis: watch::Sender<u64>,
192 index_basis: watch::Sender<u64>,
193 reports: broadcast::Sender<PeerReport>,
194 client: Mutex<Client>,
195 endpoint_index: std::sync::atomic::AtomicUsize,
198 heartbeat_ms: std::sync::atomic::AtomicU64,
201}
202
203impl Inner {
204 fn heartbeat_deadline(&self) -> Option<Duration> {
207 if let Some(timeout) = self.config.heartbeat_timeout {
208 return Some(timeout);
209 }
210 let advertised = self.heartbeat_ms.load(std::sync::atomic::Ordering::Relaxed);
211 (advertised > 0).then(|| Duration::from_millis(advertised.saturating_mul(3)))
212 }
213}
214
215pub struct Connection {
217 inner: Arc<Inner>,
218 task: tokio::task::JoinHandle<()>,
219}
220
221impl Drop for Connection {
222 fn drop(&mut self) {
223 self.task.abort();
224 }
225}
226
227async fn open_channel(config: &ConnectConfig, endpoint: &str) -> Result<Channel, PeerError> {
228 let mut endpoint = Endpoint::from_shared(endpoint.to_owned())
229 .map_err(|error| PeerError::Protocol(format!("bad endpoint: {error}")))?
230 .connect_timeout(Duration::from_secs(10));
231 if let Some(tls) = &config.tls {
232 endpoint = endpoint.tls_config(tls.clone())?;
233 }
234 Ok(endpoint.connect().await?)
235}
236
237fn make_client(channel: Channel, token: Option<String>) -> Client {
238 TransactorClient::with_interceptor(channel, TokenInterceptor::new(token))
239}
240
241impl Connection {
242 pub async fn connect(config: ConnectConfig) -> Result<Self, PeerError> {
252 if config.endpoints.is_empty() {
253 return Err(PeerError::Protocol("no endpoints configured".into()));
254 }
255 let snapshot = match &config.storage {
256 Some(storage) => load_current_snapshot(storage.as_ref(), &config.db).await?,
257 None => None,
258 };
259 let start_basis = snapshot.as_ref().map_or(0, Db::basis_t);
260 let mut last_error = PeerError::Closed;
263 let mut first = None;
264 for (index, endpoint) in config.endpoints.iter().enumerate() {
265 match open_channel(&config, endpoint).await {
266 Ok(channel) => {
267 let mut client = make_client(channel, config.token.clone());
268 match subscribe_with(&mut client, &config.db, start_basis).await {
269 Ok(stream) => {
270 first = Some((index, client, stream));
271 break;
272 }
273 Err(error) => last_error = error,
274 }
275 }
276 Err(error) => last_error = error,
277 }
278 }
279 let Some((index, client, mut stream)) = first else {
280 return Err(last_error);
281 };
282 let initial_state = snapshot.map(|db| PeerState {
283 schema: db.schema().clone(),
284 idents: db.idents().clone(),
285 interner: db.interner().clone(),
286 db,
287 instants: BTreeMap::new(),
288 });
289 let inner = Arc::new(Inner {
290 state: RwLock::new(initial_state),
291 basis: watch::channel(start_basis).0,
292 index_basis: watch::channel(0).0,
293 reports: broadcast::channel(1024).0,
294 client: Mutex::new(client),
295 endpoint_index: std::sync::atomic::AtomicUsize::new(index),
296 heartbeat_ms: std::sync::atomic::AtomicU64::new(0),
297 config,
298 });
299 let handshake_basis = pump_handshake(&inner, &mut stream).await?;
300 drain_until(&inner, &mut stream, handshake_basis).await?;
301 let task_inner = Arc::clone(&inner);
302 let task = tokio::spawn(async move {
303 run_loop(task_inner, Some(stream)).await;
304 });
305 Ok(Self { inner, task })
306 }
307
308 #[must_use]
310 pub fn db_name(&self) -> &str {
311 &self.inner.config.db
312 }
313
314 #[must_use]
321 pub fn db(&self) -> Db {
322 self.inner
323 .state
324 .read()
325 .unwrap_or_else(std::sync::PoisonError::into_inner)
326 .as_ref()
327 .expect("connection is initialized")
328 .db
329 .clone()
330 }
331
332 #[must_use]
334 pub fn basis_t(&self) -> u64 {
335 *self.inner.basis.subscribe().borrow()
336 }
337
338 #[must_use]
341 pub fn index_basis_t(&self) -> u64 {
342 *self.inner.index_basis.subscribe().borrow()
343 }
344
345 #[must_use]
347 pub fn tx_reports(&self) -> broadcast::Receiver<PeerReport> {
348 self.inner.reports.subscribe()
349 }
350
351 #[must_use]
354 pub fn tx_range(&self, start: u64, end: Option<u64>) -> Vec<TxRecord> {
355 let guard = self
356 .inner
357 .state
358 .read()
359 .unwrap_or_else(std::sync::PoisonError::into_inner);
360 let Some(state) = guard.as_ref() else {
361 return Vec::new();
362 };
363 state
364 .db
365 .tx_range(start, end)
366 .into_iter()
367 .map(|(t, datoms)| TxRecord {
368 t,
369 tx_instant: state.instants.get(&t).copied().unwrap_or_default(),
370 datoms,
371 })
372 .collect()
373 }
374
375 pub async fn transact(&self, forms: Vec<Edn>) -> Result<TxResult, PeerError> {
390 let tx_data = codec::encode_edn(&Edn::Vector(forms));
391 let deadline = tokio::time::Instant::now() + self.inner.config.failover_timeout;
392 let response = loop {
393 match self.transact_raw(tx_data.clone()).await {
394 Ok(response) => break response,
395 Err(error) if retry_is_safe(&error) && tokio::time::Instant::now() < deadline => {
396 tokio::time::sleep(self.inner.config.reconnect_min).await;
397 }
398 Err(error) => return Err(error),
399 }
400 };
401 let tempids = decode_tempids(&response.tempids)?;
402 let db_after = self.sync_to(response.basis_t).await?;
403 Ok(TxResult {
404 basis_before: response.basis_before,
405 basis_t: response.basis_t,
406 tx_instant: response.tx_instant,
407 tempids,
408 db_after,
409 })
410 }
411
412 pub async fn transact_raw(&self, tx_data: Vec<u8>) -> Result<pb::TransactResponse, PeerError> {
418 let request = pb::TransactRequest {
419 db: self.inner.config.db.clone(),
420 protocol_version: corium_protocol::PROTOCOL_VERSION,
421 tx_data,
422 };
423 let mut client = self
424 .inner
425 .client
426 .lock()
427 .unwrap_or_else(std::sync::PoisonError::into_inner)
428 .clone();
429 Ok(client.transact(request).await?.into_inner())
430 }
431
432 pub async fn sync(&self) -> Result<Db, PeerError> {
437 let mut client = self
438 .inner
439 .client
440 .lock()
441 .unwrap_or_else(std::sync::PoisonError::into_inner)
442 .clone();
443 let response = client
444 .sync(pb::SyncRequest {
445 db: self.inner.config.db.clone(),
446 t: 0,
447 })
448 .await?
449 .into_inner();
450 self.sync_to(response.basis_t).await
451 }
452
453 pub async fn sync_to(&self, t: u64) -> Result<Db, PeerError> {
458 let mut basis = self.inner.basis.subscribe();
459 loop {
460 if *basis.borrow() >= t {
461 return Ok(self.db());
462 }
463 basis.changed().await.map_err(|_| PeerError::Closed)?;
464 }
465 }
466
467 pub async fn status(&self) -> Result<pb::StatusResponse, PeerError> {
472 let mut client = self
473 .inner
474 .client
475 .lock()
476 .unwrap_or_else(std::sync::PoisonError::into_inner)
477 .clone();
478 Ok(client
479 .status(pb::StatusRequest {
480 db: self.inner.config.db.clone(),
481 })
482 .await?
483 .into_inner())
484 }
485
486 pub async fn subscribe_raw(
492 &self,
493 from_basis_t: u64,
494 ) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
495 let mut client = self
496 .inner
497 .client
498 .lock()
499 .unwrap_or_else(std::sync::PoisonError::into_inner)
500 .clone();
501 Ok(client
502 .subscribe(pb::SubscribeRequest {
503 db: self.inner.config.db.clone(),
504 protocol_version: corium_protocol::PROTOCOL_VERSION,
505 from_basis_t,
506 })
507 .await?
508 .into_inner())
509 }
510}
511
512fn retry_is_safe(error: &PeerError) -> bool {
515 match error {
516 PeerError::Transport(_) => true,
518 PeerError::Rpc(status) => match status.code() {
519 tonic::Code::FailedPrecondition => {
522 let message = status.message();
523 message.contains("standby") || message.contains("deposed")
524 }
525 tonic::Code::Unavailable => status.message().contains("connect"),
529 _ => false,
530 },
531 _ => false,
532 }
533}
534
535fn decode_tempids(bytes: &[u8]) -> Result<BTreeMap<String, EntityId>, PeerError> {
536 let Edn::Map(pairs) = codec::decode_edn(bytes)? else {
537 return Err(PeerError::Protocol("tempids must be a map".into()));
538 };
539 let mut tempids = BTreeMap::new();
540 for (key, value) in pairs {
541 let (Edn::Str(name), Edn::Long(raw)) = (key, value) else {
542 return Err(PeerError::Protocol("bad tempid entry".into()));
543 };
544 let raw = u64::try_from(raw).map_err(|_| PeerError::Protocol("bad entity id".into()))?;
545 tempids.insert(name, EntityId::from_raw(raw));
546 }
547 Ok(tempids)
548}
549
550async fn subscribe_with(
551 client: &mut Client,
552 db: &str,
553 from_basis_t: u64,
554) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
555 Ok(client
556 .subscribe(pb::SubscribeRequest {
557 db: db.to_owned(),
558 protocol_version: corium_protocol::PROTOCOL_VERSION,
559 from_basis_t,
560 })
561 .await?
562 .into_inner())
563}
564
565async fn pump_handshake(
568 inner: &Arc<Inner>,
569 stream: &mut tonic::Streaming<pb::SubscribeItem>,
570) -> Result<u64, PeerError> {
571 let first = stream
572 .message()
573 .await?
574 .and_then(|item| item.item)
575 .ok_or_else(|| PeerError::Protocol("subscription ended before handshake".into()))?;
576 let pb::subscribe_item::Item::Handshake(handshake) = first else {
577 return Err(PeerError::Protocol(
578 "subscription must begin with a handshake".into(),
579 ));
580 };
581 let local_basis = *inner.basis.subscribe().borrow();
582 if handshake.basis_t < local_basis {
583 return Err(PeerError::Protocol(format!(
584 "published snapshot basis {local_basis} is newer than transactor basis {}",
585 handshake.basis_t
586 )));
587 }
588 let (schema, idents) = codec::decode_schema(&handshake.schema)?;
589 inner.heartbeat_ms.store(
590 handshake.heartbeat_interval_ms,
591 std::sync::atomic::Ordering::Relaxed,
592 );
593 {
594 let mut guard = inner
595 .state
596 .write()
597 .unwrap_or_else(std::sync::PoisonError::into_inner);
598 if let Some(state) = guard.as_mut() {
599 state.schema = schema;
602 state.idents = idents;
603 } else {
604 let interner = KeywordInterner::default();
605 let db = Db::new(schema.clone()).with_naming(idents.clone(), interner.clone());
606 *guard = Some(PeerState {
607 schema,
608 idents,
609 interner,
610 db,
611 instants: BTreeMap::new(),
612 });
613 }
614 }
615 let _ = inner.index_basis.send_replace(handshake.index_basis_t);
616 Ok(handshake.basis_t)
617}
618
619async fn drain_until(
621 inner: &Arc<Inner>,
622 stream: &mut tonic::Streaming<pb::SubscribeItem>,
623 target: u64,
624) -> Result<(), PeerError> {
625 while *inner.basis.subscribe().borrow() < target {
626 let Some(item) = stream.message().await? else {
627 return Err(PeerError::Protocol(
628 "subscription ended during backfill".into(),
629 ));
630 };
631 if let Some(item) = item.item {
632 apply_item(inner, item)?;
633 }
634 }
635 Ok(())
636}
637
638fn apply_item(inner: &Arc<Inner>, item: pb::subscribe_item::Item) -> Result<(), PeerError> {
639 match item {
640 pb::subscribe_item::Item::Report(report) => {
641 let peer_report = {
642 let mut guard = inner
643 .state
644 .write()
645 .unwrap_or_else(std::sync::PoisonError::into_inner);
646 let state = guard
647 .as_mut()
648 .ok_or_else(|| PeerError::Protocol("report before handshake".into()))?;
649 if report.t <= state.db.basis_t() {
650 return Ok(());
651 }
652 let before = state.interner.len();
653 let datoms = codec::decode_datoms(&report.datoms, &mut state.interner)?;
654 if state.interner.len() > before {
655 state.db = state
656 .db
657 .clone()
658 .with_naming(state.idents.clone(), state.interner.clone());
659 }
660 state.db = state.db.with_transaction(report.t, &datoms);
661 state.instants.insert(report.t, report.tx_instant);
662 PeerReport {
663 t: report.t,
664 tx_instant: report.tx_instant,
665 datoms,
666 db_after: state.db.clone(),
667 }
668 };
669 let _ = inner.basis.send_replace(peer_report.t);
670 let _ = inner.reports.send(peer_report);
671 }
672 pb::subscribe_item::Item::IndexBasis(index) => {
673 let _ = inner.index_basis.send_replace(index.index_basis_t);
674 }
675 pb::subscribe_item::Item::Heartbeat(_) | pb::subscribe_item::Item::Handshake(_) => {}
676 }
677 Ok(())
678}
679
680async fn run_loop(inner: Arc<Inner>, initial: Option<tonic::Streaming<pb::SubscribeItem>>) {
687 let mut stream = initial;
688 let mut backoff = inner.config.reconnect_min;
689 let mut candidate = inner
690 .endpoint_index
691 .load(std::sync::atomic::Ordering::Relaxed);
692 loop {
693 if let Some(active) = stream.as_mut() {
694 let next = match inner.heartbeat_deadline() {
695 Some(deadline) => match tokio::time::timeout(deadline, active.message()).await {
696 Ok(next) => next,
697 Err(_elapsed) => Ok(None),
701 },
702 None => active.message().await,
703 };
704 match next {
705 Ok(Some(item)) => {
706 backoff = inner.config.reconnect_min;
707 if let Some(item) = item.item
708 && apply_item(&inner, item).is_err()
709 {
710 stream = None;
711 }
712 continue;
713 }
714 Ok(None) | Err(_) => {
715 stream = None;
716 }
717 }
718 }
719 tokio::time::sleep(backoff).await;
720 backoff = (backoff * 2).min(inner.config.reconnect_max);
721 let endpoints = &inner.config.endpoints;
722 let endpoint = &endpoints[candidate % endpoints.len()];
723 let Ok(channel) = open_channel(&inner.config, endpoint).await else {
724 candidate = (candidate + 1) % endpoints.len();
725 continue;
726 };
727 let mut client = make_client(channel, inner.config.token.clone());
728 let from = *inner.basis.subscribe().borrow();
729 match subscribe_with(&mut client, &inner.config.db, from).await {
730 Ok(mut fresh) => {
731 if pump_handshake(&inner, &mut fresh).await.is_ok() {
732 *inner
734 .client
735 .lock()
736 .unwrap_or_else(std::sync::PoisonError::into_inner) = client;
737 inner.endpoint_index.store(
738 candidate % endpoints.len(),
739 std::sync::atomic::Ordering::Relaxed,
740 );
741 stream = Some(fresh);
742 } else {
743 candidate = (candidate + 1) % endpoints.len();
744 }
745 }
746 Err(_) => {
747 candidate = (candidate + 1) % endpoints.len();
748 }
749 }
750 }
751}
752
753pub struct Admin {
755 client: CatalogClient<InterceptedService<Channel, TokenInterceptor>>,
756}
757
758impl Admin {
759 pub async fn connect(
764 endpoint: &str,
765 token: Option<String>,
766 tls: Option<ClientTlsConfig>,
767 ) -> Result<Self, PeerError> {
768 let config = ConnectConfig {
769 tls,
770 token: token.clone(),
771 ..ConnectConfig::new(endpoint, "")
772 };
773 let channel = open_channel(&config, endpoint).await?;
774 Ok(Self {
775 client: CatalogClient::with_interceptor(channel, TokenInterceptor::new(token)),
776 })
777 }
778
779 pub async fn create_database(&mut self, db: &str, schema: &[Edn]) -> Result<bool, PeerError> {
784 let response = self
785 .client
786 .create_database(pb::CreateDatabaseRequest {
787 db: db.to_owned(),
788 schema: codec::encode_edn(&Edn::Vector(schema.to_vec())),
789 })
790 .await?;
791 Ok(response.into_inner().created)
792 }
793
794 pub async fn fork_database(
802 &mut self,
803 db: &str,
804 target: &str,
805 as_of_t: Option<u64>,
806 ) -> Result<Option<u64>, PeerError> {
807 let response = self
808 .client
809 .fork_database(pb::ForkDatabaseRequest {
810 db: db.to_owned(),
811 target: target.to_owned(),
812 as_of_t: as_of_t.unwrap_or(0),
813 })
814 .await?
815 .into_inner();
816 Ok(response.created.then_some(response.basis_t))
817 }
818
819 pub async fn delete_database(&mut self, db: &str) -> Result<bool, PeerError> {
824 let response = self
825 .client
826 .delete_database(pb::DeleteDatabaseRequest { db: db.to_owned() })
827 .await?;
828 Ok(response.into_inner().deleted)
829 }
830
831 pub async fn list_databases(&mut self) -> Result<Vec<String>, PeerError> {
836 let response = self
837 .client
838 .list_databases(pb::ListDatabasesRequest {})
839 .await?;
840 Ok(response.into_inner().dbs)
841 }
842
843 pub async fn gc_deleted_databases(&mut self) -> Result<u64, PeerError> {
848 self.gc_deleted_databases_with_retention(None).await
849 }
850
851 pub async fn gc_deleted_databases_with_retention(
856 &mut self,
857 retention: Option<Duration>,
858 ) -> Result<u64, PeerError> {
859 let response = self
860 .client
861 .gc_deleted_databases(pb::GcDeletedDatabasesRequest {
862 retention_millis: retention.map(duration_millis),
863 })
864 .await?;
865 Ok(response.into_inner().swept_blobs)
866 }
867
868 pub async fn request_index(&mut self, db: &str) -> Result<u64, PeerError> {
874 let response = self
875 .client
876 .request_index(pb::RequestIndexRequest { db: db.to_owned() })
877 .await?;
878 Ok(response.into_inner().index_basis_t)
879 }
880
881 pub async fn set_index_policy(
888 &mut self,
889 db: &str,
890 update: IndexPolicySettings,
891 ) -> Result<IndexPolicySettings, PeerError> {
892 let response = self
893 .client
894 .set_index_policy(pb::SetIndexPolicyRequest {
895 db: db.to_owned(),
896 interval_ms: update.interval_ms,
897 backoff: update.backoff,
898 tail_threshold: update.tail_threshold,
899 tail_deadline_ms: update.tail_deadline_ms,
900 })
901 .await?
902 .into_inner();
903 Ok(IndexPolicySettings {
904 interval_ms: Some(response.interval_ms),
905 backoff: Some(response.backoff),
906 tail_threshold: Some(response.tail_threshold),
907 tail_deadline_ms: Some(response.tail_deadline_ms),
908 })
909 }
910}
911
912#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
916pub struct IndexPolicySettings {
917 pub interval_ms: Option<u64>,
919 pub backoff: Option<u32>,
922 pub tail_threshold: Option<u64>,
925 pub tail_deadline_ms: Option<u64>,
927}
928
929fn duration_millis(duration: Duration) -> u64 {
930 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
931}
932
933#[cfg(test)]
934mod tests {
935 use super::*;
936
937 #[test]
938 fn gc_retention_wire_value_preserves_zero_and_subseconds() {
939 assert_eq!(duration_millis(Duration::ZERO), 0);
940 assert_eq!(duration_millis(Duration::from_millis(500)), 500);
941 }
942}