klukai_types/
sqlite3_restore.rs

1use std::{
2    fs::File,
3    io::{Read, Seek},
4    os::{fd::AsRawFd, unix::prelude::FileExt},
5    path::{Path, PathBuf},
6    time::{Duration, Instant},
7};
8
9use nix::{
10    fcntl::{FcntlArg, fcntl},
11    libc::{SEEK_SET, flock},
12};
13use tracing::info;
14
15// Database file lock bytes.
16const PENDING: i64 = 0x40000000;
17const RESERVED: i64 = 0x40000001;
18const SHARED: i64 = 0x40000002;
19
20// SHM file lock bytes.
21const WRITE: i64 = 120;
22const CKPT: i64 = 121;
23const RECOVER: i64 = 122;
24const READ0: i64 = 123;
25const READ1: i64 = 124;
26const READ2: i64 = 125;
27const READ3: i64 = 126;
28const READ4: i64 = 127;
29const DMS: i64 = 128;
30
31const MIN_DB_HDR_READ_LEN: usize = 20;
32
33#[derive(Debug, thiserror::Error)]
34pub enum Error {
35    #[error(transparent)]
36    Io(#[from] std::io::Error),
37
38    #[error("acquiring lock timed out")]
39    LockTimedOut,
40
41    #[error("inconsistent restore copy, expected {expected} bytes, but copied {actual}")]
42    InconsistentCopy { expected: u64, actual: u64 },
43
44    #[error("header read was too short ({0} bytes), need at least {MIN_DB_HDR_READ_LEN}")]
45    HeaderShortRead(usize),
46
47    #[error("database read and write format mismatched: {read} != {write}")]
48    ReadWriteFormatMismatch { read: u8, write: u8 },
49}
50
51pub struct Restored {
52    pub old_len: u64,
53    pub new_len: u64,
54    pub is_wal: bool,
55}
56
57pub fn restore<P1: AsRef<Path>, P2: AsRef<Path>>(
58    src: P1,
59    dst: P2,
60    timeout: Duration,
61) -> Result<Restored, Error> {
62    let mut src_db_file = std::fs::OpenOptions::new().read(true).open(src.as_ref())?;
63    let mut dst_db_file = std::fs::OpenOptions::new()
64        .read(true)
65        .write(true)
66        .create(true)
67        .truncate(false)
68        .open(dst.as_ref())?;
69
70    let src_meta = src_db_file.metadata()?;
71    let dst_meta = dst_db_file.metadata()?;
72
73    if dst_meta.len() == 0 {
74        info!("destination is empty, just copying.");
75        copy_check(&mut src_db_file, &mut dst_db_file, src_meta.len())?;
76        return Ok(Restored {
77            old_len: dst_meta.len(),
78            new_len: src_meta.len(),
79            is_wal: false, // this is only useful in tests
80        });
81    }
82
83    let mut dst_locked = lock_all(&mut dst_db_file, dst.as_ref(), timeout)?;
84
85    let dst_journal_path = format!("{}-journal", dst.as_ref().display());
86    info!("removing the journal file at: '{dst_journal_path}'");
87    if let Err(e) = std::fs::remove_file(PathBuf::from(dst_journal_path))
88        && e.kind() != std::io::ErrorKind::NotFound
89    {
90        return Err(e.into());
91    }
92
93    if dst_locked.is_wal() {
94        let dst_wal_path = format!("{}-wal", dst.as_ref().display());
95        info!("truncating WAL file '{dst_wal_path}'");
96        std::fs::OpenOptions::new()
97            .write(true)
98            .create(true)
99            .truncate(true)
100            .open(PathBuf::from(dst_wal_path))?;
101    }
102
103    dst_db_file.seek(std::io::SeekFrom::Start(0))?;
104    src_db_file.seek(std::io::SeekFrom::Start(0))?;
105
106    info!(
107        "copying from source to destination: {}",
108        src.as_ref().display()
109    );
110
111    copy_check(&mut src_db_file, &mut dst_db_file, src_meta.len())?;
112
113    if let Locked::Wal(ref mut dst_shm_file) = dst_locked {
114        dst_shm_file.write_at(&[0; 136], 0)?;
115    }
116
117    info!("done");
118
119    Ok(Restored {
120        old_len: dst_meta.len(),
121        new_len: src_meta.len(),
122        is_wal: dst_locked.is_wal(),
123    })
124}
125
126fn copy_check(src: &mut File, dst: &mut File, len: u64) -> Result<(), Error> {
127    let n = std::io::copy(src, dst)?;
128    if n != len {
129        return Err(Error::InconsistentCopy {
130            expected: len,
131            actual: n,
132        });
133    }
134
135    dst.set_len(len)?;
136    dst.sync_all()?;
137
138    Ok(())
139}
140
141pub enum Locked {
142    Wal(File),
143    Other,
144}
145
146impl Locked {
147    fn is_wal(&self) -> bool {
148        matches!(self, Locked::Wal(_))
149    }
150}
151
152pub fn lock_all<P: AsRef<Path>>(
153    db_file: &mut File,
154    db_path: P,
155    timeout: Duration,
156) -> Result<Locked, Error> {
157    info!("locking to determine journal mode");
158
159    lock(db_file, LockType::Read, PENDING, timeout)?;
160    lock(db_file, LockType::Read, SHARED, timeout)?;
161    lock(db_file, LockType::Unlock, PENDING, timeout)?;
162
163    info!("reading database mode");
164
165    let is_wal = is_wal_mode(db_file)?;
166
167    if !is_wal {
168        info!("destination database is in a non-WAL journal mode");
169
170        lock(db_file, LockType::Write, RESERVED, timeout)?;
171        lock(db_file, LockType::Write, PENDING, timeout)?;
172        lock(db_file, LockType::Write, SHARED, timeout)?;
173        return Ok(Locked::Other);
174    }
175
176    info!("destination database is in WAL journal mode");
177
178    let shm_file = std::fs::OpenOptions::new()
179        .read(true)
180        .write(true)
181        .create(true)
182        .truncate(false)
183        .open(PathBuf::from(format!("{}-shm", db_path.as_ref().display())))?;
184
185    lock(&shm_file, LockType::Read, DMS, timeout)?;
186    lock(&shm_file, LockType::Write, WRITE, timeout)?;
187    lock(&shm_file, LockType::Write, CKPT, timeout)?;
188    lock(&shm_file, LockType::Write, RECOVER, timeout)?;
189    lock(&shm_file, LockType::Write, READ0, timeout)?;
190    lock(&shm_file, LockType::Write, READ1, timeout)?;
191    lock(&shm_file, LockType::Write, READ2, timeout)?;
192    lock(&shm_file, LockType::Write, READ3, timeout)?;
193    lock(&shm_file, LockType::Write, READ4, timeout)?;
194
195    Ok(Locked::Wal(shm_file))
196}
197
198#[allow(clippy::unnecessary_cast)] // required for other platforms
199#[derive(Debug, Clone, Copy)]
200#[repr(i16)]
201enum LockType {
202    Read = nix::libc::F_RDLCK as i16,
203    Write = nix::libc::F_WRLCK as i16,
204    Unlock = nix::libc::F_UNLCK as i16,
205}
206
207fn lock(f: &File, l_type: LockType, l_start: i64, timeout: Duration) -> Result<(), Error> {
208    let mut l_len = 1i64;
209
210    if l_start == SHARED {
211        l_len = 510;
212    }
213
214    info!(
215        "acquiring lock ({l_type:?} ({}),{l_start},{l_len}) on fd: {}",
216        l_type as i16,
217        f.as_raw_fd()
218    );
219
220    let started_at = Instant::now();
221    loop {
222        let flock = flock {
223            l_start,
224            l_len,
225            l_pid: 0,
226            l_type: l_type as i16,
227            l_whence: SEEK_SET as i16,
228        };
229        match fcntl(f, FcntlArg::F_SETLK(&flock)) {
230            Ok(_) => {
231                info!("lock acquired");
232                return Ok(());
233            }
234            Err(_e) => {
235                // still waiting on lock...
236            }
237        }
238        if started_at.elapsed() > timeout {
239            return Err(Error::LockTimedOut);
240        }
241    }
242}
243
244fn is_wal_mode(f: &mut File) -> Result<bool, Error> {
245    let mut hdr = [0u8; 100];
246
247    match f.read(&mut hdr) {
248        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(false),
249        Err(e) => return Err(e.into()),
250        Ok(n) => {
251            if n < MIN_DB_HDR_READ_LEN {
252                return Err(Error::HeaderShortRead(n));
253            }
254        }
255    }
256
257    if hdr[18] != hdr[19] {
258        return Err(Error::ReadWriteFormatMismatch {
259            read: hdr[18],
260            write: hdr[19],
261        });
262    }
263
264    Ok(hdr[18] == 2)
265}
266
267#[cfg(test)]
268mod tests {
269    use rusqlite::Connection;
270
271    use super::*;
272
273    #[test]
274    fn test_basic_non_wal() -> Result<(), Box<dyn std::error::Error>> {
275        let tmpdir = tempfile::TempDir::new()?;
276
277        let src = tmpdir.path().join("src.db");
278        let dst = tmpdir.path().join("dst.db");
279
280        // seed the dst db
281        {
282            let conn = Connection::open(&dst)?;
283
284            conn.execute_batch(
285                "CREATE TABLE foo (a INT PRIMARY KEY, b INT); INSERT INTO foo (a,b) VALUES (1,1);",
286            )?;
287        }
288
289        {
290            let conn = Connection::open(&src)?;
291            conn.execute_batch(
292                "CREATE TABLE foo (a INT PRIMARY KEY, b INT); INSERT INTO foo (a,b) VALUES (1,2);",
293            )?;
294        }
295
296        let restored = restore(src, &dst, Duration::from_secs(2))?;
297        assert!(!restored.is_wal);
298
299        {
300            let conn = Connection::open(&dst)?;
301
302            let journal_mode: String =
303                conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
304
305            assert_eq!(journal_mode, "delete");
306
307            let (a, b): (i64, i64) =
308                conn.query_row("SELECT a,b FROM foo WHERE a = ?;", [1], |row| {
309                    Ok((row.get(0)?, row.get(1)?))
310                })?;
311
312            assert_eq!(a, 1);
313            assert_eq!(b, 2);
314        }
315
316        Ok(())
317    }
318
319    #[test]
320    fn test_basic_wal() -> Result<(), Box<dyn std::error::Error>> {
321        let tmpdir = tempfile::TempDir::new()?;
322
323        let src = tmpdir.path().join("src.db");
324        let dst = tmpdir.path().join("dst.db");
325
326        // seed dst
327        {
328            let conn = Connection::open(&dst)?;
329            conn.execute_batch("PRAGMA journal_mode = WAL;")?;
330
331            conn.execute_batch(
332                "CREATE TABLE foo (a INT PRIMARY KEY, b INT); INSERT INTO foo (a,b) VALUES (1,1);",
333            )?;
334
335            let journal_mode: String =
336                conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
337
338            assert_eq!(journal_mode, "wal");
339        }
340
341        {
342            let conn = Connection::open(&src)?;
343            conn.execute_batch("PRAGMA journal_mode = WAL;")?;
344
345            conn.execute_batch(
346                "CREATE TABLE foo (a INT PRIMARY KEY, b INT); INSERT INTO foo (a,b) VALUES (1,2);",
347            )?;
348
349            let journal_mode: String =
350                conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
351
352            assert_eq!(journal_mode, "wal");
353        }
354
355        let restored = restore(src, &dst, Duration::from_secs(2))?;
356        assert!(restored.is_wal);
357
358        {
359            let conn = Connection::open(&dst)?;
360
361            let journal_mode: String =
362                conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
363
364            assert_eq!(journal_mode, "wal");
365
366            let (a, b): (i64, i64) =
367                conn.query_row("SELECT a,b FROM foo WHERE a = ?;", [1], |row| {
368                    Ok((row.get(0)?, row.get(1)?))
369                })?;
370
371            assert_eq!(a, 1);
372            assert_eq!(b, 2);
373        }
374
375        Ok(())
376    }
377}
klukai_types/sqlite3_restore.rs

klukai_types/
sqlite3_restore.rs
