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
33type Client = TransactorClient<InterceptedService<Channel, TokenInterceptor>>;
34
35#[derive(Debug, Error)]
37pub enum PeerError {
38 #[error(transparent)]
40 Transport(#[from] tonic::transport::Error),
41 #[error(transparent)]
43 Rpc(#[from] Status),
44 #[error(transparent)]
46 Codec(#[from] CodecError),
47 #[error("protocol error: {0}")]
49 Protocol(String),
50 #[error("connection closed")]
52 Closed,
53}
54
55#[derive(Clone, Debug)]
57pub struct ConnectConfig {
58 pub endpoints: Vec<String>,
64 pub db: String,
66 pub token: Option<String>,
68 pub tls: Option<ClientTlsConfig>,
70 pub reconnect_min: Duration,
72 pub reconnect_max: Duration,
74 pub failover_timeout: Duration,
80 pub heartbeat_timeout: Option<Duration>,
84}
85
86impl ConnectConfig {
87 #[must_use]
89 pub fn new(endpoint: impl Into<String>, db: impl Into<String>) -> Self {
90 Self::with_failover(vec![endpoint.into()], db)
91 }
92
93 #[must_use]
95 pub fn with_failover(endpoints: Vec<String>, db: impl Into<String>) -> Self {
96 Self {
97 endpoints,
98 db: db.into(),
99 token: None,
100 tls: None,
101 reconnect_min: Duration::from_millis(100),
102 reconnect_max: Duration::from_secs(5),
103 failover_timeout: Duration::from_secs(30),
104 heartbeat_timeout: None,
105 }
106 }
107}
108
109#[derive(Clone, Debug)]
111pub struct PeerReport {
112 pub t: u64,
114 pub tx_instant: i64,
116 pub datoms: Vec<Datom>,
118 pub db_after: Db,
120}
121
122#[derive(Clone, Debug)]
124pub struct TxResult {
125 pub basis_before: u64,
127 pub basis_t: u64,
129 pub tx_instant: i64,
131 pub tempids: BTreeMap<String, EntityId>,
133 pub db_after: Db,
135}
136
137struct PeerState {
138 schema: Schema,
139 idents: Idents,
140 interner: KeywordInterner,
141 db: Db,
142 instants: BTreeMap<u64, i64>,
143}
144
145struct Inner {
146 config: ConnectConfig,
147 state: RwLock<Option<PeerState>>,
148 basis: watch::Sender<u64>,
149 index_basis: watch::Sender<u64>,
150 reports: broadcast::Sender<PeerReport>,
151 client: Mutex<Client>,
152 endpoint_index: std::sync::atomic::AtomicUsize,
155 heartbeat_ms: std::sync::atomic::AtomicU64,
158}
159
160impl Inner {
161 fn heartbeat_deadline(&self) -> Option<Duration> {
164 if let Some(timeout) = self.config.heartbeat_timeout {
165 return Some(timeout);
166 }
167 let advertised = self.heartbeat_ms.load(std::sync::atomic::Ordering::Relaxed);
168 (advertised > 0).then(|| Duration::from_millis(advertised.saturating_mul(3)))
169 }
170}
171
172pub struct Connection {
174 inner: Arc<Inner>,
175 task: tokio::task::JoinHandle<()>,
176}
177
178impl Drop for Connection {
179 fn drop(&mut self) {
180 self.task.abort();
181 }
182}
183
184async fn open_channel(config: &ConnectConfig, endpoint: &str) -> Result<Channel, PeerError> {
185 let mut endpoint = Endpoint::from_shared(endpoint.to_owned())
186 .map_err(|error| PeerError::Protocol(format!("bad endpoint: {error}")))?
187 .connect_timeout(Duration::from_secs(10));
188 if let Some(tls) = &config.tls {
189 endpoint = endpoint.tls_config(tls.clone())?;
190 }
191 Ok(endpoint.connect().await?)
192}
193
194fn make_client(channel: Channel, token: Option<String>) -> Client {
195 TransactorClient::with_interceptor(channel, TokenInterceptor::new(token))
196}
197
198impl Connection {
199 pub async fn connect(config: ConnectConfig) -> Result<Self, PeerError> {
207 if config.endpoints.is_empty() {
208 return Err(PeerError::Protocol("no endpoints configured".into()));
209 }
210 let mut last_error = PeerError::Closed;
213 let mut first = None;
214 for (index, endpoint) in config.endpoints.iter().enumerate() {
215 match open_channel(&config, endpoint).await {
216 Ok(channel) => {
217 let mut client = make_client(channel, config.token.clone());
218 match subscribe_with(&mut client, &config.db, 0).await {
219 Ok(stream) => {
220 first = Some((index, client, stream));
221 break;
222 }
223 Err(error) => last_error = error,
224 }
225 }
226 Err(error) => last_error = error,
227 }
228 }
229 let Some((index, client, mut stream)) = first else {
230 return Err(last_error);
231 };
232 let inner = Arc::new(Inner {
233 state: RwLock::new(None),
234 basis: watch::channel(0).0,
235 index_basis: watch::channel(0).0,
236 reports: broadcast::channel(1024).0,
237 client: Mutex::new(client),
238 endpoint_index: std::sync::atomic::AtomicUsize::new(index),
239 heartbeat_ms: std::sync::atomic::AtomicU64::new(0),
240 config,
241 });
242 let handshake_basis = pump_handshake(&inner, &mut stream).await?;
243 drain_until(&inner, &mut stream, handshake_basis).await?;
244 let task_inner = Arc::clone(&inner);
245 let task = tokio::spawn(async move {
246 run_loop(task_inner, Some(stream)).await;
247 });
248 Ok(Self { inner, task })
249 }
250
251 #[must_use]
253 pub fn db_name(&self) -> &str {
254 &self.inner.config.db
255 }
256
257 #[must_use]
264 pub fn db(&self) -> Db {
265 self.inner
266 .state
267 .read()
268 .unwrap_or_else(std::sync::PoisonError::into_inner)
269 .as_ref()
270 .expect("connection is initialized")
271 .db
272 .clone()
273 }
274
275 #[must_use]
277 pub fn basis_t(&self) -> u64 {
278 *self.inner.basis.subscribe().borrow()
279 }
280
281 #[must_use]
284 pub fn index_basis_t(&self) -> u64 {
285 *self.inner.index_basis.subscribe().borrow()
286 }
287
288 #[must_use]
290 pub fn tx_reports(&self) -> broadcast::Receiver<PeerReport> {
291 self.inner.reports.subscribe()
292 }
293
294 #[must_use]
297 pub fn tx_range(&self, start: u64, end: Option<u64>) -> Vec<TxRecord> {
298 let guard = self
299 .inner
300 .state
301 .read()
302 .unwrap_or_else(std::sync::PoisonError::into_inner);
303 let Some(state) = guard.as_ref() else {
304 return Vec::new();
305 };
306 state
307 .db
308 .tx_range(start, end)
309 .into_iter()
310 .map(|(t, datoms)| TxRecord {
311 t,
312 tx_instant: state.instants.get(&t).copied().unwrap_or_default(),
313 datoms,
314 })
315 .collect()
316 }
317
318 pub async fn transact(&self, forms: Vec<Edn>) -> Result<TxResult, PeerError> {
333 let tx_data = codec::encode_edn(&Edn::Vector(forms));
334 let deadline = tokio::time::Instant::now() + self.inner.config.failover_timeout;
335 let response = loop {
336 match self.transact_raw(tx_data.clone()).await {
337 Ok(response) => break response,
338 Err(error) if retry_is_safe(&error) && tokio::time::Instant::now() < deadline => {
339 tokio::time::sleep(self.inner.config.reconnect_min).await;
340 }
341 Err(error) => return Err(error),
342 }
343 };
344 let tempids = decode_tempids(&response.tempids)?;
345 let db_after = self.sync_to(response.basis_t).await?;
346 Ok(TxResult {
347 basis_before: response.basis_before,
348 basis_t: response.basis_t,
349 tx_instant: response.tx_instant,
350 tempids,
351 db_after,
352 })
353 }
354
355 pub async fn transact_raw(&self, tx_data: Vec<u8>) -> Result<pb::TransactResponse, PeerError> {
361 let request = pb::TransactRequest {
362 db: self.inner.config.db.clone(),
363 protocol_version: corium_protocol::PROTOCOL_VERSION,
364 tx_data,
365 };
366 let mut client = self
367 .inner
368 .client
369 .lock()
370 .unwrap_or_else(std::sync::PoisonError::into_inner)
371 .clone();
372 Ok(client.transact(request).await?.into_inner())
373 }
374
375 pub async fn sync(&self) -> Result<Db, PeerError> {
380 let mut client = self
381 .inner
382 .client
383 .lock()
384 .unwrap_or_else(std::sync::PoisonError::into_inner)
385 .clone();
386 let response = client
387 .sync(pb::SyncRequest {
388 db: self.inner.config.db.clone(),
389 t: 0,
390 })
391 .await?
392 .into_inner();
393 self.sync_to(response.basis_t).await
394 }
395
396 pub async fn sync_to(&self, t: u64) -> Result<Db, PeerError> {
401 let mut basis = self.inner.basis.subscribe();
402 loop {
403 if *basis.borrow() >= t {
404 return Ok(self.db());
405 }
406 basis.changed().await.map_err(|_| PeerError::Closed)?;
407 }
408 }
409
410 pub async fn status(&self) -> Result<pb::StatusResponse, PeerError> {
415 let mut client = self
416 .inner
417 .client
418 .lock()
419 .unwrap_or_else(std::sync::PoisonError::into_inner)
420 .clone();
421 Ok(client
422 .status(pb::StatusRequest {
423 db: self.inner.config.db.clone(),
424 })
425 .await?
426 .into_inner())
427 }
428
429 pub async fn subscribe_raw(
435 &self,
436 from_basis_t: u64,
437 ) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
438 let mut client = self
439 .inner
440 .client
441 .lock()
442 .unwrap_or_else(std::sync::PoisonError::into_inner)
443 .clone();
444 Ok(client
445 .subscribe(pb::SubscribeRequest {
446 db: self.inner.config.db.clone(),
447 protocol_version: corium_protocol::PROTOCOL_VERSION,
448 from_basis_t,
449 })
450 .await?
451 .into_inner())
452 }
453}
454
455fn retry_is_safe(error: &PeerError) -> bool {
458 match error {
459 PeerError::Transport(_) => true,
461 PeerError::Rpc(status) => match status.code() {
462 tonic::Code::FailedPrecondition => {
465 let message = status.message();
466 message.contains("standby") || message.contains("deposed")
467 }
468 tonic::Code::Unavailable => status.message().contains("connect"),
472 _ => false,
473 },
474 _ => false,
475 }
476}
477
478fn decode_tempids(bytes: &[u8]) -> Result<BTreeMap<String, EntityId>, PeerError> {
479 let Edn::Map(pairs) = codec::decode_edn(bytes)? else {
480 return Err(PeerError::Protocol("tempids must be a map".into()));
481 };
482 let mut tempids = BTreeMap::new();
483 for (key, value) in pairs {
484 let (Edn::Str(name), Edn::Long(raw)) = (key, value) else {
485 return Err(PeerError::Protocol("bad tempid entry".into()));
486 };
487 let raw = u64::try_from(raw).map_err(|_| PeerError::Protocol("bad entity id".into()))?;
488 tempids.insert(name, EntityId::from_raw(raw));
489 }
490 Ok(tempids)
491}
492
493async fn subscribe_with(
494 client: &mut Client,
495 db: &str,
496 from_basis_t: u64,
497) -> Result<tonic::Streaming<pb::SubscribeItem>, PeerError> {
498 Ok(client
499 .subscribe(pb::SubscribeRequest {
500 db: db.to_owned(),
501 protocol_version: corium_protocol::PROTOCOL_VERSION,
502 from_basis_t,
503 })
504 .await?
505 .into_inner())
506}
507
508async fn pump_handshake(
511 inner: &Arc<Inner>,
512 stream: &mut tonic::Streaming<pb::SubscribeItem>,
513) -> Result<u64, PeerError> {
514 let first = stream
515 .message()
516 .await?
517 .and_then(|item| item.item)
518 .ok_or_else(|| PeerError::Protocol("subscription ended before handshake".into()))?;
519 let pb::subscribe_item::Item::Handshake(handshake) = first else {
520 return Err(PeerError::Protocol(
521 "subscription must begin with a handshake".into(),
522 ));
523 };
524 let (schema, idents) = codec::decode_schema(&handshake.schema)?;
525 inner.heartbeat_ms.store(
526 handshake.heartbeat_interval_ms,
527 std::sync::atomic::Ordering::Relaxed,
528 );
529 {
530 let mut guard = inner
531 .state
532 .write()
533 .unwrap_or_else(std::sync::PoisonError::into_inner);
534 if let Some(state) = guard.as_mut() {
535 state.schema = schema;
538 state.idents = idents;
539 } else {
540 let interner = KeywordInterner::default();
541 let db = Db::new(schema.clone()).with_naming(idents.clone(), interner.clone());
542 *guard = Some(PeerState {
543 schema,
544 idents,
545 interner,
546 db,
547 instants: BTreeMap::new(),
548 });
549 }
550 }
551 let _ = inner.index_basis.send_replace(handshake.index_basis_t);
552 Ok(handshake.basis_t)
553}
554
555async fn drain_until(
557 inner: &Arc<Inner>,
558 stream: &mut tonic::Streaming<pb::SubscribeItem>,
559 target: u64,
560) -> Result<(), PeerError> {
561 while *inner.basis.subscribe().borrow() < target {
562 let Some(item) = stream.message().await? else {
563 return Err(PeerError::Protocol(
564 "subscription ended during backfill".into(),
565 ));
566 };
567 if let Some(item) = item.item {
568 apply_item(inner, item)?;
569 }
570 }
571 Ok(())
572}
573
574fn apply_item(inner: &Arc<Inner>, item: pb::subscribe_item::Item) -> Result<(), PeerError> {
575 match item {
576 pb::subscribe_item::Item::Report(report) => {
577 let peer_report = {
578 let mut guard = inner
579 .state
580 .write()
581 .unwrap_or_else(std::sync::PoisonError::into_inner);
582 let state = guard
583 .as_mut()
584 .ok_or_else(|| PeerError::Protocol("report before handshake".into()))?;
585 if report.t <= state.db.basis_t() {
586 return Ok(());
587 }
588 let before = state.interner.len();
589 let datoms = codec::decode_datoms(&report.datoms, &mut state.interner)?;
590 if state.interner.len() > before {
591 state.db = state
592 .db
593 .clone()
594 .with_naming(state.idents.clone(), state.interner.clone());
595 }
596 state.db = state.db.with_transaction(report.t, &datoms);
597 state.instants.insert(report.t, report.tx_instant);
598 PeerReport {
599 t: report.t,
600 tx_instant: report.tx_instant,
601 datoms,
602 db_after: state.db.clone(),
603 }
604 };
605 let _ = inner.basis.send_replace(peer_report.t);
606 let _ = inner.reports.send(peer_report);
607 }
608 pb::subscribe_item::Item::IndexBasis(index) => {
609 let _ = inner.index_basis.send_replace(index.index_basis_t);
610 }
611 pb::subscribe_item::Item::Heartbeat(_) | pb::subscribe_item::Item::Handshake(_) => {}
612 }
613 Ok(())
614}
615
616async fn run_loop(inner: Arc<Inner>, initial: Option<tonic::Streaming<pb::SubscribeItem>>) {
623 let mut stream = initial;
624 let mut backoff = inner.config.reconnect_min;
625 let mut candidate = inner
626 .endpoint_index
627 .load(std::sync::atomic::Ordering::Relaxed);
628 loop {
629 if let Some(active) = stream.as_mut() {
630 let next = match inner.heartbeat_deadline() {
631 Some(deadline) => match tokio::time::timeout(deadline, active.message()).await {
632 Ok(next) => next,
633 Err(_elapsed) => Ok(None),
637 },
638 None => active.message().await,
639 };
640 match next {
641 Ok(Some(item)) => {
642 backoff = inner.config.reconnect_min;
643 if let Some(item) = item.item {
644 if apply_item(&inner, item).is_err() {
645 stream = None;
646 }
647 }
648 continue;
649 }
650 Ok(None) | Err(_) => {
651 stream = None;
652 }
653 }
654 }
655 tokio::time::sleep(backoff).await;
656 backoff = (backoff * 2).min(inner.config.reconnect_max);
657 let endpoints = &inner.config.endpoints;
658 let endpoint = &endpoints[candidate % endpoints.len()];
659 let Ok(channel) = open_channel(&inner.config, endpoint).await else {
660 candidate = (candidate + 1) % endpoints.len();
661 continue;
662 };
663 let mut client = make_client(channel, inner.config.token.clone());
664 let from = *inner.basis.subscribe().borrow();
665 match subscribe_with(&mut client, &inner.config.db, from).await {
666 Ok(mut fresh) => {
667 if pump_handshake(&inner, &mut fresh).await.is_ok() {
668 *inner
670 .client
671 .lock()
672 .unwrap_or_else(std::sync::PoisonError::into_inner) = client;
673 inner.endpoint_index.store(
674 candidate % endpoints.len(),
675 std::sync::atomic::Ordering::Relaxed,
676 );
677 stream = Some(fresh);
678 } else {
679 candidate = (candidate + 1) % endpoints.len();
680 }
681 }
682 Err(_) => {
683 candidate = (candidate + 1) % endpoints.len();
684 }
685 }
686 }
687}
688
689pub struct Admin {
691 client: CatalogClient<InterceptedService<Channel, TokenInterceptor>>,
692}
693
694impl Admin {
695 pub async fn connect(
700 endpoint: &str,
701 token: Option<String>,
702 tls: Option<ClientTlsConfig>,
703 ) -> Result<Self, PeerError> {
704 let config = ConnectConfig {
705 tls,
706 token: token.clone(),
707 ..ConnectConfig::new(endpoint, "")
708 };
709 let channel = open_channel(&config, endpoint).await?;
710 Ok(Self {
711 client: CatalogClient::with_interceptor(channel, TokenInterceptor::new(token)),
712 })
713 }
714
715 pub async fn create_database(&mut self, db: &str, schema: &[Edn]) -> Result<bool, PeerError> {
720 let response = self
721 .client
722 .create_database(pb::CreateDatabaseRequest {
723 db: db.to_owned(),
724 schema: codec::encode_edn(&Edn::Vector(schema.to_vec())),
725 })
726 .await?;
727 Ok(response.into_inner().created)
728 }
729
730 pub async fn delete_database(&mut self, db: &str) -> Result<bool, PeerError> {
735 let response = self
736 .client
737 .delete_database(pb::DeleteDatabaseRequest { db: db.to_owned() })
738 .await?;
739 Ok(response.into_inner().deleted)
740 }
741
742 pub async fn list_databases(&mut self) -> Result<Vec<String>, PeerError> {
747 let response = self
748 .client
749 .list_databases(pb::ListDatabasesRequest {})
750 .await?;
751 Ok(response.into_inner().dbs)
752 }
753
754 pub async fn gc_deleted_databases(&mut self) -> Result<u64, PeerError> {
759 self.gc_deleted_databases_with_retention(None).await
760 }
761
762 pub async fn gc_deleted_databases_with_retention(
767 &mut self,
768 retention: Option<Duration>,
769 ) -> Result<u64, PeerError> {
770 let response = self
771 .client
772 .gc_deleted_databases(pb::GcDeletedDatabasesRequest {
773 retention_millis: retention.map(duration_millis),
774 })
775 .await?;
776 Ok(response.into_inner().swept_blobs)
777 }
778}
779
780fn duration_millis(duration: Duration) -> u64 {
781 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
782}
783
784#[cfg(test)]
785mod tests {
786 use super::*;
787
788 #[test]
789 fn gc_retention_wire_value_preserves_zero_and_subseconds() {
790 assert_eq!(duration_millis(Duration::ZERO), 0);
791 assert_eq!(duration_millis(Duration::from_millis(500)), 500);
792 }
793}