Skip to main content

klukai_types/
agent.rs

1use std::{
2    cmp,
3    collections::{BTreeMap, HashMap, HashSet, btree_map},
4    fmt,
5    future::Future,
6    io,
7    net::SocketAddr,
8    ops::{Deref, DerefMut, RangeInclusive},
9    path::{Path, PathBuf},
10    sync::{
11        Arc,
12        atomic::{AtomicUsize, Ordering},
13    },
14    time::{Duration, Instant},
15};
16
17use antithesis_sdk::{assert_always, assert_unreachable};
18use arc_swap::ArcSwap;
19use camino::Utf8PathBuf;
20use compact_str::{CompactString, ToCompactString};
21use indexmap::IndexMap;
22use metrics::{gauge, histogram};
23use parking_lot::RwLock;
24use rangemap::RangeInclusiveSet;
25use rusqlite::{Connection, OpenFlags, OptionalExtension, Transaction, named_params};
26use serde::{Deserialize, Serialize};
27use serde_json::json;
28use tokio::{
29    runtime::Handle,
30    sync::{Semaphore, oneshot},
31};
32use tokio::{
33    sync::{
34        AcquireError, OwnedRwLockWriteGuard as OwnedTokioRwLockWriteGuard, OwnedSemaphorePermit,
35        RwLock as TokioRwLock, RwLockReadGuard as TokioRwLockReadGuard,
36        RwLockWriteGuard as TokioRwLockWriteGuard,
37    },
38    time::timeout,
39};
40use tokio_util::sync::{CancellationToken, DropGuard};
41use tracing::{debug, error, trace, warn};
42
43use crate::tripwire::Tripwire;
44
45use crate::{
46    actor::{Actor, ActorId, ClusterId},
47    base::{CrsqlDbVersion, CrsqlDbVersionRange, CrsqlSeq, CrsqlSeqRange},
48    broadcast::{BroadcastInput, ChangeSource, ChangeV1, FocaInput, Timestamp},
49    channel::{CorroSender, bounded},
50    config::Config,
51    pubsub::SubsManager,
52    schema::Schema,
53    sqlite::{
54        CrConn, Migration, SqlitePool, SqlitePoolError, rusqlite_to_crsqlite,
55        rusqlite_to_crsqlite_write, setup_conn,
56    },
57    sqlite_pool,
58    updates::UpdatesManager,
59};
60
61use super::members::Members;
62
63#[derive(Clone)]
64pub struct Agent(Arc<AgentInner>);
65
66pub struct AgentConfig {
67    pub actor_id: ActorId,
68    pub pool: SplitPool,
69    pub config: ArcSwap<Config>,
70    pub gossip_addr: SocketAddr,
71    pub external_addr: Option<SocketAddr>,
72    pub api_addr: SocketAddr,
73    pub members: RwLock<Members>,
74    pub clock: Arc<uhlc::HLC>,
75
76    pub booked: Booked,
77
78    pub tx_bcast: CorroSender<BroadcastInput>,
79    pub tx_apply: CorroSender<(ActorId, CrsqlDbVersion)>,
80    pub tx_clear_buf: CorroSender<(ActorId, CrsqlDbVersionRange)>,
81    pub tx_changes: CorroSender<(ChangeV1, ChangeSource)>,
82    pub tx_foca: CorroSender<FocaInput>,
83
84    pub write_sema: Arc<Semaphore>,
85
86    pub schema: RwLock<Schema>,
87    pub cluster_id: ClusterId,
88
89    pub subs_manager: SubsManager,
90
91    pub updates_manager: UpdatesManager,
92
93    pub tripwire: Tripwire,
94}
95
96pub struct AgentInner {
97    actor_id: ActorId,
98    pool: SplitPool,
99    config: ArcSwap<Config>,
100    gossip_addr: SocketAddr,
101    external_addr: Option<SocketAddr>,
102    api_addr: SocketAddr,
103    members: RwLock<Members>,
104    clock: Arc<uhlc::HLC>,
105    booked: Booked,
106    tx_bcast: CorroSender<BroadcastInput>,
107    tx_apply: CorroSender<(ActorId, CrsqlDbVersion)>,
108    tx_clear_buf: CorroSender<(ActorId, CrsqlDbVersionRange)>,
109    tx_changes: CorroSender<(ChangeV1, ChangeSource)>,
110    tx_foca: CorroSender<FocaInput>,
111    write_sema: Arc<Semaphore>,
112    schema: RwLock<Schema>,
113    cluster_id: ArcSwap<ClusterId>,
114    limits: Limits,
115    subs_manager: SubsManager,
116    updates_manager: UpdatesManager,
117}
118
119#[derive(Debug, Clone)]
120pub struct Limits {
121    pub sync: Arc<Semaphore>,
122}
123
124impl Agent {
125    pub fn new(config: AgentConfig) -> Self {
126        Self(Arc::new(AgentInner {
127            actor_id: config.actor_id,
128            pool: config.pool,
129            config: config.config,
130            gossip_addr: config.gossip_addr,
131            external_addr: config.external_addr,
132            api_addr: config.api_addr,
133            members: config.members,
134            clock: config.clock,
135            booked: config.booked,
136            tx_bcast: config.tx_bcast,
137            tx_apply: config.tx_apply,
138            tx_clear_buf: config.tx_clear_buf,
139            tx_changes: config.tx_changes,
140            tx_foca: config.tx_foca,
141            write_sema: config.write_sema,
142            schema: config.schema,
143            cluster_id: ArcSwap::from_pointee(config.cluster_id),
144            limits: Limits {
145                sync: Arc::new(Semaphore::new(3)),
146            },
147            subs_manager: config.subs_manager,
148            updates_manager: config.updates_manager,
149        }))
150    }
151
152    pub fn actor<C: Into<Option<ClusterId>>>(&self, cluster_id: C) -> Actor {
153        Actor::new(
154            self.0.actor_id,
155            self.external_addr().unwrap_or_else(|| self.gossip_addr()),
156            self.clock().new_timestamp().into(),
157            cluster_id.into().unwrap_or_else(|| self.cluster_id()),
158        )
159    }
160
161    /// Return a borrowed [SqlitePool]
162    pub fn pool(&self) -> &SplitPool {
163        &self.0.pool
164    }
165
166    pub fn actor_id(&self) -> ActorId {
167        self.0.actor_id
168    }
169
170    pub fn clock(&self) -> &Arc<uhlc::HLC> {
171        &self.0.clock
172    }
173
174    pub fn gossip_addr(&self) -> SocketAddr {
175        self.0.gossip_addr
176    }
177
178    pub fn external_addr(&self) -> Option<SocketAddr> {
179        self.0.external_addr
180    }
181
182    pub fn api_addr(&self) -> SocketAddr {
183        self.0.api_addr
184    }
185
186    pub fn tx_bcast(&self) -> &CorroSender<BroadcastInput> {
187        &self.0.tx_bcast
188    }
189
190    pub fn tx_apply(&self) -> &CorroSender<(ActorId, CrsqlDbVersion)> {
191        &self.0.tx_apply
192    }
193
194    pub fn tx_changes(&self) -> &CorroSender<(ChangeV1, ChangeSource)> {
195        &self.0.tx_changes
196    }
197
198    pub fn tx_clear_buf(&self) -> &CorroSender<(ActorId, CrsqlDbVersionRange)> {
199        &self.0.tx_clear_buf
200    }
201
202    pub fn tx_foca(&self) -> &CorroSender<FocaInput> {
203        &self.0.tx_foca
204    }
205
206    pub fn write_sema(&self) -> &Arc<Semaphore> {
207        &self.0.write_sema
208    }
209
210    pub async fn write_permit(&self) -> Result<OwnedSemaphorePermit, AcquireError> {
211        self.0.write_sema.clone().acquire_owned().await
212    }
213
214    pub fn write_permit_blocking(&self) -> Result<OwnedSemaphorePermit, AcquireError> {
215        Handle::current().block_on(self.0.write_sema.clone().acquire_owned())
216    }
217
218    pub fn booked(&self) -> &Booked {
219        &self.0.booked
220    }
221
222    pub fn members(&self) -> &RwLock<Members> {
223        &self.0.members
224    }
225
226    pub fn schema(&self) -> &RwLock<Schema> {
227        &self.0.schema
228    }
229
230    pub fn db_path(&self) -> Utf8PathBuf {
231        self.0.config.load().db.path.clone()
232    }
233
234    pub fn config(&self) -> arc_swap::Guard<Arc<Config>, arc_swap::strategy::DefaultStrategy> {
235        self.0.config.load()
236    }
237
238    pub fn set_config(&self, new_conf: Config) {
239        self.0.config.store(Arc::new(new_conf))
240    }
241
242    pub fn limits(&self) -> &Limits {
243        &self.0.limits
244    }
245
246    pub fn subs_manager(&self) -> &SubsManager {
247        &self.0.subs_manager
248    }
249
250    pub fn updates_manager(&self) -> &UpdatesManager {
251        &self.0.updates_manager
252    }
253
254    pub fn set_cluster_id(&self, cluster_id: ClusterId) {
255        self.0.cluster_id.store(Arc::new(cluster_id));
256    }
257
258    pub fn cluster_id(&self) -> ClusterId {
259        *self.0.cluster_id.load().as_ref()
260    }
261
262    pub fn update_clock_with_timestamp(
263        &self,
264        actor_id: ActorId,
265        ts: Timestamp,
266    ) -> Result<(), String> {
267        let id = actor_id
268            .try_into()
269            .map_err(|e| format!("could not convert ActorId to uhlc ID: {e}"))?;
270        self.clock()
271            .update_with_timestamp(&uhlc::Timestamp::new(ts.0, id))
272    }
273}
274
275pub fn migrate(clock: Arc<uhlc::HLC>, conn: &mut Connection) -> rusqlite::Result<()> {
276    let migrations: Vec<Box<dyn Migration>> = vec![
277        Box::new(init_migration as fn(&Transaction) -> rusqlite::Result<()>),
278        Box::new(crsqlite_v0_17_migration(clock)),
279    ];
280
281    crate::sqlite::migrate(conn, migrations)
282}
283
284fn init_migration(tx: &Transaction) -> rusqlite::Result<()> {
285    tx.execute_batch(
286        r#"
287            -- key/value for internal corrosion data (e.g. 'schema_version' => INT)
288            -- can be pre-created before corrosion starts so we get a cluster id.
289            CREATE TABLE IF NOT EXISTS __corro_state (key TEXT NOT NULL PRIMARY KEY, value);
290
291            -- internal per-db-version seq bookkeeping
292            CREATE TABLE  __corro_seq_bookkeeping (
293                -- remote actor / site id
294                site_id BLOB NOT NULL,
295                -- remote internal version
296                db_version INTEGER NOT NULL,
297
298                -- start and end seq for this bookkept record
299                start_seq INTEGER NOT NULL,
300                end_seq INTEGER NOT NULL,
301
302                last_seq INTEGER NOT NULL,
303
304                -- timestamp, need to propagate...
305                ts TEXT NOT NULL,
306
307                PRIMARY KEY (site_id, db_version, start_seq)
308            ) WITHOUT ROWID;
309
310            -- buffered changes (similar schema as crsql_changes)
311            CREATE TABLE __corro_buffered_changes (
312                "table" TEXT NOT NULL,
313                pk BLOB NOT NULL,
314                cid TEXT NOT NULL,
315                val ANY, -- shouldn't matter I don't think
316                col_version INTEGER NOT NULL,
317                db_version INTEGER NOT NULL,
318                site_id BLOB NOT NULL, -- this differs from crsql_changes, we'll never buffer our own
319                seq INTEGER NOT NULL,
320                cl INTEGER NOT NULL, -- causal length
321                ts TEXT NOT NULL,
322
323                PRIMARY KEY (site_id, db_version, seq)
324            ) WITHOUT ROWID;
325
326            -- SWIM memberships
327            CREATE TABLE __corro_members (
328                actor_id BLOB PRIMARY KEY NOT NULL,
329                address TEXT NOT NULL,
330
331                foca_state JSON,
332
333                rtt_min INTEGER,
334                updated_at DATETIME NOT NULL DEFAULT 0
335            ) WITHOUT ROWID;
336
337            -- tracked corrosion schema
338            CREATE TABLE IF NOT EXISTS __corro_schema (
339                tbl_name TEXT NOT NULL,
340                type TEXT NOT NULL,
341                name TEXT NOT NULL,
342                sql TEXT NOT NULL,
343
344                source TEXT NOT NULL,
345
346                PRIMARY KEY (tbl_name, type, name)
347            ) WITHOUT ROWID;
348
349            -- store known needed versions
350            CREATE TABLE __corro_bookkeeping_gaps (
351                actor_id BLOB NOT NULL,
352                start INTEGER NOT NULL,
353                end INTEGER NOT NULL,
354
355                PRIMARY KEY (actor_id, start)
356            ) WITHOUT ROWID;
357        "#,
358    )?;
359
360    let _value: i64 = tx.query_row(
361        "SELECT crsql_config_set('merge-equal-values', 1);",
362        [],
363        |row| row.get(0),
364    )?;
365
366    Ok(())
367}
368
369// since crsqlite 0.17, ts is now stored as TEXT in clock tables
370// also sets the new 'merge-equal-values' config to true.
371fn crsqlite_v0_17_migration(
372    clock: Arc<uhlc::HLC>,
373) -> impl Fn(&Transaction) -> rusqlite::Result<()> {
374    let ts = Timestamp::from(clock.new_timestamp()).as_u64().to_string();
375
376    move |tx: &Transaction| -> rusqlite::Result<()> {
377        let tables: Vec<String> = tx.prepare("SELECT tbl_name FROM sqlite_master WHERE type='table' AND tbl_name LIKE '%__crsql_clock'")?.query_map([], |row| row.get(0))?.collect::<rusqlite::Result<Vec<_>>>()?;
378
379        for table in tables {
380            let indexes: Vec<String> = tx
381                .prepare(&format!(
382                    "SELECT sql FROM sqlite_master WHERE type='index' AND name LIKE '{table}%'"
383                ))?
384                .query_map([], |row| row.get(0))?
385                .collect::<rusqlite::Result<Vec<_>>>()?;
386
387            tx.execute_batch(
388                &format!(r#"
389                    CREATE TABLE {table}_new (
390                        key INTEGER NOT NULL,
391                        col_name TEXT NOT NULL,
392                        col_version INTEGER NOT NULL,
393                        db_version INTEGER NOT NULL,
394                        site_id INTEGER NOT NULL DEFAULT 0,
395                        seq INTEGER NOT NULL,
396                        ts TEXT NOT NULL DEFAULT '0',
397                        PRIMARY KEY (key, col_name)
398                    ) WITHOUT ROWID, STRICT;
399
400                    INSERT INTO {table}_new (key, col_name, col_version, db_version, site_id, seq, ts) SELECT key, col_name, col_version, db_version, site_id, seq, '{ts}' FROM {table};
401
402                    ALTER TABLE {table} RENAME TO {table}_old;
403                    ALTER TABLE {table}_new RENAME TO {table};
404
405                    DROP TABLE {table}_old;
406
407                    CREATE INDEX IF NOT EXISTS corro_{table}_site_id_dbv ON {table} (site_id, db_version);
408                "#),
409            )?;
410
411            // recreate the indexes
412            for sql in indexes {
413                tx.execute_batch(&sql)?;
414            }
415        }
416
417        Ok(())
418    }
419}
420
421#[derive(Debug, Clone)]
422pub struct SplitPool(Arc<SplitPoolInner>);
423
424#[derive(Debug)]
425struct SplitPoolInner {
426    path: PathBuf,
427    write_sema: Arc<Semaphore>,
428
429    read: SqlitePool,
430    write: SqlitePool,
431
432    priority_tx: CorroSender<oneshot::Sender<DropGuard>>,
433    normal_tx: CorroSender<oneshot::Sender<DropGuard>>,
434    low_tx: CorroSender<oneshot::Sender<DropGuard>>,
435}
436
437#[derive(Debug, thiserror::Error)]
438pub enum PoolError {
439    #[error(transparent)]
440    Pool(#[from] SqlitePoolError),
441    #[error("queue is closed")]
442    QueueClosed,
443    #[error("callback is closed")]
444    CallbackClosed,
445    #[error("could not acquire write permit")]
446    Permit(#[from] AcquireError),
447    #[error("timed out acquiring write permit while {op:?}")]
448    TimedOut { op: String },
449}
450
451#[derive(Debug, thiserror::Error)]
452pub enum ChangeError {
453    #[error("could not acquire pooled write connection: {0}")]
454    Pool(#[from] PoolError),
455    #[error("could not acquire pooled read connection: {0}")]
456    SqlitePool(#[from] SqlitePoolError),
457    #[error("rusqlite: {source} (actor_id: {actor_id:?}, version: {version:?})")]
458    Rusqlite {
459        source: rusqlite::Error,
460        actor_id: Option<ActorId>,
461        version: Option<CrsqlDbVersion>,
462    },
463    #[error("non-contiguous empties range delete")]
464    NonContiguousDelete,
465}
466
467#[derive(Debug, thiserror::Error)]
468pub enum SplitPoolCreateError {
469    #[error(transparent)]
470    Pool(#[from] sqlite_pool::CreatePoolError),
471    #[error(transparent)]
472    Io(#[from] io::Error),
473    #[error(transparent)]
474    Rusqlite(#[from] rusqlite::Error),
475}
476
477impl SplitPool {
478    pub async fn create<P: AsRef<Path>>(
479        path: P,
480        write_sema: Arc<Semaphore>,
481    ) -> Result<Self, SplitPoolCreateError> {
482        let rw_pool = sqlite_pool::Config::new(path.as_ref())
483            .max_size(1)
484            .create_pool_transform(rusqlite_to_crsqlite_write)?;
485
486        debug!("built RW pool");
487
488        let ro_pool = sqlite_pool::Config::new(path.as_ref())
489            .read_only()
490            .max_size(20)
491            .create_pool_transform(rusqlite_to_crsqlite)?;
492        debug!("built RO pool");
493
494        Ok(Self::new(
495            path.as_ref().to_owned(),
496            write_sema,
497            ro_pool,
498            rw_pool,
499        ))
500    }
501
502    fn new(path: PathBuf, write_sema: Arc<Semaphore>, read: SqlitePool, write: SqlitePool) -> Self {
503        let (priority_tx, mut priority_rx) = bounded(256, "priority");
504        let (normal_tx, mut normal_rx) = bounded(512, "normal");
505        let (low_tx, mut low_rx) = bounded(1024, "low");
506
507        tokio::spawn(async move {
508            loop {
509                let (tx, channel) = tokio::select! {
510                    biased;
511
512                    Some(tx) = priority_rx.recv() => (tx, "priority"),
513                    Some(tx) = normal_rx.recv() => (tx, "normal"),
514                    Some(tx) = low_rx.recv() => (tx, "low"),
515                };
516
517                wait_conn_drop(tx, channel).await
518            }
519        });
520
521        Self(Arc::new(SplitPoolInner {
522            path,
523            write_sema,
524            read,
525            write,
526            priority_tx,
527            normal_tx,
528            low_tx,
529        }))
530    }
531
532    pub fn emit_metrics(&self) {
533        let read_state = self.0.read.status();
534        gauge!("corro.sqlite.pool.read.connections").set(read_state.size as f64);
535        gauge!("corro.sqlite.pool.read.connections.available").set(read_state.available as f64);
536        gauge!("corro.sqlite.pool.read.connections.waiting").set(read_state.waiting as f64);
537
538        let write_state = self.0.write.status();
539        gauge!("corro.sqlite.pool.write.connections").set(write_state.size as f64);
540        gauge!("corro.sqlite.pool.write.connections.available").set(write_state.available as f64);
541        gauge!("corro.sqlite.pool.write.connections.waiting").set(write_state.waiting as f64);
542
543        let available_permit = self.0.write_sema.available_permits();
544        gauge!("corro.sqlite.write.permits.available").set(available_permit as f64);
545    }
546
547    // get a read-only connection
548    #[tracing::instrument(skip(self), level = "debug")]
549    pub async fn read(&self) -> Result<sqlite_pool::Connection<CrConn>, SqlitePoolError> {
550        self.0.read.get().await
551    }
552
553    #[tracing::instrument(skip(self), level = "debug")]
554    pub fn read_blocking(&self) -> Result<sqlite_pool::Connection<CrConn>, SqlitePoolError> {
555        Handle::current().block_on(self.0.read.get())
556    }
557
558    pub fn drain_read(&self) {
559        self.0.read.retain(|_, _| false);
560    }
561
562    #[tracing::instrument(skip(self), level = "debug")]
563    pub fn dedicated(&self) -> rusqlite::Result<Connection> {
564        let conn = rusqlite::Connection::open(&self.0.path)?;
565        setup_conn(&conn)?;
566        Ok(conn)
567    }
568
569    #[tracing::instrument(skip(self), level = "debug")]
570    pub fn client_dedicated(&self) -> rusqlite::Result<CrConn> {
571        let conn = rusqlite::Connection::open(&self.0.path)?;
572        rusqlite_to_crsqlite_write(conn)
573    }
574
575    #[tracing::instrument(skip(self), level = "debug")]
576    pub fn client_dedicated_readonly(&self) -> rusqlite::Result<CrConn> {
577        let conn = rusqlite::Connection::open_with_flags(
578            &self.0.path,
579            OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
580        )?;
581        rusqlite_to_crsqlite(conn)
582    }
583
584    // get a high priority write connection (e.g. client input)
585    #[tracing::instrument(skip(self), level = "debug")]
586    pub async fn write_priority(&self) -> Result<WriteConn, PoolError> {
587        self.write_inner(&self.0.priority_tx, "priority").await
588    }
589
590    // get a normal priority write connection (e.g. sync process)
591    #[tracing::instrument(skip(self), level = "debug")]
592    pub async fn write_normal(&self) -> Result<WriteConn, PoolError> {
593        self.write_inner(&self.0.normal_tx, "normal").await
594    }
595
596    // get a low priority write connection (e.g. background tasks)
597    #[tracing::instrument(skip(self), level = "debug")]
598    pub async fn write_low(&self) -> Result<WriteConn, PoolError> {
599        self.write_inner(&self.0.low_tx, "low").await
600    }
601
602    async fn write_inner(
603        &self,
604        chan: &CorroSender<oneshot::Sender<DropGuard>>,
605        queue: &'static str,
606    ) -> Result<WriteConn, PoolError> {
607        let (tx, rx) = oneshot::channel();
608        let max_timeout = Duration::from_secs(5 * 60);
609
610        timeout_fut("tx to oneshot channel", max_timeout, chan.send(tx))
611            .await?
612            .map_err(|_| PoolError::QueueClosed)?;
613
614        let start = Instant::now();
615
616        let _drop_guard = timeout_fut("rx from oneshot channel", max_timeout, rx)
617            .await?
618            .map_err(|_| PoolError::CallbackClosed)?;
619
620        histogram!("corro.sqlite.pool.queue.seconds", "queue" => queue)
621            .record(start.elapsed().as_secs_f64());
622        let conn = timeout_fut("acquiring write conn", max_timeout, self.0.write.get()).await??;
623
624        let start = Instant::now();
625        let _permit = timeout_fut(
626            "acquiring write semaphore",
627            max_timeout,
628            self.0.write_sema.clone().acquire_owned(),
629        )
630        .await??;
631
632        histogram!("corro.sqlite.write_permit.acquisition.seconds")
633            .record(start.elapsed().as_secs_f64());
634
635        Ok(WriteConn {
636            conn,
637            _drop_guard,
638            _permit,
639        })
640    }
641}
642
643async fn wait_conn_drop(tx: oneshot::Sender<DropGuard>, channel: &'static str) {
644    let cancel = CancellationToken::new();
645
646    if let Err(_e) = tx.send(cancel.clone().drop_guard()) {
647        error!("could not send back drop guard for pooled conn, oneshot channel likely closed");
648        return;
649    }
650
651    let mut interval = tokio::time::interval(Duration::from_secs(5 * 60));
652    // skip first tick
653    interval.tick().await;
654    let start = Instant::now();
655
656    loop {
657        tokio::select! {
658            _ = cancel.cancelled() => {
659                break;
660            }
661            _ = interval.tick() => {
662                let details = json!({
663                    "channel": channel,
664                    "elapsed": start.elapsed().as_secs_f64()
665                });
666                assert_unreachable!(
667                    "wait_conn_drop has been running for too long",
668                    &details
669                );
670
671                let elapsed = start.elapsed();
672                warn!("wait_conn_drop has been running since {elapsed:?}, token_is_cancelled - {:?}, channel - {channel}", cancel.is_cancelled());
673                continue;
674            }
675        }
676    }
677}
678
679async fn timeout_fut<T, F>(op: &'static str, duration: Duration, fut: F) -> Result<T, PoolError>
680where
681    F: Future<Output = T>,
682{
683    timeout(duration, fut)
684        .await
685        .map_err(|_| PoolError::TimedOut { op: op.to_string() })
686}
687
688pub struct WriteConn {
689    conn: sqlite_pool::Connection<CrConn>,
690    _drop_guard: DropGuard,
691    _permit: OwnedSemaphorePermit,
692}
693
694impl Deref for WriteConn {
695    type Target = sqlite_pool::Connection<CrConn>;
696
697    fn deref(&self) -> &Self::Target {
698        &self.conn
699    }
700}
701
702impl DerefMut for WriteConn {
703    fn deref_mut(&mut self) -> &mut Self::Target {
704        &mut self.conn
705    }
706}
707pub struct CountedTokioRwLock<T> {
708    registry: LockRegistry,
709    lock: Arc<TokioRwLock<T>>,
710}
711
712impl<T> CountedTokioRwLock<T> {
713    fn new(registry: LockRegistry, value: T) -> Self {
714        Self {
715            registry,
716            lock: Arc::new(TokioRwLock::new(value)),
717        }
718    }
719
720    #[tracing::instrument(skip_all, level = "debug")]
721    pub async fn write<C: fmt::Display, E: Into<Option<C>>>(
722        &self,
723        label: &'static str,
724        extra: E,
725    ) -> CountedTokioRwLockWriteGuard<'_, T> {
726        self.registry.acquire_write(label, extra, &self.lock).await
727    }
728
729    #[tracing::instrument(skip_all, level = "debug")]
730    pub async fn write_owned<C: fmt::Display, E: Into<Option<C>>>(
731        &self,
732        label: &'static str,
733        extra: E,
734    ) -> CountedOwnedTokioRwLockWriteGuard<T> {
735        self.registry
736            .acquire_write_owned(label, extra, self.lock.clone())
737            .await
738    }
739
740    #[tracing::instrument(skip_all, level = "debug")]
741    pub fn blocking_write<C: fmt::Display, E: Into<Option<C>>>(
742        &self,
743        label: &'static str,
744        extra: E,
745    ) -> CountedTokioRwLockWriteGuard<'_, T> {
746        self.registry
747            .acquire_blocking_write(label, extra, &self.lock)
748    }
749
750    #[tracing::instrument(skip_all, level = "debug")]
751    pub fn blocking_write_owned<C: fmt::Display, E: Into<Option<C>>>(
752        &self,
753        label: &'static str,
754        extra: E,
755    ) -> CountedOwnedTokioRwLockWriteGuard<T> {
756        self.registry
757            .acquire_blocking_write_owned(label, extra, self.lock.clone())
758    }
759
760    #[tracing::instrument(skip_all, level = "debug")]
761    pub fn blocking_read<C: fmt::Display, E: Into<Option<C>>>(
762        &self,
763        label: &'static str,
764        extra: E,
765    ) -> CountedTokioRwLockReadGuard<'_, T> {
766        self.registry
767            .acquire_blocking_read(label, extra, &self.lock)
768    }
769
770    #[tracing::instrument(skip_all, level = "debug")]
771    pub async fn read<C: fmt::Display, E: Into<Option<C>>>(
772        &self,
773        label: &'static str,
774        extra: E,
775    ) -> CountedTokioRwLockReadGuard<'_, T> {
776        self.registry.acquire_read(label, extra, &self.lock).await
777    }
778
779    pub fn registry(&self) -> &LockRegistry {
780        &self.registry
781    }
782}
783
784pub struct CountedTokioRwLockWriteGuard<'a, T> {
785    lock: TokioRwLockWriteGuard<'a, T>,
786    _tracker: LockTracker,
787}
788
789impl<'a, T> Deref for CountedTokioRwLockWriteGuard<'a, T> {
790    type Target = TokioRwLockWriteGuard<'a, T>;
791
792    fn deref(&self) -> &Self::Target {
793        &self.lock
794    }
795}
796
797impl<'a, T> DerefMut for CountedTokioRwLockWriteGuard<'a, T> {
798    fn deref_mut(&mut self) -> &mut Self::Target {
799        &mut self.lock
800    }
801}
802
803pub struct CountedOwnedTokioRwLockWriteGuard<T> {
804    lock: OwnedTokioRwLockWriteGuard<T>,
805    _tracker: LockTracker,
806}
807
808impl<T> Deref for CountedOwnedTokioRwLockWriteGuard<T> {
809    type Target = OwnedTokioRwLockWriteGuard<T>;
810
811    fn deref(&self) -> &Self::Target {
812        &self.lock
813    }
814}
815
816impl<T> DerefMut for CountedOwnedTokioRwLockWriteGuard<T> {
817    fn deref_mut(&mut self) -> &mut Self::Target {
818        &mut self.lock
819    }
820}
821
822pub struct CountedTokioRwLockReadGuard<'a, T> {
823    lock: TokioRwLockReadGuard<'a, T>,
824    _tracker: LockTracker,
825}
826
827impl<'a, T> Deref for CountedTokioRwLockReadGuard<'a, T> {
828    type Target = TokioRwLockReadGuard<'a, T>;
829
830    fn deref(&self) -> &Self::Target {
831        &self.lock
832    }
833}
834
835impl<'a, T> DerefMut for CountedTokioRwLockReadGuard<'a, T> {
836    fn deref_mut(&mut self) -> &mut Self::Target {
837        &mut self.lock
838    }
839}
840
841type LockId = usize;
842
843#[derive(Debug, Clone)]
844pub struct LockMeta {
845    pub label: &'static str,
846    pub extra: Option<CompactString>,
847    pub kind: LockKind,
848    pub state: LockState,
849    pub started_at: Instant,
850}
851
852#[derive(Default, Clone)]
853pub struct LockRegistry {
854    id_gen: Arc<AtomicUsize>,
855    pub map: Arc<RwLock<IndexMap<LockId, LockMeta>>>,
856}
857
858impl LockRegistry {
859    fn remove(&self, id: &LockId) {
860        self.map.write().swap_remove(id);
861    }
862
863    async fn acquire_write<'a, T, C: fmt::Display, E: Into<Option<C>>>(
864        &self,
865        label: &'static str,
866        extra: E,
867        lock: &'a TokioRwLock<T>,
868    ) -> CountedTokioRwLockWriteGuard<'a, T> {
869        let id = self.gen_id();
870        self.insert_lock(
871            id,
872            LockMeta {
873                label,
874                extra: extra.into().map(|d| d.to_compact_string()),
875                kind: LockKind::Write,
876                state: LockState::Acquiring,
877                started_at: Instant::now(),
878            },
879        );
880        let _tracker = LockTracker {
881            id,
882            registry: self.clone(),
883        };
884        let w = lock.write().await;
885        self.set_lock_state(&id, LockState::Locked);
886        CountedTokioRwLockWriteGuard { lock: w, _tracker }
887    }
888
889    async fn acquire_write_owned<T, C: fmt::Display, E: Into<Option<C>>>(
890        &self,
891        label: &'static str,
892        extra: E,
893        lock: Arc<TokioRwLock<T>>,
894    ) -> CountedOwnedTokioRwLockWriteGuard<T> {
895        let id = self.gen_id();
896        self.insert_lock(
897            id,
898            LockMeta {
899                label,
900                extra: extra.into().map(|d| d.to_compact_string()),
901                kind: LockKind::Write,
902                state: LockState::Acquiring,
903                started_at: Instant::now(),
904            },
905        );
906        let _tracker = LockTracker {
907            id,
908            registry: self.clone(),
909        };
910        let w = lock.write_owned().await;
911        self.set_lock_state(&id, LockState::Locked);
912        CountedOwnedTokioRwLockWriteGuard { lock: w, _tracker }
913    }
914
915    fn acquire_blocking_write<'a, T, C: fmt::Display, E: Into<Option<C>>>(
916        &self,
917        label: &'static str,
918        extra: E,
919        lock: &'a TokioRwLock<T>,
920    ) -> CountedTokioRwLockWriteGuard<'a, T> {
921        let id = self.gen_id();
922        self.insert_lock(
923            id,
924            LockMeta {
925                label,
926                extra: extra.into().map(|d| d.to_compact_string()),
927                kind: LockKind::Write,
928                state: LockState::Acquiring,
929                started_at: Instant::now(),
930            },
931        );
932        let _tracker = LockTracker {
933            id,
934            registry: self.clone(),
935        };
936        let w = lock.blocking_write();
937        self.set_lock_state(&id, LockState::Locked);
938        CountedTokioRwLockWriteGuard { lock: w, _tracker }
939    }
940
941    fn acquire_blocking_write_owned<T, C: fmt::Display, E: Into<Option<C>>>(
942        &self,
943        label: &'static str,
944        extra: E,
945        lock: Arc<TokioRwLock<T>>,
946    ) -> CountedOwnedTokioRwLockWriteGuard<T> {
947        let id = self.gen_id();
948        self.insert_lock(
949            id,
950            LockMeta {
951                label,
952                extra: extra.into().map(|d| d.to_compact_string()),
953                kind: LockKind::Write,
954                state: LockState::Acquiring,
955                started_at: Instant::now(),
956            },
957        );
958        let _tracker = LockTracker {
959            id,
960            registry: self.clone(),
961        };
962        let w = loop {
963            if let Ok(w) = lock.clone().try_write_owned() {
964                break w;
965            }
966            // don't instantly loop
967            std::thread::sleep(Duration::from_millis(1));
968        };
969        self.set_lock_state(&id, LockState::Locked);
970        CountedOwnedTokioRwLockWriteGuard { lock: w, _tracker }
971    }
972
973    async fn acquire_read<'a, T, C: fmt::Display, E: Into<Option<C>>>(
974        &self,
975        label: &'static str,
976        extra: E,
977        lock: &'a TokioRwLock<T>,
978    ) -> CountedTokioRwLockReadGuard<'a, T> {
979        let id = self.gen_id();
980        self.insert_lock(
981            id,
982            LockMeta {
983                label,
984                extra: extra
985                    .into()
986                    .map(|d| d.to_compact_string())
987                    .map(|d| d.to_compact_string()),
988                kind: LockKind::Read,
989                state: LockState::Acquiring,
990                started_at: Instant::now(),
991            },
992        );
993        let _tracker = LockTracker {
994            id,
995            registry: self.clone(),
996        };
997        let w = lock.read().await;
998        self.set_lock_state(&id, LockState::Locked);
999        CountedTokioRwLockReadGuard { lock: w, _tracker }
1000    }
1001
1002    fn acquire_blocking_read<'a, T, C: fmt::Display, E: Into<Option<C>>>(
1003        &self,
1004        label: &'static str,
1005        extra: E,
1006        lock: &'a TokioRwLock<T>,
1007    ) -> CountedTokioRwLockReadGuard<'a, T> {
1008        let id = self.gen_id();
1009        self.insert_lock(
1010            id,
1011            LockMeta {
1012                label,
1013                extra: extra.into().map(|d| d.to_compact_string()),
1014                kind: LockKind::Read,
1015                state: LockState::Acquiring,
1016                started_at: Instant::now(),
1017            },
1018        );
1019        let _tracker = LockTracker {
1020            id,
1021            registry: self.clone(),
1022        };
1023        let w = lock.blocking_read();
1024        self.set_lock_state(&id, LockState::Locked);
1025        CountedTokioRwLockReadGuard { lock: w, _tracker }
1026    }
1027
1028    fn set_lock_state(&self, id: &LockId, state: LockState) {
1029        if let Some(meta) = self.map.write().get_mut(id) {
1030            meta.state = state
1031        }
1032    }
1033
1034    fn insert_lock(&self, id: LockId, meta: LockMeta) {
1035        self.map.write().insert(id, meta);
1036    }
1037
1038    fn gen_id(&self) -> LockId {
1039        self.id_gen.fetch_add(1, Ordering::Release) + 1
1040    }
1041}
1042
1043#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
1044#[serde(rename_all = "snake_case")]
1045pub enum LockState {
1046    Acquiring,
1047    Locked,
1048}
1049
1050#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
1051#[serde(rename_all = "snake_case")]
1052pub enum LockKind {
1053    Read,
1054    Write,
1055}
1056
1057struct LockTracker {
1058    id: LockId,
1059    registry: LockRegistry,
1060}
1061
1062impl Drop for LockTracker {
1063    fn drop(&mut self) {
1064        self.registry.remove(&self.id)
1065    }
1066}
1067
1068#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
1069pub struct PartialVersion {
1070    // range of sequences recorded
1071    pub seqs: RangeInclusiveSet<CrsqlSeq>,
1072    // actual last sequence originally produced
1073    pub last_seq: CrsqlSeq,
1074    // timestamp when the change was produced by the source
1075    pub ts: Timestamp,
1076}
1077
1078impl PartialVersion {
1079    pub fn is_complete(&self) -> bool {
1080        self.seqs.gaps(&self.full_range()).count() == 0
1081    }
1082
1083    pub fn full_range(&self) -> RangeInclusive<CrsqlSeq> {
1084        CrsqlSeq(1)..=self.last_seq
1085    }
1086}
1087
1088#[derive(Debug, Clone, Eq, PartialEq)]
1089pub struct CurrentVersion {
1090    pub db_version: CrsqlDbVersion,
1091    pub last_seq: CrsqlSeq,
1092    pub ts: Timestamp,
1093}
1094
1095#[derive(Debug, Clone, Eq, PartialEq)]
1096pub enum KnownDbVersion {
1097    Cleared,
1098    Current(CurrentVersion),
1099    Partial(PartialVersion),
1100}
1101
1102#[derive(Debug)]
1103pub struct VersionsSnapshot {
1104    actor_id: ActorId,
1105    needed: RangeInclusiveSet<CrsqlDbVersion>,
1106    partials: BTreeMap<CrsqlDbVersion, PartialVersion>,
1107    max: Option<CrsqlDbVersion>,
1108}
1109
1110impl VersionsSnapshot {
1111    pub fn needed(&self) -> &RangeInclusiveSet<CrsqlDbVersion> {
1112        &self.needed
1113    }
1114
1115    pub fn insert_gaps(&mut self, db_versions: RangeInclusiveSet<CrsqlDbVersion>) {
1116        self.needed.extend(db_versions);
1117    }
1118
1119    pub fn insert_db(
1120        &mut self,         // only because we want 1 mt a time here
1121        conn: &Connection, // usually a `Transaction`
1122        db_versions: RangeInclusiveSet<CrsqlDbVersion>,
1123    ) -> rusqlite::Result<()> {
1124        trace!("wants to insert into db {db_versions:?}");
1125        let mut changes = self.compute_gaps_change(db_versions);
1126
1127        trace!(actor_id = %self.actor_id, "delete: {:?}", changes.remove_ranges);
1128        trace!(actor_id = %self.actor_id, "new: {:?}", changes.insert_set);
1129
1130        // those are actual ranges we had stored and will change, remove them from the DB
1131        for range in std::mem::take(&mut changes.remove_ranges) {
1132            debug!(actor_id = %self.actor_id, "deleting {range:?}");
1133            let count = conn
1134                .prepare_cached("DELETE FROM __corro_bookkeeping_gaps WHERE actor_id = :actor_id AND start = :start AND end = :end")?
1135                .execute(named_params! {
1136                    ":actor_id": self.actor_id,
1137                    ":start": range.start(),
1138                    ":end": range.end()
1139                })?;
1140            if count != 1 {
1141                warn!(actor_id = %self.actor_id, "did not delete gap from db: {range:?}");
1142            }
1143            let details = json!({"count": count, "range": range});
1144            assert_always!(count == 1, "ineffective deletion of gaps in-db", &details);
1145            for version in CrsqlDbVersionRange::from(&range) {
1146                self.partials.remove(&version);
1147            }
1148            self.needed.remove(range);
1149        }
1150
1151        for range in std::mem::take(&mut changes.insert_set) {
1152            debug!(actor_id = %self.actor_id, "inserting {range:?}");
1153            let res = conn
1154                .prepare_cached(
1155                    "INSERT INTO __corro_bookkeeping_gaps VALUES (:actor_id, :start, :end)",
1156                )?
1157                .execute(named_params! {
1158                    ":actor_id": self.actor_id,
1159                    ":start": range.start(),
1160                    ":end": range.end()
1161                });
1162
1163            if let Err(e) = res {
1164                let (actor_id, start, end) : (ActorId, CrsqlDbVersion, CrsqlDbVersion) = conn.query_row("SELECT actor_id, start, end FROM __corro_bookkeeping_gaps WHERE actor_id = :actor_id AND start = :start", named_params! {
1165                    ":actor_id": self.actor_id,
1166                    ":start": range.start(),
1167                }, |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))?;
1168
1169                warn!("already had gaps entry! actor_id: {actor_id}, start: {start}, end: {end}");
1170
1171                return Err(e);
1172            }
1173            self.needed.insert(range);
1174        }
1175
1176        self.max = changes.max.take();
1177
1178        Ok(())
1179    }
1180
1181    fn compute_gaps_change(&self, versions: RangeInclusiveSet<CrsqlDbVersion>) -> GapsChanges {
1182        trace!("needed: {:?}", self.needed);
1183
1184        let mut changes = GapsChanges {
1185            // set as the current max
1186            max: self.max,
1187
1188            insert_set: Default::default(),
1189            remove_ranges: Default::default(),
1190        };
1191
1192        for versions in versions.clone() {
1193            // only update the max if it's bigger
1194            changes.max = cmp::max(changes.max, Some(*versions.end()));
1195
1196            // iterate all partially or fully overlapping changes
1197            for range in self.needed.overlapping(&versions) {
1198                trace!(actor_id = %self.actor_id, "overlapping: {range:?}");
1199                // insert the overlapping range in the set (collapses ajoining ranges)
1200                changes.insert_set.insert(range.clone());
1201                // remove the range, they'll be added back from the insert set ^
1202                changes.remove_ranges.insert(range.clone());
1203            }
1204
1205            // check if there's a previous range with an end version = start version - 1
1206            if let Some(range) = self.needed.get(&CrsqlDbVersion(versions.start().0 - 1)) {
1207                trace!(actor_id = %self.actor_id, "got a start - 1: {range:?}");
1208                // insert the collapsible range
1209                changes.insert_set.insert(range.clone());
1210                // remove the collapsible range
1211                changes.remove_ranges.insert(range.clone());
1212            }
1213
1214            // check if there's a next range with an start version = end version + 1
1215            if let Some(range) = self.needed.get(&CrsqlDbVersion(versions.end().0 + 1)) {
1216                trace!(actor_id = %self.actor_id, "got a end + 1: {range:?}");
1217                // insert the collapsible range
1218                changes.insert_set.insert(range.clone());
1219                // remove the collapsible range
1220                changes.remove_ranges.insert(range.clone());
1221            }
1222
1223            // either a max or 0
1224            // TODO: figure out if we want to use 0 instead of None in the struct by default
1225            let current_max = self.max.unwrap_or_default();
1226
1227            // check if there's a gap created between our current max and the start version we just inserted
1228            let gap_start = current_max + 1;
1229            if gap_start < *versions.start() {
1230                let range = gap_start..=*versions.start();
1231                trace!("inserting gap between max + 1 and start: {range:?}");
1232                changes.insert_set.insert(range.clone());
1233                for range in self.needed.overlapping(&range) {
1234                    changes.insert_set.insert(range.clone());
1235                    changes.remove_ranges.insert(range.clone());
1236                }
1237            }
1238        }
1239
1240        for versions in versions {
1241            // we now know the applied versions
1242            changes.insert_set.remove(versions.clone());
1243        }
1244
1245        changes
1246    }
1247}
1248
1249// this struct must be drained!
1250impl Drop for VersionsSnapshot {
1251    fn drop(&mut self) {
1252        trace!("dropping snapshot: {self:?}");
1253        if !self.needed.is_empty() {
1254            warn!(actor_id = %self.actor_id, "needed versions from snapshot were not drained: {:?}", self.needed);
1255        }
1256        debug_assert!(
1257            self.needed.is_empty(),
1258            "needed versions from snapshot were not drained"
1259        );
1260        if !self.partials.is_empty() {
1261            warn!(actor_id = %self.actor_id, "partials were not drained: {:?}", self.partials);
1262        }
1263        debug_assert!(
1264            self.partials.is_empty(),
1265            "partials from snapshot were not drained!"
1266        );
1267        debug_assert!(self.max.is_none(), "max value was not applied");
1268    }
1269}
1270
1271#[derive(Clone, Debug, Eq, PartialEq)]
1272pub struct BookedVersions {
1273    actor_id: ActorId,
1274    pub partials: BTreeMap<CrsqlDbVersion, PartialVersion>,
1275    needed: RangeInclusiveSet<CrsqlDbVersion>,
1276    max: Option<CrsqlDbVersion>,
1277}
1278
1279impl BookedVersions {
1280    pub fn new(actor_id: ActorId) -> Self {
1281        Self {
1282            actor_id,
1283            partials: Default::default(),
1284            needed: Default::default(),
1285            max: Default::default(),
1286        }
1287    }
1288
1289    pub fn actor_id(&self) -> ActorId {
1290        self.actor_id
1291    }
1292
1293    pub fn from_conn(conn: &Connection, actor_id: ActorId) -> rusqlite::Result<Self> {
1294        trace!("from_conn");
1295        let mut bv = BookedVersions::new(actor_id);
1296
1297        // fetch the biggest version we know, a partial version might override
1298        // this below
1299        bv.max = conn
1300            .prepare_cached("SELECT db_version FROM crsql_db_versions WHERE site_id = ?")?
1301            .query_row([actor_id], |row| row.get(0))
1302            .optional()?;
1303
1304        {
1305            // fetch known partial sequences
1306            let mut prepped = conn.prepare_cached(
1307            "SELECT db_version, start_seq, end_seq, last_seq, ts FROM __corro_seq_bookkeeping WHERE site_id = ?",
1308            )?;
1309            let mut rows = prepped.query([actor_id])?;
1310
1311            loop {
1312                let row = rows.next()?;
1313                match row {
1314                    None => break,
1315                    Some(row) => {
1316                        let version = row.get(0)?;
1317                        // NOTE: use normal insert logic to have a consistent behavior
1318                        bv.insert_partial(
1319                            version,
1320                            PartialVersion {
1321                                seqs: RangeInclusiveSet::from_iter(vec![row.get(1)?..=row.get(2)?]),
1322                                last_seq: row.get(3)?,
1323                                ts: row.get(4)?,
1324                            },
1325                        );
1326                    }
1327                }
1328            }
1329        }
1330
1331        let mut snap = bv.snapshot();
1332
1333        {
1334            // fetch the sync's needed version gaps
1335            let mut prepped = conn.prepare_cached(
1336                "SELECT start, end FROM __corro_bookkeeping_gaps WHERE actor_id = ?",
1337            )?;
1338            let mut rows = prepped.query([actor_id])?;
1339
1340            loop {
1341                let row = rows.next()?;
1342                match row {
1343                    None => break,
1344                    Some(row) => {
1345                        let start_v = row.get(0)?;
1346                        let end_v = row.get(1)?;
1347
1348                        // TODO: don't do this manually...
1349                        snap.needed.insert(start_v..=end_v);
1350                        // max for booked versions shouldn't come from gaps
1351                        if Some(end_v) > snap.max {
1352                            warn!(%actor_id, %start_v, %end_v, max = ?snap.max, "max for actor is less than gap");
1353                        }
1354                    }
1355                }
1356            }
1357        }
1358
1359        bv.commit_snapshot(snap);
1360
1361        Ok(bv)
1362    }
1363
1364    pub fn contains_version(&self, version: &CrsqlDbVersion) -> bool {
1365        // corrosion knows about a version if...
1366
1367        // it's not in the list of needed versions
1368        !self.needed.iter().any(|range| range.start() <= version && version <= range.end()) &&
1369        // and the last known version is bigger than the requested version
1370        self.max.unwrap_or_default() >= *version
1371        // we don't need to look at partials because if we have a partial
1372        // then it fulfills the previous conditions
1373    }
1374
1375    pub fn get_partial(&self, version: &CrsqlDbVersion) -> Option<&PartialVersion> {
1376        self.partials.get(version)
1377    }
1378
1379    pub fn contains(&self, version: CrsqlDbVersion, seqs: Option<CrsqlSeqRange>) -> bool {
1380        self.contains_version(&version)
1381            && seqs
1382                .map(|mut check_seqs| match self.partials.get(&version) {
1383                    Some(partial) => check_seqs.all(|seq| partial.seqs.contains(&seq)),
1384                    // if `contains_version` is true but we don't have a partial version,
1385                    // then we must have it as a fully applied or cleared version
1386                    None => true,
1387                })
1388                .unwrap_or(true)
1389    }
1390
1391    #[inline]
1392    pub fn contains_all(
1393        &self,
1394        versions: impl Into<CrsqlDbVersionRange>,
1395        seqs: Option<CrsqlSeqRange>,
1396    ) -> bool {
1397        versions.into().all(|version| self.contains(version, seqs))
1398    }
1399
1400    pub fn last(&self) -> Option<CrsqlDbVersion> {
1401        self.max
1402    }
1403
1404    pub fn commit_snapshot(&mut self, mut snap: VersionsSnapshot) {
1405        debug!("comitting snapshot");
1406        self.needed = std::mem::take(&mut snap.needed);
1407        self.partials = std::mem::take(&mut snap.partials);
1408        self.max = snap.max.take();
1409    }
1410
1411    pub fn snapshot(&self) -> VersionsSnapshot {
1412        trace!("creating snapshot");
1413        VersionsSnapshot {
1414            actor_id: self.actor_id,
1415            needed: self.needed.clone(),
1416            partials: self.partials.clone(),
1417            max: self.max,
1418        }
1419    }
1420
1421    // used when the commit has succeeded
1422    pub fn insert_partial(
1423        &mut self,
1424        version: CrsqlDbVersion,
1425        partial: PartialVersion,
1426    ) -> PartialVersion {
1427        debug!(actor_id = %self.actor_id, "insert partial {version:?}");
1428
1429        match self.partials.entry(version) {
1430            btree_map::Entry::Vacant(entry) => {
1431                self.max = cmp::max(self.max, Some(version));
1432                entry.insert(partial).clone()
1433            }
1434            btree_map::Entry::Occupied(mut entry) => {
1435                let got = entry.get_mut();
1436                got.seqs.extend(partial.seqs);
1437                got.clone()
1438            }
1439        }
1440    }
1441
1442    pub fn needed(&self) -> &RangeInclusiveSet<CrsqlDbVersion> {
1443        &self.needed
1444    }
1445}
1446
1447#[derive(Debug)]
1448pub struct GapsChanges {
1449    max: Option<CrsqlDbVersion>,
1450    insert_set: RangeInclusiveSet<CrsqlDbVersion>,
1451    remove_ranges: HashSet<RangeInclusive<CrsqlDbVersion>>,
1452}
1453
1454pub type BookedInner = Arc<CountedTokioRwLock<BookedVersions>>;
1455
1456#[derive(Clone)]
1457pub struct Booked(BookedInner);
1458
1459impl Booked {
1460    pub fn new(versions: BookedVersions, registry: LockRegistry) -> Self {
1461        Self(Arc::new(CountedTokioRwLock::new(registry, versions)))
1462    }
1463
1464    pub async fn read<D: fmt::Display, E: Into<Option<D>>>(
1465        &self,
1466        label: &'static str,
1467        extra: E,
1468    ) -> CountedTokioRwLockReadGuard<'_, BookedVersions> {
1469        self.0.read(label, extra).await
1470    }
1471
1472    pub async fn write<C: fmt::Display, E: Into<Option<C>>>(
1473        &self,
1474        label: &'static str,
1475        extra: E,
1476    ) -> CountedTokioRwLockWriteGuard<'_, BookedVersions> {
1477        self.0.write(label, extra).await
1478    }
1479
1480    pub async fn write_owned<C: fmt::Display, E: Into<Option<C>>>(
1481        &self,
1482        label: &'static str,
1483        extra: E,
1484    ) -> CountedOwnedTokioRwLockWriteGuard<BookedVersions> {
1485        self.0.write_owned(label, extra).await
1486    }
1487
1488    pub fn blocking_write<C: fmt::Display, E: Into<Option<C>>>(
1489        &self,
1490        label: &'static str,
1491        extra: E,
1492    ) -> CountedTokioRwLockWriteGuard<'_, BookedVersions> {
1493        self.0.blocking_write(label, extra)
1494    }
1495
1496    pub fn blocking_read<C: fmt::Display, E: Into<Option<C>>>(
1497        &self,
1498        label: &'static str,
1499        extra: E,
1500    ) -> CountedTokioRwLockReadGuard<'_, BookedVersions> {
1501        self.0.blocking_read(label, extra)
1502    }
1503
1504    pub fn blocking_write_owned<C: fmt::Display, E: Into<Option<C>>>(
1505        &self,
1506        label: &'static str,
1507        extra: E,
1508    ) -> CountedOwnedTokioRwLockWriteGuard<BookedVersions> {
1509        self.0.blocking_write_owned(label, extra)
1510    }
1511}
1512
1513#[derive(Default)]
1514pub struct BookieInner {
1515    map: HashMap<ActorId, Booked>,
1516    registry: LockRegistry,
1517}
1518
1519impl BookieInner {
1520    pub fn ensure(&mut self, actor_id: ActorId) -> Booked {
1521        self.map
1522            .entry(actor_id)
1523            .or_insert_with(|| {
1524                Booked(Arc::new(CountedTokioRwLock::new(
1525                    self.registry.clone(),
1526                    BookedVersions::new(actor_id),
1527                )))
1528            })
1529            .clone()
1530    }
1531
1532    pub fn replace_actor(&mut self, actor_id: ActorId, bv: BookedVersions) {
1533        self.map.insert(
1534            actor_id,
1535            Booked(Arc::new(CountedTokioRwLock::new(self.registry.clone(), bv))),
1536        );
1537    }
1538}
1539
1540impl Deref for BookieInner {
1541    type Target = HashMap<ActorId, Booked>;
1542
1543    fn deref(&self) -> &Self::Target {
1544        &self.map
1545    }
1546}
1547
1548impl DerefMut for BookieInner {
1549    fn deref_mut(&mut self) -> &mut Self::Target {
1550        &mut self.map
1551    }
1552}
1553
1554#[derive(Clone)]
1555pub struct Bookie(Arc<CountedTokioRwLock<BookieInner>>);
1556
1557impl Bookie {
1558    pub fn new(map: HashMap<ActorId, BookedVersions>) -> Self {
1559        let registry = LockRegistry::default();
1560        Self::new_with_registry(map, registry)
1561    }
1562
1563    pub fn new_with_registry(
1564        map: HashMap<ActorId, BookedVersions>,
1565        registry: LockRegistry,
1566    ) -> Self {
1567        Self(Arc::new(CountedTokioRwLock::new(
1568            registry.clone(),
1569            BookieInner {
1570                map: map
1571                    .into_iter()
1572                    .map(|(k, v)| (k, Booked::new(v, registry.clone())))
1573                    .collect(),
1574                registry,
1575            },
1576        )))
1577    }
1578
1579    pub async fn read<C: fmt::Display, E: Into<Option<C>>>(
1580        &self,
1581        label: &'static str,
1582        extra: E,
1583    ) -> CountedTokioRwLockReadGuard<'_, BookieInner> {
1584        self.0.read(label, extra).await
1585    }
1586
1587    pub async fn write<C: fmt::Display, E: Into<Option<C>>>(
1588        &self,
1589        label: &'static str,
1590        extra: E,
1591    ) -> CountedTokioRwLockWriteGuard<'_, BookieInner> {
1592        self.0.write(label, extra).await
1593    }
1594
1595    pub fn blocking_write<C: fmt::Display, E: Into<Option<C>>>(
1596        &self,
1597        label: &'static str,
1598        extra: E,
1599    ) -> CountedTokioRwLockWriteGuard<'_, BookieInner> {
1600        self.0.blocking_write(label, extra)
1601    }
1602
1603    pub fn registry(&self) -> &LockRegistry {
1604        self.0.registry()
1605    }
1606}
1607
1608#[cfg(test)]
1609mod tests {
1610    use super::*;
1611    use crate::{dbvr, dbvri};
1612    use rangemap::range_inclusive_set;
1613
1614    #[test]
1615    fn test_booked_insert_db() -> rusqlite::Result<()> {
1616        _ = tracing_subscriber::fmt::try_init();
1617
1618        let mut conn = CrConn::init(Connection::open_in_memory()?)?;
1619        setup_conn(&conn)?;
1620        let clock = Arc::new(uhlc::HLC::default());
1621        migrate(clock, &mut conn)?;
1622
1623        let actor_id = ActorId::default();
1624        let mut bv = BookedVersions::new(actor_id);
1625
1626        let mut all = RangeInclusiveSet::new();
1627
1628        insert_everywhere(
1629            &conn,
1630            &mut bv,
1631            &mut all,
1632            range_inclusive_set![dbvri!(1, 20)],
1633        )?;
1634        expect_gaps(&conn, &bv, &all, vec![])?;
1635
1636        insert_everywhere(
1637            &conn,
1638            &mut bv,
1639            &mut all,
1640            range_inclusive_set![dbvri!(1, 10)],
1641        )?;
1642        expect_gaps(&conn, &bv, &all, vec![])?;
1643
1644        // try from an empty state again
1645        let mut bv = BookedVersions::new(actor_id);
1646        let mut all = RangeInclusiveSet::new();
1647
1648        // create 2:=3 gap
1649        insert_everywhere(
1650            &conn,
1651            &mut bv,
1652            &mut all,
1653            range_inclusive_set![dbvri!(1, 1), dbvri!(4, 4)],
1654        )?;
1655        expect_gaps(&conn, &bv, &all, vec![dbvr!(2, 3)])?;
1656
1657        // fill gap
1658        insert_everywhere(
1659            &conn,
1660            &mut bv,
1661            &mut all,
1662            range_inclusive_set![dbvri!(3, 3), dbvri!(2, 2)],
1663        )?;
1664        expect_gaps(&conn, &bv, &all, vec![])?;
1665
1666        // try from an empty state again
1667        let mut bv = BookedVersions::new(actor_id);
1668        let mut all = RangeInclusiveSet::new();
1669
1670        // insert a non-1 first version
1671        insert_everywhere(
1672            &conn,
1673            &mut bv,
1674            &mut all,
1675            range_inclusive_set![dbvri!(5, 20)],
1676        )?;
1677        expect_gaps(&conn, &bv, &all, vec![dbvr!(1, 4)])?;
1678
1679        // insert a further change that does not overlap a gap
1680        insert_everywhere(&conn, &mut bv, &mut all, range_inclusive_set![dbvri!(6, 7)])?;
1681        expect_gaps(&conn, &bv, &all, vec![dbvr!(1, 4)])?;
1682
1683        // insert a further change that does overlap a gap
1684        insert_everywhere(&conn, &mut bv, &mut all, range_inclusive_set![dbvri!(3, 7)])?;
1685        expect_gaps(&conn, &bv, &all, vec![dbvr!(1, 2)])?;
1686
1687        insert_everywhere(&conn, &mut bv, &mut all, range_inclusive_set![dbvri!(1, 2)])?;
1688        expect_gaps(&conn, &bv, &all, vec![])?;
1689
1690        insert_everywhere(
1691            &conn,
1692            &mut bv,
1693            &mut all,
1694            range_inclusive_set![dbvri!(25, 25)],
1695        )?;
1696        expect_gaps(&conn, &bv, &all, vec![dbvr!(21, 24)])?;
1697
1698        insert_everywhere(
1699            &conn,
1700            &mut bv,
1701            &mut all,
1702            range_inclusive_set![dbvri!(30, 35)],
1703        )?;
1704        expect_gaps(&conn, &bv, &all, vec![dbvr!(21, 24), dbvr!(26, 29)])?;
1705
1706        // NOTE: overlapping partially from the end
1707
1708        insert_everywhere(
1709            &conn,
1710            &mut bv,
1711            &mut all,
1712            range_inclusive_set![dbvri!(19, 22)],
1713        )?;
1714        expect_gaps(&conn, &bv, &all, vec![dbvr!(23, 24), dbvr!(26, 29)])?;
1715
1716        // NOTE: overlapping partially from the start
1717
1718        insert_everywhere(
1719            &conn,
1720            &mut bv,
1721            &mut all,
1722            range_inclusive_set![dbvri!(24, 25)],
1723        )?;
1724        expect_gaps(&conn, &bv, &all, vec![dbvr!(23, 23), dbvr!(26, 29)])?;
1725
1726        // NOTE: overlapping 2 ranges
1727
1728        insert_everywhere(
1729            &conn,
1730            &mut bv,
1731            &mut all,
1732            range_inclusive_set![dbvri!(23, 27)],
1733        )?;
1734        expect_gaps(&conn, &bv, &all, vec![dbvr!(28, 29)])?;
1735
1736        // NOTE: ineffective insert of already known ranges
1737
1738        insert_everywhere(
1739            &conn,
1740            &mut bv,
1741            &mut all,
1742            range_inclusive_set![dbvri!(1, 20)],
1743        )?;
1744        expect_gaps(&conn, &bv, &all, vec![dbvr!(28, 29)])?;
1745
1746        // NOTE: overlapping no ranges, but encompassing a full range
1747
1748        insert_everywhere(
1749            &conn,
1750            &mut bv,
1751            &mut all,
1752            range_inclusive_set![dbvri!(27, 30)],
1753        )?;
1754        expect_gaps(&conn, &bv, &all, vec![])?;
1755
1756        // NOTE: touching multiple ranges, partially
1757
1758        // create gap 36..=39
1759        insert_everywhere(
1760            &conn,
1761            &mut bv,
1762            &mut all,
1763            range_inclusive_set![dbvri!(40, 45)],
1764        )?;
1765        // create gap 46..=49
1766        insert_everywhere(
1767            &conn,
1768            &mut bv,
1769            &mut all,
1770            range_inclusive_set![dbvri!(50, 55)],
1771        )?;
1772
1773        insert_everywhere(
1774            &conn,
1775            &mut bv,
1776            &mut all,
1777            range_inclusive_set![dbvri!(38, 47)],
1778        )?;
1779        expect_gaps(&conn, &bv, &all, vec![dbvr!(36, 37), dbvr!(48, 49)])?;
1780
1781        // test loading a bv from the conn, they should be identical!
1782        let mut bv2 = BookedVersions::from_conn(&conn, actor_id)?;
1783        // manually set the last version because there's nothing in `__corro_bookkeeping`
1784        bv2.max = Some(CrsqlDbVersion(55));
1785
1786        assert_eq!(bv, bv2);
1787
1788        Ok(())
1789    }
1790
1791    fn insert_everywhere(
1792        conn: &Connection,
1793        bv: &mut BookedVersions,
1794        all_versions: &mut RangeInclusiveSet<CrsqlDbVersion>,
1795        versions: RangeInclusiveSet<CrsqlDbVersion>,
1796    ) -> rusqlite::Result<()> {
1797        all_versions.extend(versions.clone());
1798        let mut snap = bv.snapshot();
1799        snap.insert_db(conn, versions)?;
1800        bv.commit_snapshot(snap);
1801        Ok(())
1802    }
1803
1804    fn expect_gaps(
1805        conn: &Connection,
1806        bv: &BookedVersions,
1807        all_versions: &RangeInclusiveSet<CrsqlDbVersion>,
1808        expected: Vec<CrsqlDbVersionRange>,
1809    ) -> rusqlite::Result<()> {
1810        let gaps: Vec<(ActorId, CrsqlDbVersion, CrsqlDbVersion)> = conn
1811            .prepare_cached("SELECT actor_id, start, end FROM __corro_bookkeeping_gaps")?
1812            .query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))?
1813            .collect::<Result<Vec<_>, _>>()?;
1814
1815        assert_eq!(
1816            gaps,
1817            expected
1818                .clone()
1819                .into_iter()
1820                .map(|expected| (bv.actor_id, expected.start(), expected.end()))
1821                .collect::<Vec<_>>()
1822        );
1823
1824        for range in all_versions.iter() {
1825            assert!(bv.contains_all(range.clone(), None));
1826        }
1827
1828        for range in expected {
1829            for v in range {
1830                assert!(!bv.contains(v, None), "expected not to contain {v}");
1831                assert!(bv.needed.contains(&v), "expected needed to contain {v:?}");
1832            }
1833        }
1834
1835        assert_eq!(
1836            bv.max,
1837            all_versions.iter().next_back().map(|range| *range.end()),
1838            "expected last version not to increment"
1839        );
1840
1841        Ok(())
1842    }
1843}