1use crate::Db;
2use rusqlite::{Transaction, TransactionBehavior};
3use std::time::{Duration, Instant};
4
5pub const JOB_SLICE_BUDGET_QUIET: Duration = Duration::from_millis(200);
6pub const JOB_SLICE_BUDGET_SHARED: Duration = Duration::from_millis(25);
7pub const JOB_SLICE_BREATH: Duration = Duration::from_millis(10);
8pub const SHARED_SLICES: usize = 20;
9pub const PEN_CONTENDED: Duration = Duration::from_millis(5);
10
11pub fn write_sliced<T>(
12 db: &Db,
13 items: impl IntoIterator<Item = T>,
14 apply: impl FnMut(&Transaction<'_>, T) -> rusqlite::Result<()>,
15) -> rusqlite::Result<usize> {
16 write_sliced_with(
17 db,
18 items,
19 apply,
20 JOB_SLICE_BUDGET_QUIET,
21 JOB_SLICE_BUDGET_SHARED,
22 JOB_SLICE_BREATH,
23 )
24}
25
26pub fn write_sliced_with<T>(
27 db: &Db,
28 items: impl IntoIterator<Item = T>,
29 mut apply: impl FnMut(&Transaction<'_>, T) -> rusqlite::Result<()>,
30 quiet: Duration,
31 shared: Duration,
32 breath: Duration,
33) -> rusqlite::Result<usize> {
34 let mut items = items.into_iter().peekable();
35 let mut applied = 0;
36 let mut shared_left = 0;
37 while items.peek().is_some() {
38 {
39 let connection = db.lock().unwrap_or_else(|poison| poison.into_inner());
40 let started = Instant::now();
41 let transaction =
42 Transaction::new_unchecked(&connection, TransactionBehavior::Immediate)?;
43 if started.elapsed() >= PEN_CONTENDED {
44 shared_left = SHARED_SLICES;
45 }
46 let budget = if shared_left == 0 {
47 quiet
48 } else {
49 shared_left -= 1;
50 shared
51 };
52 let opened = Instant::now();
53 for item in items.by_ref() {
54 apply(&transaction, item)?;
55 applied += 1;
56 if opened.elapsed() >= budget {
57 break;
58 }
59 }
60 transaction.commit()?;
61 }
62 if items.peek().is_some() && !breath.is_zero() {
63 std::thread::sleep(breath);
64 }
65 }
66 Ok(applied)
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use rusqlite::Connection;
73 use std::path::PathBuf;
74 use std::sync::{Arc, Mutex};
75
76 #[test]
77 fn the_five_slice_constants_keep_the_pen_ruling() {
78 assert_eq!(JOB_SLICE_BUDGET_QUIET, Duration::from_millis(200));
79 assert_eq!(JOB_SLICE_BUDGET_SHARED, Duration::from_millis(25));
80 assert_eq!(JOB_SLICE_BREATH, Duration::from_millis(10));
81 assert_eq!(SHARED_SLICES, 20);
82 assert_eq!(PEN_CONTENDED, Duration::from_millis(5));
83 }
84
85 #[test]
86 fn write_sliced_lands_every_item_across_many_slices() {
87 let connection = Connection::open_in_memory().unwrap();
88 connection
89 .execute_batch("CREATE TABLE values_seen(value INTEGER);")
90 .unwrap();
91 let db = Arc::new(Mutex::new(connection));
92 let applied = write_sliced_with(
93 &db,
94 0..20,
95 |transaction, value| {
96 transaction.execute("INSERT INTO values_seen VALUES (?1)", [value])?;
97 std::thread::sleep(Duration::from_millis(1));
98 Ok(())
99 },
100 Duration::from_millis(2),
101 Duration::from_millis(1),
102 Duration::ZERO,
103 )
104 .unwrap();
105 assert_eq!(applied, 20);
106 let connection = db.lock().unwrap();
107 assert_eq!(
108 connection
109 .query_row("SELECT COUNT(*) FROM values_seen", [], |row| row
110 .get::<_, i64>(0))
111 .unwrap(),
112 20
113 );
114 }
115
116 #[test]
117 fn write_sliced_keeps_committed_slices_and_rolls_back_only_the_failing_one() {
118 let connection = Connection::open_in_memory().unwrap();
119 connection
120 .execute_batch("CREATE TABLE values_seen(value INTEGER);")
121 .unwrap();
122 let db = Arc::new(Mutex::new(connection));
123 let result = write_sliced_with(
124 &db,
125 0..10,
126 |transaction, value| {
127 if value == 5 {
128 return Err(rusqlite::Error::InvalidQuery);
129 }
130 transaction.execute("INSERT INTO values_seen VALUES (?1)", [value])?;
131 std::thread::sleep(Duration::from_millis(1));
132 Ok(())
133 },
134 Duration::from_millis(2),
135 Duration::from_millis(1),
136 Duration::ZERO,
137 );
138 assert!(result.is_err());
139 let connection = db.lock().unwrap();
140 let kept = connection
141 .prepare("SELECT value FROM values_seen ORDER BY value")
142 .unwrap()
143 .query_map([], |row| row.get::<_, i64>(0))
144 .unwrap()
145 .collect::<rusqlite::Result<Vec<_>>>()
146 .unwrap();
147 assert!(!kept.is_empty());
148 assert!(kept.iter().all(|value| *value < 5));
149 }
150
151 #[test]
152 fn a_competing_writer_waits_no_longer_than_one_quiet_slice_and_breath() {
153 let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tmp/testscratch");
154 std::fs::create_dir_all(&root).unwrap();
155 let holder = tempfile::Builder::new()
156 .prefix("sliced-competing-writer-")
157 .tempdir_in(root)
158 .unwrap();
159 let path = holder.path().join("desk.db");
160 let connection = Connection::open(&path).unwrap();
161 connection.busy_timeout(Duration::from_secs(1)).unwrap();
162 connection
163 .pragma_update(None, "journal_mode", "WAL")
164 .unwrap();
165 connection
166 .execute_batch("CREATE TABLE values_seen(value INTEGER);")
167 .unwrap();
168 let db = Arc::new(Mutex::new(connection));
169 let mut companion = crate::open_companion(&path, None).unwrap();
170 companion
171 .busy_handler(Some(|_| {
172 std::thread::sleep(Duration::from_millis(1));
173 true
174 }))
175 .unwrap();
176 let (started_tx, started_rx) = std::sync::mpsc::channel();
177 let competitor = std::thread::spawn(move || {
178 started_rx.recv().unwrap();
179 let started = Instant::now();
180 let transaction = companion
181 .transaction_with_behavior(TransactionBehavior::Immediate)
182 .unwrap();
183 transaction
184 .execute("INSERT INTO values_seen VALUES (-1)", [])
185 .unwrap();
186 transaction.commit().unwrap();
187 started.elapsed()
188 });
189 let mut signalled = false;
190 let applied = write_sliced(&db, 0..1_000, |transaction, value| {
191 if !signalled {
192 started_tx.send(()).unwrap();
193 signalled = true;
194 }
195 transaction.execute("INSERT INTO values_seen VALUES (?1)", [value])?;
196 std::thread::sleep(Duration::from_millis(1));
197 Ok(())
198 })
199 .unwrap();
200 let waited = competitor.join().unwrap();
201 let ruled_budget = JOB_SLICE_BUDGET_QUIET + JOB_SLICE_BREATH;
202 let generous_wall_clock_bound = ruled_budget + Duration::from_millis(100);
203 eprintln!(
204 "US2_WAIT_MS={} US2_RULE_MS={} US2_BOUND_MS={}",
205 waited.as_millis(),
206 ruled_budget.as_millis(),
207 generous_wall_clock_bound.as_millis()
208 );
209 assert_eq!(applied, 1_000);
210 assert!(waited <= generous_wall_clock_bound);
211 }
212}