1use crate::{
9 agent::{TO_CLEAR_COUNT, handlers},
10 api::public::{
11 api_v1_db_schema, api_v1_queries, api_v1_table_stats, api_v1_transactions,
12 pubsub::{api_v1_sub_by_id, api_v1_subs},
13 update::SharedUpdateBroadcastCache,
14 },
15 transport::Transport,
16};
17
18use antithesis_sdk::assert_sometimes;
19use klukai_types::{
20 actor::{Actor, ActorId},
21 agent::{Agent, Bookie, ChangeError, CurrentVersion, KnownDbVersion, PartialVersion},
22 api::TableName,
23 base::{CrsqlDbVersion, CrsqlDbVersionRange, CrsqlSeq},
24 broadcast::{ChangeSource, ChangeV1, Changeset, ChangesetParts, FocaCmd, FocaInput},
25 channel::CorroReceiver,
26 config::AuthzConfig,
27 pubsub::SubsManager,
28 tripwire::Outcome,
29 updates::{match_changes, match_changes_from_db_version},
30};
31
32use super::BcastCache;
33use crate::api::public::update::api_v1_updates;
34use antithesis_sdk::{assert_always, assert_unreachable};
35use axum::{
36 BoxError, Extension, Router,
37 error_handling::HandleErrorLayer,
38 extract::{DefaultBodyLimit, Request},
39 routing::{get, post},
40};
41use axum_extra::TypedHeader;
42use foca::Member;
43use headers::{Authorization, authorization::Bearer};
44use hyper::StatusCode;
45use klukai_types::{
47 broadcast::Timestamp,
48 spawn::spawn_counted,
49 sqlite_pool::{Committable, InterruptibleTransaction},
50 tripwire::{PreemptibleFutureExt, Tripwire},
51};
52use metrics::{counter, histogram};
53use rangemap::{RangeInclusiveMap, RangeInclusiveSet};
54use rusqlite::{Connection, named_params, params};
55use serde_json::json;
56use std::{
57 cmp,
58 collections::{BTreeMap, HashSet},
59 convert::Infallible,
60 net::SocketAddr,
61 ops::{Deref, RangeInclusive},
62 path::Path,
63 sync::{Arc, atomic::AtomicI64},
64 time::{Duration, Instant},
65};
66use tokio::{
67 net::TcpListener,
68 task::{JoinHandle, block_in_place},
69};
70use tower::{limit::ConcurrencyLimitLayer, load_shed::LoadShedLayer};
71use tower_http::trace::TraceLayer;
72use tracing::{debug, error, info, trace, warn};
73
74pub async fn initialise_foca(agent: &Agent) {
75 let states = load_member_states(agent).await;
76 if !states.is_empty() {
77 let mut foca_states = BTreeMap::<SocketAddr, Member<Actor>>::new();
78
79 {
80 let mut members = agent.members().write();
82 for (address, foca_state) in states {
83 members.by_addr.insert(address, foca_state.id().id());
84 if matches!(foca_state.state(), foca::State::Suspect) {
85 continue;
86 }
87
88 match foca_states.get(&address) {
92 Some(state) if state.id().ts() < foca_state.id().ts() => {
93 foca_states.insert(address, foca_state);
94 }
95 None => {
96 foca_states.insert(address, foca_state);
97 }
98 _ => {}
99 }
100 }
101 }
102
103 if let Err(e) = agent
104 .tx_foca()
105 .send(FocaInput::ApplyMany(foca_states.into_values().collect()))
106 .await
107 {
108 error!(
109 "Failed to queue initial foca state: {e:?}, cluster membership states will be broken!"
110 );
111 }
112
113 let agent = agent.clone();
114 tokio::task::spawn(async move {
115 let scatter = rand::random::<u64>() % 15;
118 tokio::time::sleep(Duration::from_secs(25 + scatter)).await;
119
120 async fn apply_rejoin(agent: &Agent) -> eyre::Result<()> {
121 let (cb_tx, cb_rx) = tokio::sync::oneshot::channel();
122 agent
123 .tx_foca()
124 .send(FocaInput::Cmd(FocaCmd::Rejoin(cb_tx)))
125 .await?;
126 cb_rx
129 .await?
130 .map_err(|foca_err| eyre::eyre!("foca error: {foca_err}"))
131 }
132
133 if let Err(e) = apply_rejoin(&agent).await {
134 error!("failed to execute cluster rejoin: {e:?}");
135 }
136 });
137 } else {
138 warn!("No existing cluster member state to load! This seems sus");
139 }
140}
141
142pub async fn load_member_states(agent: &Agent) -> Vec<(SocketAddr, Member<Actor>)> {
144 match agent.pool().read().await {
145 Ok(conn) => block_in_place(|| {
146 match conn.prepare("SELECT address,foca_state FROM __corro_members") {
147 Ok(mut prepped) => {
148 match prepped
149 .query_map([], |row| Ok((
150 row.get::<_, String>(0)?.parse().map_err(|e| rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e)))?,
151 row.get::<_, String>(1)?
152 ))
153 )
154 .and_then(|rows| rows.collect::<rusqlite::Result<Vec<(SocketAddr, String)>>>())
155 {
156 Ok(members) => {
157 members.into_iter().filter_map(|(address, state)| match serde_json::from_str::<foca::Member<Actor>>(state.as_str()) {
158 Ok(fs) => Some((address, fs)),
159 Err(e) => {
160 error!("could not deserialize foca member state: {e} (json: {state})");
161 None
162 }
163 }).collect::<Vec<(SocketAddr, Member<Actor>)>>()
164 }
165 Err(e) => {
166 error!("could not query for foca member states: {e}");
167 vec![]
168 },
169 }
170 }
171 Err(e) => {
172 error!("could not prepare query for foca member states: {e}");
173 vec![]
174 }
175 }
176 }),
177 Err(e) => {
178 error!("could not acquire conn for foca member states: {e}");
179 vec![]
180 }
181 }
182}
183
184pub async fn setup_http_api_handler(
185 agent: &Agent,
186 tripwire: &Tripwire,
187 subs_bcast_cache: BcastCache,
188 updates_bcast_cache: SharedUpdateBroadcastCache,
189 subs_manager: &SubsManager,
190 api_listeners: Vec<TcpListener>,
191) -> eyre::Result<Vec<JoinHandle<()>>> {
192 let api = Router::<()>::new()
193 .route(
195 "/v1/transactions",
196 post(api_v1_transactions).route_layer(
197 tower::ServiceBuilder::new()
198 .layer(HandleErrorLayer::new(|_error: BoxError| async {
199 Ok::<_, Infallible>((
200 StatusCode::SERVICE_UNAVAILABLE,
201 "max concurrency limit reached".to_string(),
202 ))
203 }))
204 .layer(LoadShedLayer::new())
205 .layer(ConcurrencyLimitLayer::new(128)),
206 ),
207 )
208 .route(
210 "/v1/queries",
211 post(api_v1_queries).route_layer(
212 tower::ServiceBuilder::new()
213 .layer(HandleErrorLayer::new(|_error: BoxError| async {
214 Ok::<_, Infallible>((
215 StatusCode::SERVICE_UNAVAILABLE,
216 "max concurrency limit reached".to_string(),
217 ))
218 }))
219 .layer(LoadShedLayer::new())
220 .layer(ConcurrencyLimitLayer::new(128)),
221 ),
222 )
223 .route(
224 "/v1/subscriptions",
225 post(api_v1_subs).route_layer(
226 tower::ServiceBuilder::new()
227 .layer(HandleErrorLayer::new(|_error: BoxError| async {
228 Ok::<_, Infallible>((
229 StatusCode::SERVICE_UNAVAILABLE,
230 "max concurrency limit reached".to_string(),
231 ))
232 }))
233 .layer(LoadShedLayer::new())
234 .layer(ConcurrencyLimitLayer::new(128)),
235 ),
236 )
237 .route(
238 "/v1/updates/{table}",
239 post(api_v1_updates).route_layer(
240 tower::ServiceBuilder::new()
241 .layer(HandleErrorLayer::new(|_error: BoxError| async {
242 Ok::<_, Infallible>((
243 StatusCode::SERVICE_UNAVAILABLE,
244 "max concurrency limit reached".to_string(),
245 ))
246 }))
247 .layer(LoadShedLayer::new())
248 .layer(ConcurrencyLimitLayer::new(128)),
249 ),
250 )
251 .route(
252 "/v1/subscriptions/{id}",
253 get(api_v1_sub_by_id).route_layer(
254 tower::ServiceBuilder::new()
255 .layer(HandleErrorLayer::new(|_error: BoxError| async {
256 Ok::<_, Infallible>((
257 StatusCode::SERVICE_UNAVAILABLE,
258 "max concurrency limit reached".to_string(),
259 ))
260 }))
261 .layer(LoadShedLayer::new())
262 .layer(ConcurrencyLimitLayer::new(128)),
263 ),
264 )
265 .route(
266 "/v1/migrations",
267 post(api_v1_db_schema).route_layer(
268 tower::ServiceBuilder::new()
269 .layer(HandleErrorLayer::new(|_error: BoxError| async {
270 Ok::<_, Infallible>((
271 StatusCode::SERVICE_UNAVAILABLE,
272 "max concurrency limit reached".to_string(),
273 ))
274 }))
275 .layer(LoadShedLayer::new())
276 .layer(ConcurrencyLimitLayer::new(4)),
277 ),
278 )
279 .route(
280 "/v1/table_stats",
281 post(api_v1_table_stats).route_layer(
282 tower::ServiceBuilder::new()
283 .layer(HandleErrorLayer::new(|_error: BoxError| async {
284 Ok::<_, Infallible>((
285 StatusCode::SERVICE_UNAVAILABLE,
286 "max concurrency limit reached".to_string(),
287 ))
288 }))
289 .layer(LoadShedLayer::new())
290 .layer(ConcurrencyLimitLayer::new(4)),
291 ),
292 )
293 .layer(axum::middleware::from_fn(require_authz))
294 .layer(
295 tower::ServiceBuilder::new()
296 .layer(Extension(Arc::new(AtomicI64::new(0))))
297 .layer(Extension(agent.clone()))
298 .layer(Extension(subs_bcast_cache))
299 .layer(Extension(updates_bcast_cache))
300 .layer(Extension(subs_manager.clone()))
301 .layer(Extension(tripwire.clone())),
302 )
303 .layer(DefaultBodyLimit::disable())
304 .layer(TraceLayer::new_for_http());
305
306 let mut handles: Vec<JoinHandle<()>> = vec![];
307
308 for api_listener in api_listeners {
309 let api_addr = api_listener.local_addr()?;
310 info!("Starting API listener on tcp/{api_addr}");
311
312 let app = api
313 .clone()
314 .into_make_service_with_connect_info::<SocketAddr>();
315 let tripwire_clone = tripwire.clone();
316
317 let fut = axum::serve(api_listener, app).with_graceful_shutdown(async move {
318 tripwire_clone.await;
319 info!("corrosion api http tripped {api_addr}");
320 });
321
322 handles.push(spawn_counted(async move {
323 if let Err(e) = fut.await {
324 error!("API server error: {e}");
325 }
326 info!("corrosion api is done");
327 }));
328 }
329
330 Ok(handles)
331}
332
333async fn require_authz(
334 Extension(agent): Extension<Agent>,
335 maybe_authz_header: Option<TypedHeader<Authorization<Bearer>>>,
336 request: Request,
337 next: axum::middleware::Next,
338) -> Result<axum::response::Response, axum::http::StatusCode> {
339 let passed = if let Some(ref authz) = agent.config().api.authorization {
340 match authz {
341 AuthzConfig::BearerToken(token) => maybe_authz_header
342 .map(|h| h.token() == token)
343 .unwrap_or(false),
344 }
345 } else {
346 true
347 };
348
349 if !passed {
350 return Err(axum::http::StatusCode::UNAUTHORIZED);
351 }
352
353 Ok(next.run(request).await)
354}
355
356pub async fn sync_loop(agent: Agent, bookie: Bookie, transport: Transport, mut tripwire: Tripwire) {
363 let min_sync_backoff = Duration::from_secs(agent.config().perf.min_sync_backoff as u64);
364 let max_sync_backoff = Duration::from_secs(agent.config().perf.max_sync_backoff as u64);
365 let mut sync_backoff = klukai_types::backoff::Backoff::new(0)
366 .timeout_range(min_sync_backoff, max_sync_backoff)
367 .iter();
368 let next_sync_at = tokio::time::sleep(sync_backoff.next().unwrap());
369 tokio::pin!(next_sync_at);
370
371 loop {
372 tokio::select! {
373 biased;
374
375 _ = &mut next_sync_at => {},
376 _ = &mut tripwire => {
377 break;
378 }
379 };
380
381 match tokio::time::timeout(
383 Duration::from_secs(300),
384 handlers::handle_sync(&agent, &bookie, &transport),
385 )
386 .preemptible(&mut tripwire)
387 .await
388 {
389 klukai_types::tripwire::Outcome::Preempted(_) => {
390 warn!("aborted sync by tripwire");
391 break;
392 }
393 klukai_types::tripwire::Outcome::Completed(res) => match res {
394 Ok(Err(e)) => {
395 error!("could not sync: {e}");
396 }
398 Err(_e) => {
399 warn!("timed out waiting for sync to complete!");
400 }
401 Ok(Ok(_)) => {}
402 },
403 }
404 next_sync_at
405 .as_mut()
406 .reset(tokio::time::Instant::now() + sync_backoff.next().unwrap());
407 }
408}
409
410pub async fn apply_fully_buffered_changes_loop(
411 agent: Agent,
412 bookie: Bookie,
413 mut rx_apply: CorroReceiver<(ActorId, CrsqlDbVersion)>,
414 mut tripwire: Tripwire,
415) {
416 info!("Starting apply_fully_buffered_changes loop");
417
418 let tx_timeout: Duration = Duration::from_secs(agent.config().perf.sql_tx_timeout as u64);
419 while let Outcome::Completed(Some((actor_id, version))) =
420 rx_apply.recv().preemptible(&mut tripwire).await
421 {
422 debug!(%actor_id, %version, "picked up background apply of buffered changes");
423 match process_fully_buffered_changes(&agent, &bookie, actor_id, version, tx_timeout).await {
424 Ok(false) => {
425 warn!(%actor_id, %version, "did not apply buffered changes");
426 }
427 Ok(true) => {
428 debug!(%actor_id, %version, "succesfully applied buffered changes");
429 }
430 Err(e) => {
431 error!(%actor_id, %version, "could not apply fully buffered changes: {e}");
432 }
433 }
434 }
435
436 info!("fully_buffered_changes_loop ended");
437}
438
439pub async fn clear_buffered_meta_loop(
441 agent: Agent,
442 mut rx_partials: CorroReceiver<(ActorId, CrsqlDbVersionRange)>,
443) {
444 let tx_timeout: Duration = Duration::from_secs(agent.config().perf.sql_tx_timeout as u64);
445
446 while let Some((actor_id, versions)) = rx_partials.recv().await {
447 let pool = agent.pool().clone();
448 let self_actor_id = agent.actor_id();
449 tokio::spawn(async move {
450 loop {
451 let res = {
452 let mut conn = pool.write_low().await?;
453
454 block_in_place(|| {
455 let tx = InterruptibleTransaction::new(
456 conn.immediate_transaction()?,
457 Some(tx_timeout),
458 "clear_buffered_meta",
459 );
460
461 let seq_count = tx
465 .prepare_cached("DELETE FROM __corro_seq_bookkeeping WHERE (site_id, db_version, start_seq) IN (SELECT site_id, db_version, start_seq FROM __corro_seq_bookkeeping WHERE site_id = ? AND db_version >= ? AND db_version <= ? LIMIT ?)")?
466 .execute(params![actor_id, versions.start(), versions.end(), TO_CLEAR_COUNT])?;
467
468 let buf_count = tx
470 .prepare_cached("DELETE FROM __corro_buffered_changes WHERE (site_id, db_version, seq) IN (SELECT site_id, db_version, seq FROM __corro_buffered_changes WHERE site_id = ? AND db_version >= ? AND db_version <= ? LIMIT ?)")?
471 .execute(params![actor_id, versions.start(), versions.end(), TO_CLEAR_COUNT])?;
472
473 tx.commit()?;
474
475 Ok::<_, rusqlite::Error>((buf_count, seq_count))
476 })
477 };
478
479 match res {
480 Ok((buf_count, seq_count)) => {
481 if buf_count + seq_count > 0 {
482 assert_sometimes!(true, "Corrosion clears buffered meta");
483 info!(%actor_id, %self_actor_id, "cleared {} buffered meta rows for versions {versions:?}", buf_count + seq_count);
484 }
485 if buf_count < TO_CLEAR_COUNT && seq_count < TO_CLEAR_COUNT {
486 break;
487 }
488 }
489 Err(e) => {
490 error!(%actor_id, "could not clear buffered meta for versions {versions:?}: {e}");
491 }
492 }
493
494 tokio::time::sleep(Duration::from_secs(2)).await;
495 }
496
497 Ok::<_, eyre::Report>(())
498 });
499 }
500}
501
502#[tracing::instrument(skip_all, err)]
503pub fn process_single_version<T: Deref<Target = rusqlite::Connection> + Committable>(
504 agent: &Agent,
505 tx: &mut InterruptibleTransaction<T>,
506 change: ChangeV1,
507) -> rusqlite::Result<(KnownDbVersion, Changeset)> {
508 let ChangeV1 {
509 actor_id,
510 changeset,
511 } = change;
512
513 let versions = changeset.versions();
514
515 let sp = tx.savepoint()?;
516 let mut changes_per_table = BTreeMap::new();
517 let (known, changeset) = if changeset.is_complete() {
518 let (known, changeset, table) = process_complete_version(
519 agent.clone(),
520 &sp,
521 actor_id,
522 versions,
523 changeset
524 .into_parts()
525 .expect("no changeset parts, this shouldn't be happening!"),
526 )?;
527
528 if check_buffered_meta_to_clear(&sp, actor_id, changeset.versions())?
529 && let Err(e) = agent
530 .tx_clear_buf()
531 .try_send((actor_id, changeset.versions()))
532 {
533 error!("could not schedule buffered meta clear: {e}");
534 }
535 changes_per_table = table;
536
537 (known, changeset)
538 } else {
539 let parts = changeset.into_parts().unwrap();
540 let known = process_incomplete_version(&sp, actor_id, &parts)?;
541
542 (known, parts.into())
543 };
544
545 sp.commit()?;
546
547 for (table_name, count) in changes_per_table {
548 counter!("corro.changes.committed", "table" => table_name.to_string(), "source" => "remote").increment(count);
549 }
550
551 Ok((known, changeset))
552}
553
554#[tracing::instrument(skip(agent, bookie), err)]
555pub async fn process_fully_buffered_changes(
556 agent: &Agent,
557 bookie: &Bookie,
558 actor_id: ActorId,
559 version: CrsqlDbVersion,
560 tx_timeout: Duration,
561) -> Result<bool, ChangeError> {
562 let rows_impacted = {
563 let mut conn = agent.pool().write_normal().await?;
564
565 debug!(%actor_id, %version, "acquired write (normal) connection to process fully buffered changes");
566 assert_sometimes!(true, "Corrosion processes fully buffered changes");
567 let booked = {
568 bookie
569 .write("process_fully_buffered(ensure)", actor_id.as_simple())
570 .await
571 .ensure(actor_id)
572 };
573
574 let mut bookedw = booked
575 .write(
576 "process_fully_buffered(booked writer)",
577 actor_id.as_simple(),
578 )
579 .await;
580 debug!(%actor_id, %version, "acquired Booked write lock to process fully buffered changes");
581
582 block_in_place(|| {
583 let last_seq = {
584 match bookedw.partials.get(&version) {
585 Some(PartialVersion { seqs, last_seq, .. }) => {
586 if seqs.gaps(&(CrsqlSeq(0)..=*last_seq)).count() != 0 {
587 error!(%actor_id, %version, "found sequence gaps: {:?}, aborting!", seqs.gaps(&(CrsqlSeq(0)..=*last_seq)).collect::<RangeInclusiveSet<CrsqlSeq>>());
588 return Ok(false);
590 }
591 *last_seq
592 }
593 None => {
594 warn!(%actor_id, %version, "version not found in cache, returning");
595 return Ok(false);
596 }
597 }
598 };
599
600 let base_tx = conn
601 .immediate_transaction()
602 .map_err(|source| ChangeError::Rusqlite {
603 source,
604 actor_id: Some(actor_id),
605 version: Some(version),
606 })?;
607
608 let tx = InterruptibleTransaction::new(
609 base_tx,
610 Some(tx_timeout),
611 "process_buffered_changes",
612 );
613
614 info!(%actor_id, %version, "Processing buffered changes to crsql_changes (actor: {actor_id}, version: {version}, last_seq: {last_seq})");
615
616 let rows_present: bool = tx.prepare_cached("SELECT EXISTS (SELECT 1 FROM __corro_buffered_changes WHERE site_id = ? AND db_version = ?)")
617 .map_err(|source| ChangeError::Rusqlite{source, actor_id: Some(actor_id), version: Some(version)})?
618 .query_row(params![actor_id, version], |row| row.get(0))
619 .map_err(|source| ChangeError::Rusqlite{source, actor_id: Some(actor_id), version: Some(version)})?;
620
621 let start = Instant::now();
622
623 if rows_present {
624 let count = tx
626 .prepare_cached(
627 r#"
628 INSERT INTO crsql_changes ("table", pk, cid, val, col_version, db_version, site_id, cl, seq, ts)
629 SELECT "table", pk, cid, val, col_version, db_version, site_id, cl, seq, ts
630 FROM __corro_buffered_changes
631 WHERE site_id = ?
632 AND db_version = ?
633 ORDER BY db_version ASC, seq ASC
634 "#,
635 ).map_err(|source| ChangeError::Rusqlite{source, actor_id: Some(actor_id), version: Some(version)})?
636 .execute(params![actor_id.as_bytes(), version]).map_err(|source| ChangeError::Rusqlite{source, actor_id: Some(actor_id), version: Some(version)})?;
637 info!(%actor_id, %version, "Inserted {count} rows from buffered into crsql_changes in {:?}", start.elapsed());
638 } else {
639 info!(%actor_id, %version, "No buffered rows, skipped insertion into crsql_changes");
640 }
641
642 if let Err(e) = agent
643 .tx_clear_buf()
644 .try_send((actor_id, CrsqlDbVersionRange::single(version)))
645 {
646 error!("could not schedule buffered data clear: {e}");
647 }
648
649 let rows_impacted: i64 = tx
650 .prepare_cached("SELECT crsql_rows_impacted()")
651 .map_err(|source| ChangeError::Rusqlite {
652 source,
653 actor_id: Some(actor_id),
654 version: Some(version),
655 })?
656 .query_row((), |row| row.get(0))
657 .map_err(|source| ChangeError::Rusqlite {
658 source,
659 actor_id: Some(actor_id),
660 version: Some(version),
661 })?;
662
663 debug!(%actor_id, %version, "rows impacted by buffered changes insertion: {rows_impacted}");
664
665 let mut snap = bookedw.snapshot();
666 snap.insert_db(&tx, [version..=version].into())
667 .map_err(|source| ChangeError::Rusqlite {
668 source,
669 actor_id: Some(actor_id),
670 version: Some(version),
671 })?;
672
673 tx.commit().map_err(|source| ChangeError::Rusqlite {
674 source,
675 actor_id: Some(actor_id),
676 version: Some(version),
677 })?;
678
679 bookedw.commit_snapshot(snap);
680
681 Ok::<_, ChangeError>(rows_impacted > 0)
682 })
683 }?;
684
685 if rows_impacted {
686 let conn = agent.pool().read().await?;
687 block_in_place(|| {
688 if let Err(e) =
689 match_changes_from_db_version(agent.subs_manager(), &conn, version, actor_id)
690 {
691 error!(%version, "could not match changes for subs from db version: {e}");
692 }
693 });
694
695 block_in_place(|| {
696 if let Err(e) =
697 match_changes_from_db_version(agent.updates_manager(), &conn, version, actor_id)
698 {
699 error!(%version, "could not match changes for updates from db version: {e}");
700 }
701 });
702 }
703
704 Ok(rows_impacted)
705}
706
707#[tracing::instrument(skip(agent, bookie, changes), err)]
708pub async fn process_multiple_changes(
709 agent: Agent,
710 bookie: Bookie,
711 changes: Vec<(ChangeV1, ChangeSource, Instant)>,
712 tx_timeout: Duration,
713) -> Result<(), ChangeError> {
714 let start = Instant::now();
715 counter!("corro.agent.changes.processing.started").increment(changes.len() as u64);
716 debug!(self_actor_id = %agent.actor_id(), "processing multiple changes, len: {}", changes.iter().map(|(change, _, _)| cmp::max(change.len(), 1)).sum::<usize>());
717 trace!(self_actor_id = %agent.actor_id(), "changes: {changes:?}");
718
719 const PROCESSING_WARN_THRESHOLD: Duration = Duration::from_secs(5);
720
721 let mut seen = HashSet::new();
722 let mut unknown_changes: BTreeMap<_, Vec<_>> = BTreeMap::new();
723 for (change, src, queued_at) in changes {
724 histogram!("corro.agent.changes.queued.seconds").record(queued_at.elapsed());
725 let versions = change.versions();
726 let seqs = change.seqs();
727
728 if !seen.insert((change.actor_id, versions, seqs)) {
729 continue;
730 }
731
732 let booked_writer = {
733 bookie
734 .write(
735 "process_multiple_changes(ensure)",
736 change.actor_id.as_simple(),
737 )
738 .await
739 .ensure(change.actor_id)
740 };
741 if booked_writer
742 .read(
743 "process_multiple_changes(contains_all?)",
744 change.actor_id.as_simple(),
745 )
746 .await
747 .contains_all(change.versions(), change.seqs())
748 {
749 continue;
750 }
751
752 unknown_changes
753 .entry(change.actor_id)
754 .or_default()
755 .push((change, src));
756 }
757 let elapsed = start.elapsed();
758 if elapsed >= PROCESSING_WARN_THRESHOLD {
759 warn!("process_multiple_changes: removing duplicates took too long - {elapsed:?}");
760 }
761
762 let mut conn = agent.pool().write_normal().await?;
763
764 let changesets = block_in_place(|| {
765 let start = Instant::now();
766 let tx = conn
767 .immediate_transaction()
768 .map_err(|source| ChangeError::Rusqlite {
769 source,
770 actor_id: None,
771 version: None,
772 })?;
773
774 let mut tx =
775 InterruptibleTransaction::new(tx, Some(tx_timeout), "process_multiple_changes");
776 let mut processed: BTreeMap<ActorId, Vec<_>> = BTreeMap::new();
777 let mut changesets = vec![];
778
779 let mut count = 0;
780
781 let sub_start = Instant::now();
782 for (actor_id, changes) in unknown_changes {
783 let booked = {
784 bookie
785 .blocking_write(
786 "process_multiple_changes(for_actor_blocking)",
787 actor_id.as_simple(),
788 )
789 .ensure(actor_id)
790 };
791 let booked_write = booked.blocking_write(
792 "process_multiple_changes(booked writer, unknown changes)",
793 actor_id.as_simple(),
794 );
795
796 let max = booked_write.last();
797 let mut seen = RangeInclusiveMap::new();
798
799 for (change, src) in changes {
800 trace!("handling a single changeset: {change:?}");
801 let seqs = change.seqs();
802 if booked_write.contains_all(change.versions(), change.seqs()) {
803 trace!("previously unknown versions are now deemed known, aborting inserts");
804 continue;
805 }
806
807 let versions = change.versions();
808
809 if versions.clone().all(|version| match seqs {
811 Some(mut check_seqs) => match seen.get(&version) {
812 Some(maybe_partial) => match maybe_partial {
813 Some(PartialVersion { seqs, .. }) => {
814 check_seqs.all(|seq| seqs.contains(&seq))
815 }
816 None => true,
818 },
819 None => false,
820 },
821 None => seen.contains_key(&version),
822 }) {
823 continue;
824 }
825
826 let known = if change.is_complete() && change.is_empty() {
828 let versions = change.versions();
829 let end = versions.end();
830 if Some(end) > max {
833 process_empty_version(&tx, change.actor_id, &end).map_err(|e| {
834 ChangeError::Rusqlite {
835 source: e,
836 actor_id: Some(change.actor_id),
837 version: Some(end),
838 }
839 })?;
840 }
841 KnownDbVersion::Cleared
842 } else {
843 if let Some(seqs) = change.seqs()
844 && seqs.end() < seqs.start()
845 {
846 warn!(%actor_id, versions = ?change.versions(), "received an invalid change, seqs start is greater than seqs end: {seqs:?}");
847 continue;
848 }
849
850 let (known, changeset) = {
851 match process_single_version(&agent, &mut tx, change) {
852 Ok(res) => {
853 count += 1;
854 res
855 }
856 Err(e) => {
857 error!("error processing single version: {e}");
858 if e.sqlite_error_code().is_some_and(|code| {
859 code != rusqlite::ErrorCode::DiskFull
860 && code != rusqlite::ErrorCode::OperationInterrupted
861 }) {
862 let details = json!({"error": e.to_string()});
863 assert_unreachable!("error committing transaction", &details);
864 }
865 if tx.is_autocommit() {
867 error!(
868 "error processing single version: {e} and transaction was rolled back"
869 );
870 return Err(ChangeError::Rusqlite {
871 source: e,
872 actor_id: None,
873 version: None,
874 });
875 } else {
876 error!("error processing single version: {e}");
877 continue;
878 }
879 }
880 }
881 };
882
883 if let KnownDbVersion::Current(CurrentVersion { db_version, .. }) = &known {
884 changesets.push((actor_id, changeset, *db_version, src));
886 }
887
888 known
889 };
890
891 debug!(%actor_id, self_actor_id = %agent.actor_id(), ?versions, "got known to insert: {known:?}");
892 let partial = match known {
893 KnownDbVersion::Partial(partial) => Some(partial),
894 _ => None,
895 };
896
897 seen.insert(versions.into(), partial.clone());
898 processed
899 .entry(actor_id)
900 .or_default()
901 .push((versions, partial));
902 }
903 }
904
905 let elapsed = sub_start.elapsed();
906 if elapsed >= PROCESSING_WARN_THRESHOLD {
907 warn!("process_multiple_changes:: process_single_version took too long - {elapsed:?}");
908 }
909
910 let sub_start = Instant::now();
911 let mut snapshots = BTreeMap::new();
912
913 for (actor_id, processed) in processed.iter() {
914 debug!(%actor_id, self_actor_id = %agent.actor_id(), "processing {} changesets", processed.len());
915
916 let booked = {
917 bookie
918 .blocking_write(
919 "process_multiple_changes(for_actor_blocking)",
920 actor_id.as_simple(),
921 )
922 .ensure(*actor_id)
923 };
924
925 let mut snap = match snapshots.remove(actor_id) {
927 Some(snap) => snap,
928 None => {
929 let booked_write = booked.blocking_write(
930 "process_multiple_changes(booked writer, during processed)",
931 actor_id.as_simple(),
932 );
933 booked_write.snapshot()
934 }
935 };
936
937 snap.insert_db(
938 &tx,
939 processed
940 .iter()
941 .map(|(versions, _)| versions.into())
942 .collect(),
943 )
944 .map_err(|source| ChangeError::Rusqlite {
945 source,
946 actor_id: Some(*actor_id),
947 version: None,
948 })?;
949
950 snapshots.insert(*actor_id, snap);
951 }
952
953 debug!("inserted {count} new changesets");
954
955 tx.commit().map_err(|source| {
956 if source.sqlite_error_code().is_some_and(|code| {
958 code != rusqlite::ErrorCode::DiskFull
959 && code != rusqlite::ErrorCode::OperationInterrupted
960 }) {
961 let details =
962 json!({"elapsed": elapsed.as_secs_f32(), "error": source.to_string()});
963 assert_unreachable!("error committing transaction", &details);
964 }
965 ChangeError::Rusqlite {
966 source,
967 actor_id: None,
968 version: None,
969 }
970 })?;
971
972 let elapsed = sub_start.elapsed();
973 if elapsed >= PROCESSING_WARN_THRESHOLD {
974 warn!("process_multiple_changes: commiting transaction took too long - {elapsed:?}");
975 }
976
977 for (_, changeset, _, _) in changesets.iter() {
978 if let Some(ts) = changeset.ts() {
979 let dur = (agent.clock().new_timestamp().get_time() - ts.0).to_duration();
980 histogram!("corro.agent.changes.commit.lag.seconds").record(dur);
981 }
982 }
983
984 debug!("committed {count} changes in {:?}", start.elapsed());
985
986 let sub_start = Instant::now();
987 for (actor_id, processed) in processed {
988 let booked = {
989 bookie
990 .blocking_write(
991 "process_multiple_changes(for_actor_blocking)",
992 actor_id.as_simple(),
993 )
994 .ensure(actor_id)
995 };
996 let mut booked_write = booked.blocking_write(
997 "process_multiple_changes(booked writer, before apply needed)",
998 actor_id.as_simple(),
999 );
1000
1001 if let Some(snap) = snapshots.remove(&actor_id) {
1002 booked_write.commit_snapshot(snap);
1003 }
1004
1005 for (versions, partial) in processed {
1006 let version = versions.start();
1007 if let Some(partial) = partial {
1008 let PartialVersion { seqs, last_seq, .. } =
1009 booked_write.insert_partial(version, partial);
1010
1011 let full_seqs_range = CrsqlSeq(0)..=last_seq;
1012 let gaps_count = seqs.gaps(&full_seqs_range).count();
1013 if gaps_count == 0 {
1014 debug!(%actor_id, %version, "we now have all versions, notifying for background jobber to insert buffered changes! seqs: {seqs:?}, expected full seqs: {full_seqs_range:?}");
1016 let tx_apply = agent.tx_apply().clone();
1017 tokio::spawn(async move {
1018 if let Err(e) = tx_apply.send((actor_id, version)).await {
1019 error!(
1020 "could not send trigger for applying fully buffered changes later: {e}"
1021 );
1022 }
1023 });
1024 } else {
1025 debug!(%actor_id, %version, "still have {gaps_count} gaps in partially buffered seqs: {:?}", seqs.gaps(&full_seqs_range).collect::<Vec<_>>());
1026 }
1027 }
1028 }
1029 }
1030
1031 let elapsed = sub_start.elapsed();
1032 let details = json!({"elapsed": elapsed.as_secs_f32()});
1033 assert_always!(
1034 elapsed < Duration::from_secs(60),
1035 "process_multiple_changes took too long",
1036 &details
1037 );
1038 if elapsed >= PROCESSING_WARN_THRESHOLD {
1039 warn!("process_multiple_changes: commiting snapshots took too long - {elapsed:?}");
1040 }
1041
1042 Ok::<_, ChangeError>(changesets)
1043 })?;
1044
1045 let mut change_chunk_size = 0;
1046
1047 for (_actor_id, changeset, db_version, _src) in changesets {
1048 change_chunk_size += changeset.len();
1049 match_changes(agent.subs_manager(), &changeset, db_version);
1050 match_changes(agent.updates_manager(), &changeset, db_version);
1051 }
1052
1053 histogram!("corro.agent.changes.processing.time.seconds", "source" => "remote")
1054 .record(start.elapsed());
1055 histogram!("corro.agent.changes.processing.chunk_size").record(change_chunk_size as f64);
1056
1057 Ok(())
1058}
1059
1060#[tracing::instrument(skip(tx), err)]
1061pub fn process_empty_version<T: Deref<Target = rusqlite::Connection> + Committable>(
1062 tx: &InterruptibleTransaction<T>,
1063 actor_id: ActorId,
1064 end_version: &CrsqlDbVersion,
1065) -> rusqlite::Result<()> {
1066 let _ = tx
1067 .prepare_cached("SELECT crsql_set_db_version(?, ?)")?
1068 .query_row((actor_id, end_version), |row| row.get::<_, String>(0))?;
1069
1070 Ok(())
1071}
1072
1073#[tracing::instrument(skip(sp, parts), err)]
1074pub fn process_incomplete_version<T: Deref<Target = rusqlite::Connection> + Committable>(
1075 sp: &InterruptibleTransaction<T>,
1076 actor_id: ActorId,
1077 parts: &ChangesetParts,
1078) -> rusqlite::Result<KnownDbVersion> {
1079 let ChangesetParts {
1080 version,
1081 changes,
1082 seqs,
1083 last_seq,
1084 ts,
1085 } = parts;
1086
1087 let mut changes_per_table = BTreeMap::new();
1088
1089 debug!(%actor_id, %version, "incomplete change, seqs: {seqs:?}, last_seq: {last_seq:?}, len: {}", changes.len());
1090 assert_sometimes!(true, "Corrosion processes incomplete changes");
1091 let mut inserted = 0;
1092 for change in changes.iter() {
1093 trace!("buffering change! {change:?}");
1094
1095 let new_insertion = sp
1097 .prepare_cached(
1098 r#"
1099 INSERT INTO __corro_buffered_changes
1100 ("table", pk, cid, val, col_version, db_version, site_id, cl, seq, ts)
1101 VALUES
1102 (:table, :pk, :cid, :val, :col_version, :db_version, :site_id, :cl, :seq, :ts)
1103 ON CONFLICT (site_id, db_version, seq)
1104 DO NOTHING
1105 "#,
1106 )?
1107 .execute(named_params! {
1108 ":table": change.table.as_str(),
1109 ":pk": change.pk,
1110 ":cid": change.cid.as_str(),
1111 ":val": &change.val,
1112 ":col_version": change.col_version,
1113 ":db_version": change.db_version,
1114 ":site_id": &change.site_id,
1115 ":cl": change.cl,
1116 ":seq": change.seq,
1117 ":ts": ts,
1118 })?;
1119
1120 inserted += new_insertion;
1121
1122 if let Some(counter) = changes_per_table.get_mut(&change.table) {
1123 *counter += 1;
1124 } else {
1125 changes_per_table.insert(change.table.clone(), 1);
1126 }
1127 }
1128
1129 debug!(%actor_id, %version, "buffered {inserted} changes");
1130
1131 let deleted: Vec<RangeInclusive<CrsqlSeq>> = sp
1132 .prepare_cached(
1133 "
1134 DELETE FROM __corro_seq_bookkeeping
1135 WHERE site_id = :actor_id AND db_version = :db_version AND
1136 (
1137 -- [:start]---[start_seq]---[:end]
1138 ( start_seq BETWEEN :start AND :end ) OR
1139
1140 -- [start_seq]---[:start]---[:end]---[end_seq]
1141 ( start_seq <= :start AND end_seq >= :end ) OR
1142
1143 -- [:start]---[start_seq]---[:end]---[end_seq]
1144 ( start_seq <= :end AND end_seq >= :end ) OR
1145
1146 -- [:start]---[end_seq]---[:end]
1147 ( end_seq BETWEEN :start AND :end ) OR
1148
1149 -- ---[:end][start_seq]---[end_seq]
1150 ( start_seq = :end + 1 AND end_seq ) OR
1151
1152 -- [end_seq][:start]---
1153 ( end_seq = :start - 1 )
1154 )
1155 RETURNING start_seq, end_seq
1156 ",
1157 )?
1158 .query_map(
1159 named_params![
1160 ":actor_id": actor_id,
1161 ":db_version": version,
1162 ":start": seqs.start_int(),
1163 ":end": seqs.end_int(),
1164 ],
1165 |row| Ok(row.get(0)?..=row.get(1)?),
1166 )
1167 .and_then(|rows| rows.collect::<rusqlite::Result<Vec<_>>>())?;
1168
1169 let mut new_ranges = RangeInclusiveSet::from_iter(deleted);
1171 new_ranges.insert((*seqs).into());
1172
1173 let details = json!({"new_ranges": new_ranges});
1174 assert_always!(
1175 new_ranges.len() == 1,
1176 "deleted non-contiguous seq ranges!",
1177 &details
1178 );
1179 if new_ranges.len() > 1 {
1181 warn!("deleted non-contiguous seq ranges! {new_ranges:?}");
1182 return Err(rusqlite::Error::StatementChangedRows(new_ranges.len()));
1184 }
1185
1186 for range in new_ranges.clone() {
1188 sp
1189 .prepare_cached(
1190 "
1191 INSERT INTO __corro_seq_bookkeeping (site_id, db_version, start_seq, end_seq, last_seq, ts)
1192 VALUES (?, ?, ?, ?, ?, ?);
1193 ",
1194 )?
1195 .execute(params![actor_id, version, range.start(), range.end(), last_seq, ts])?;
1196 }
1197
1198 for (table_name, count) in changes_per_table {
1199 counter!("corro.changes.committed", "table" => table_name.to_string(), "source" => "remote").increment(count);
1200 }
1201
1202 Ok(KnownDbVersion::Partial(PartialVersion {
1203 seqs: new_ranges,
1204 last_seq: *last_seq,
1205 ts: *ts,
1206 }))
1207}
1208
1209#[tracing::instrument(skip(agent, sp, parts), err)]
1210pub fn process_complete_version<T: Deref<Target = rusqlite::Connection> + Committable>(
1211 agent: Agent,
1212 sp: &InterruptibleTransaction<T>,
1213 actor_id: ActorId,
1214 versions: CrsqlDbVersionRange,
1215 parts: ChangesetParts,
1216) -> rusqlite::Result<(KnownDbVersion, Changeset, BTreeMap<TableName, u64>)> {
1217 let ChangesetParts {
1218 version,
1219 changes,
1220 seqs,
1221 last_seq,
1222 ts,
1223 } = parts;
1224
1225 let len = changes.len();
1226
1227 debug!(%actor_id, %version, "complete change, applying right away! seqs: {seqs:?}, last_seq: {last_seq}, changes len: {len}, db version: {version}");
1228
1229 let details = json!({"len": len, "seqs": seqs.start_int(), "seqs_end": seqs.end_int()});
1230 assert_always!(
1231 len <= seqs.len(),
1232 "number of changes is greater than the seq num",
1233 &details
1234 );
1235 debug_assert!(
1236 len <= seqs.len(),
1237 "change from actor {actor_id} version {version} has len {len} but seqs range is {seqs:?} and last_seq is {last_seq}"
1238 );
1239
1240 let mut impactful_changeset = vec![];
1241
1242 let mut last_rows_impacted = 0;
1243
1244 let mut changes_per_table = BTreeMap::new();
1245
1246 for change in changes {
1247 trace!("inserting change! {change:?}");
1248
1249 sp.prepare_cached(
1250 r#"
1251 INSERT INTO crsql_changes
1252 ("table", pk, cid, val, col_version, db_version, site_id, cl, seq, ts)
1253 VALUES
1254 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
1255 "#,
1256 )?
1257 .execute(params![
1258 change.table.as_str(),
1259 change.pk,
1260 change.cid.as_str(),
1261 &change.val,
1262 change.col_version,
1263 change.db_version,
1264 &change.site_id,
1265 change.cl,
1266 change.seq,
1268 ts,
1269 ])?;
1270 let rows_impacted: i64 = sp
1271 .prepare_cached("SELECT crsql_rows_impacted()")?
1272 .query_row((), |row| row.get(0))?;
1273
1274 if rows_impacted > last_rows_impacted {
1275 trace!("inserted the change into crsql_changes");
1276 impactful_changeset.push(change);
1277 if let Some(c) = impactful_changeset.last() {
1278 if let Some(counter) = changes_per_table.get_mut(&c.table) {
1279 *counter += 1;
1280 } else {
1281 changes_per_table.insert(c.table.clone(), 1);
1282 }
1283 }
1284 }
1285 last_rows_impacted = rows_impacted;
1286 }
1287
1288 let (known_version, new_changeset) = if impactful_changeset.is_empty() {
1289 (
1290 KnownDbVersion::Cleared,
1291 Changeset::Empty {
1292 versions,
1293 ts: Some(Timestamp::from(agent.clock().new_timestamp())),
1294 },
1295 )
1296 } else {
1297 (
1298 KnownDbVersion::Current(CurrentVersion {
1299 db_version: version,
1300 last_seq,
1301 ts,
1302 }),
1303 Changeset::Full {
1304 version,
1305 changes: impactful_changeset,
1306 seqs,
1307 last_seq,
1308 ts,
1309 },
1310 )
1311 };
1312
1313 Ok::<_, rusqlite::Error>((known_version, new_changeset, changes_per_table))
1314}
1315
1316pub fn check_buffered_meta_to_clear(
1317 conn: &Connection,
1318 actor_id: ActorId,
1319 versions: CrsqlDbVersionRange,
1320) -> rusqlite::Result<bool> {
1321 let should_clear: bool = conn.prepare_cached("SELECT EXISTS(SELECT 1 FROM __corro_buffered_changes WHERE site_id = ? AND db_version >= ? AND db_version <= ?)")?.query_row(params![actor_id, versions.start(), versions.end()], |row| row.get(0))?;
1322 if should_clear {
1323 return Ok(true);
1324 }
1325
1326 conn.prepare_cached("SELECT EXISTS(SELECT 1 FROM __corro_seq_bookkeeping WHERE site_id = ? AND db_version >= ? AND db_version <= ?)")?.query_row(params![actor_id, versions.start(), versions.end()], |row| row.get(0))
1327}
1328
1329pub fn log_at_pow_10(msg: &str, count: &mut u64) {
1330 *count += 1;
1331 if is_pow_10(*count) {
1332 warn!("{msg} (log count: {count})")
1333 }
1334 if *count == 100000000 {
1336 *count = 0;
1337 }
1338}
1339
1340#[inline]
1341fn is_pow_10(i: u64) -> bool {
1342 matches!(
1343 i,
1344 1 | 10 | 100 | 1000 | 10000 | 1000000 | 10000000 | 100000000
1345 )
1346}
1347
1348pub async fn read_files_from_paths<P: AsRef<Path>>(
1349 schema_paths: &[P],
1350) -> eyre::Result<Vec<String>> {
1351 let mut contents = vec![];
1352
1353 for schema_path in schema_paths.iter() {
1354 match tokio::fs::metadata(schema_path).await {
1355 Ok(meta) => {
1356 if meta.is_dir() {
1357 match tokio::fs::read_dir(schema_path).await {
1358 Ok(mut dir) => {
1359 let mut entries = vec![];
1360
1361 while let Ok(Some(entry)) = dir.next_entry().await {
1362 entries.push(entry);
1363 }
1364
1365 let mut entries: Vec<_> = entries
1366 .into_iter()
1367 .filter_map(|entry| {
1368 entry.path().extension().and_then(|ext| {
1369 if ext == "sql" { Some(entry) } else { None }
1370 })
1371 })
1372 .collect();
1373
1374 entries.sort_by_key(|entry| entry.path());
1375
1376 for entry in entries.iter() {
1377 match tokio::fs::read_to_string(entry.path()).await {
1378 Ok(s) => {
1379 contents.push(s);
1380 }
1381 Err(e) => {
1382 warn!(
1383 "could not read schema file '{}', error: {e}",
1384 entry.path().display()
1385 );
1386 }
1387 }
1388 }
1389 }
1390 Err(e) => {
1391 warn!(
1392 "could not read dir '{}', error: {e}",
1393 schema_path.as_ref().display()
1394 );
1395 }
1396 }
1397 } else if meta.is_file() {
1398 match tokio::fs::read_to_string(schema_path).await {
1399 Ok(s) => {
1400 contents.push(s);
1401 }
1403 Err(e) => {
1404 warn!(
1405 "could not read schema file '{}', error: {e}",
1406 schema_path.as_ref().display()
1407 );
1408 }
1409 }
1410 }
1411 }
1412
1413 Err(e) => {
1414 warn!(
1415 "could not read schema file meta '{}', error: {e}",
1416 schema_path.as_ref().display()
1417 );
1418 }
1419 }
1420 }
1421
1422 Ok(contents)
1423}