Skip to main content

kache_store/
index_compaction.rs

1//! Reclaiming free SQLite pages from the index file.
2
3use crate::file_hash::FileHashCache;
4use rusqlite::{Connection, ErrorCode};
5
6/// Free pages below this many bytes are left for SQLite to reuse.
7const MIN_RECLAIM_BYTES: u64 = 64 << 20;
8/// Room kept on top of the rewritten pages for journal and filesystem slack.
9const VACUUM_HEADROOM_BYTES: u64 = 64 << 20;
10
11/// Page accounting of the index file, read without modifying it.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct IndexPageStats {
14    pub pages: u64,
15    pub free_pages: u64,
16    pub page_size: u64,
17}
18
19impl IndexPageStats {
20    pub fn file_bytes(&self) -> u64 {
21        self.pages.saturating_mul(self.page_size)
22    }
23
24    pub fn free_bytes(&self) -> u64 {
25        self.free_pages.saturating_mul(self.page_size)
26    }
27
28    pub fn live_bytes(&self) -> u64 {
29        self.pages
30            .saturating_sub(self.free_pages)
31            .saturating_mul(self.page_size)
32    }
33
34    /// Worth a VACUUM only when the freelist is both large and at least a
35    /// quarter of the file.
36    pub fn should_compact(&self) -> bool {
37        self.free_bytes() >= MIN_RECLAIM_BYTES && self.free_pages >= self.pages.div_ceil(4)
38    }
39
40    /// Free disk bytes a VACUUM needs. VACUUM copies only live pages: once
41    /// into a temporary database and once into the WAL. The size of the file
42    /// itself does not matter, so a 45 GB index with 1 GB live needs about
43    /// 2 GB, not 90 GB.
44    pub fn vacuum_space_needed(&self) -> u64 {
45        self.live_bytes()
46            .saturating_mul(2)
47            .saturating_add(VACUUM_HEADROOM_BYTES)
48    }
49
50    pub fn fits_in(&self, free_disk_bytes: u64) -> bool {
51        free_disk_bytes >= self.vacuum_space_needed()
52    }
53}
54
55/// Result of one compaction attempt. Only I/O failures are errors.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum IndexCompaction {
58    /// The freelist is below the thresholds.
59    NotNeeded,
60    /// Another connection holds the index. Nothing was lost; retry later.
61    Busy,
62    /// The volume cannot hold the rewritten copy. Free pages stay reusable.
63    InsufficientSpace { needed: u64, free: u64 },
64    /// File size in bytes before and after.
65    Compacted { before: u64, after: u64 },
66}
67
68/// Page accounting for any connection to the index, including read-only ones.
69pub fn index_page_stats(db: &Connection) -> rusqlite::Result<IndexPageStats> {
70    let pragma = |name: &str| -> rusqlite::Result<u64> {
71        db.pragma_query_value(None, name, |row| row.get::<_, u32>(0).map(u64::from))
72    };
73    Ok(IndexPageStats {
74        pages: pragma("page_count")?,
75        free_pages: pragma("freelist_count")?,
76        page_size: pragma("page_size")?,
77    })
78}
79
80fn is_busy(err: &rusqlite::Error) -> bool {
81    matches!(
82        err.sqlite_error_code(),
83        Some(ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
84    )
85}
86
87/// Where a compaction attempt stopped. The two busy stages differ only in
88/// the message `compact_sparse_index` shows.
89enum Attempt {
90    Done(IndexCompaction),
91    BusyBeforeVacuum,
92    BusyAfterVacuum,
93}
94
95const BUSY_BEFORE_VACUUM: &str = "index is busy; retry compaction after builds finish";
96const BUSY_AFTER_VACUUM: &str = "index compacted but WAL is busy; retry repair after builds finish";
97const NO_SPACE: &str = "not enough free space to compact the index; its free pages remain reusable";
98
99/// Turn index contention into `when_busy`; every other error stays an error.
100fn busy_as<T>(result: rusqlite::Result<T>, when_busy: T) -> rusqlite::Result<T> {
101    match result {
102        Err(err) if is_busy(&err) => Ok(when_busy),
103        other => other,
104    }
105}
106
107/// TRUNCATE checkpoint; `Ok(false)` when another connection blocked it.
108fn checkpoint(db: &Connection) -> rusqlite::Result<bool> {
109    let ran = db.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |row| {
110        row.get::<_, i64>(0).map(|busy| busy == 0)
111    });
112    busy_as(ran, false)
113}
114
115fn attempt(db: &Connection) -> anyhow::Result<Attempt> {
116    use anyhow::Context;
117    let stats = index_page_stats(db)?;
118    if !stats.should_compact() {
119        return Ok(Attempt::Done(IndexCompaction::NotNeeded));
120    }
121    let path = std::path::Path::new(db.path().context("index has no filesystem path")?);
122    let free = kache_fs::volume_usage(path)
123        .context("could not check free space for index compaction")?
124        .free;
125    if !stats.fits_in(free) {
126        return Ok(Attempt::Done(IndexCompaction::InsufficientSpace {
127            needed: stats.vacuum_space_needed(),
128            free,
129        }));
130    }
131    if !checkpoint(db)? {
132        return Ok(Attempt::BusyBeforeVacuum);
133    }
134    let before = std::fs::metadata(path)?.len();
135    if !busy_as(db.execute_batch("VACUUM").map(|()| true), false)? {
136        return Ok(Attempt::BusyBeforeVacuum);
137    }
138    // The file shrinks only once the vacuumed pages leave the WAL.
139    if !checkpoint(db)? {
140        return Ok(Attempt::BusyAfterVacuum);
141    }
142    let after = std::fs::metadata(path)?.len();
143    Ok(Attempt::Done(IndexCompaction::Compacted { before, after }))
144}
145
146impl FileHashCache<'_> {
147    pub fn index_page_stats(&self) -> rusqlite::Result<IndexPageStats> {
148        index_page_stats(self.db())
149    }
150
151    /// VACUUM the index when its freelist is large. Index reads continue
152    /// while it runs but writes wait, so callers pick a quiet moment; a
153    /// contended index reports `Busy` instead of failing.
154    pub fn compact_index(&self) -> anyhow::Result<IndexCompaction> {
155        Ok(match attempt(self.db())? {
156            Attempt::Done(outcome) => outcome,
157            Attempt::BusyBeforeVacuum | Attempt::BusyAfterVacuum => IndexCompaction::Busy,
158        })
159    }
160
161    /// `compact_index` for explicit repair: anything short of a compaction
162    /// that was needed is an error the user can read.
163    pub fn compact_sparse_index(&self) -> anyhow::Result<Option<(u64, u64)>> {
164        match attempt(self.db())? {
165            Attempt::Done(IndexCompaction::NotNeeded) => Ok(None),
166            Attempt::Done(IndexCompaction::Compacted { before, after }) => {
167                Ok(Some((before, after)))
168            }
169            Attempt::Done(IndexCompaction::InsufficientSpace { .. }) => anyhow::bail!(NO_SPACE),
170            Attempt::Done(IndexCompaction::Busy) | Attempt::BusyBeforeVacuum => {
171                anyhow::bail!(BUSY_BEFORE_VACUUM)
172            }
173            Attempt::BusyAfterVacuum => anyhow::bail!(BUSY_AFTER_VACUUM),
174        }
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    fn stats(pages: u64, free_pages: u64, page_size: u64) -> IndexPageStats {
183        IndexPageStats {
184            pages,
185            free_pages,
186            page_size,
187        }
188    }
189
190    /// A 45 GB index whose migration left 44.2 GB on the freelist.
191    const SPARSE_HOST: IndexPageStats = IndexPageStats {
192        pages: 11_019_675,
193        free_pages: 10_803_505,
194        page_size: 4096,
195    };
196
197    /// An index with 80 MiB of dropped legacy rows and one live table.
198    #[cfg(unix)]
199    fn sparse_index(path: &std::path::Path) -> FileHashCache<'static> {
200        let db = Connection::open(path).unwrap();
201        db.execute_batch(
202            "CREATE TABLE legacy(payload BLOB);
203             INSERT INTO legacy VALUES (zeroblob(83886080));
204             CREATE TABLE kept(value TEXT);
205             INSERT INTO kept VALUES ('live');",
206        )
207        .unwrap();
208        drop(db);
209        let cache = FileHashCache::open(path).unwrap();
210        cache.db().execute_batch("DROP TABLE legacy").unwrap();
211        cache
212    }
213
214    #[test]
215    fn index_compaction_requires_both_size_and_fraction() {
216        assert!(!stats(100_000, 25_000, 1).should_compact());
217        assert!(!stats(1_000_000, 20_000, 4096).should_compact());
218        assert!(stats(65_536, 16_384, 4096).should_compact());
219        assert!(!stats(65_537, 16_384, 4096).should_compact());
220        assert!(!stats(65_532, 16_383, 4096).should_compact());
221    }
222
223    #[test]
224    fn page_stats_split_the_file_into_free_and_live_bytes() {
225        let s = stats(10, 3, 4096);
226        assert_eq!(s.file_bytes(), 40_960);
227        assert_eq!(s.free_bytes(), 12_288);
228        assert_eq!(s.live_bytes(), 28_672);
229        assert_eq!(s.vacuum_space_needed(), 2 * 28_672 + (64 << 20));
230        assert_eq!(stats(3, 10, 4096).live_bytes(), 0);
231        assert_eq!(stats(u64::MAX, 0, 4096).vacuum_space_needed(), u64::MAX);
232    }
233
234    #[test]
235    fn space_guard_follows_live_bytes_not_file_size() {
236        assert_eq!(SPARSE_HOST.file_bytes(), 45_136_588_800);
237        assert_eq!(SPARSE_HOST.live_bytes(), 885_432_320);
238        assert!(SPARSE_HOST.should_compact());
239        let needed = SPARSE_HOST.vacuum_space_needed();
240        assert_eq!(needed, 1_837_973_504);
241        assert!(SPARSE_HOST.fits_in(3 << 30));
242        assert!(!SPARSE_HOST.fits_in(1 << 30));
243        assert!(SPARSE_HOST.fits_in(needed));
244        assert!(!SPARSE_HOST.fits_in(needed - 1));
245    }
246
247    #[test]
248    fn only_contention_counts_as_busy() {
249        let failure =
250            |code: i32| rusqlite::Error::SqliteFailure(rusqlite::ffi::Error::new(code), None);
251        assert!(is_busy(&failure(rusqlite::ffi::SQLITE_BUSY)));
252        assert!(is_busy(&failure(rusqlite::ffi::SQLITE_LOCKED)));
253        assert!(!is_busy(&failure(rusqlite::ffi::SQLITE_IOERR)));
254        assert!(!is_busy(&failure(rusqlite::ffi::SQLITE_FULL)));
255        assert!(!is_busy(&rusqlite::Error::QueryReturnedNoRows));
256
257        assert_eq!(busy_as(Err(failure(rusqlite::ffi::SQLITE_BUSY)), 7), Ok(7));
258        assert_eq!(
259            busy_as(Err(failure(rusqlite::ffi::SQLITE_LOCKED)), 7),
260            Ok(7)
261        );
262        assert_eq!(busy_as(Ok(3), 7), Ok(3));
263        let io = busy_as(Err(failure(rusqlite::ffi::SQLITE_IOERR)), 7).unwrap_err();
264        assert_eq!(io.sqlite_error_code(), Some(ErrorCode::SystemIoFailure));
265    }
266
267    #[test]
268    fn fresh_index_needs_no_compaction() {
269        let dir = tempfile::tempdir().unwrap();
270        let cache = FileHashCache::open(&dir.path().join("index.db")).unwrap();
271        let stats = cache.index_page_stats().unwrap();
272        assert!(stats.pages > 0);
273        assert_eq!(stats.free_pages, 0);
274        assert!(!stats.should_compact());
275        assert_eq!(cache.compact_index().unwrap(), IndexCompaction::NotNeeded);
276        assert_eq!(cache.compact_sparse_index().unwrap(), None);
277    }
278
279    #[test]
280    #[cfg(unix)]
281    fn dropped_table_is_reported_then_reclaimed() {
282        let dir = tempfile::tempdir().unwrap();
283        let path = dir.path().join("index.db");
284        let cache = sparse_index(&path);
285        let sparse = cache.index_page_stats().unwrap();
286        assert!(sparse.free_bytes() >= 80 << 20);
287        assert!(sparse.live_bytes() < 1 << 20);
288        assert!(sparse.should_compact());
289        assert_eq!(index_page_stats(cache.db()).unwrap(), sparse);
290
291        let IndexCompaction::Compacted { before, after } = cache.compact_index().unwrap() else {
292            panic!("large freelist must compact");
293        };
294        assert!(before >= 80 << 20);
295        assert!(after < 1 << 20);
296        assert_eq!(after, std::fs::metadata(&path).unwrap().len());
297        let compacted = cache.index_page_stats().unwrap();
298        assert_eq!(compacted.free_pages, 0);
299        assert_eq!(compacted.file_bytes(), after);
300        let kept: String = cache
301            .db()
302            .query_row("SELECT value FROM kept", [], |row| row.get(0))
303            .unwrap();
304        assert_eq!(kept, "live");
305        assert_eq!(cache.compact_index().unwrap(), IndexCompaction::NotNeeded);
306    }
307
308    #[test]
309    #[cfg(unix)]
310    fn open_reader_makes_compaction_busy_and_leaves_the_file_alone() {
311        let dir = tempfile::tempdir().unwrap();
312        let path = dir.path().join("index.db");
313        let cache = sparse_index(&path);
314        cache.db().pragma_update(None, "busy_timeout", "0").unwrap();
315        let reader = Connection::open(&path).unwrap();
316        reader.execute_batch("BEGIN").unwrap();
317        let _: String = reader
318            .query_row("SELECT value FROM kept", [], |row| row.get(0))
319            .unwrap();
320
321        let before = cache.index_page_stats().unwrap();
322        assert_eq!(cache.compact_index().unwrap(), IndexCompaction::Busy);
323        assert_eq!(cache.index_page_stats().unwrap(), before);
324        assert!(std::fs::metadata(&path).unwrap().len() >= 80 << 20);
325        let err = cache.compact_sparse_index().unwrap_err();
326        assert_eq!(err.to_string(), BUSY_BEFORE_VACUUM);
327
328        reader.execute_batch("COMMIT").unwrap();
329        assert!(matches!(
330            cache.compact_index().unwrap(),
331            IndexCompaction::Compacted { .. }
332        ));
333    }
334}