1use std::io::{Error, ErrorKind, Result};
2use std::path::{Path, PathBuf};
3use std::sync::Mutex;
4use std::time::Duration;
5
6use rusqlite::{params, Connection, OptionalExtension, TransactionBehavior};
7use serde_json::{json, Value};
8
9use crate::runtime::checkpoint_codec::{checkpoint_from_json, validate_checkpoint};
10use crate::runtime::state::{
11 check_claim, checkpoint_status_value, clear_claim, validate_claim, validate_renew, Checkpoint,
12 LeaseOperationClock, StateStore, StateStoreSpec,
13};
14
15const SELECT_CHECKPOINT: &str = "SELECT task_id, cycle_index, status, messages, cycles, \
16 shared_state, revision, claim_token, claimed_cycle, lease_expires_at_ms, terminal_result, budget_usage \
17 FROM checkpoints";
18
19#[derive(Debug)]
20pub struct SqliteStateStore {
21 connection: Mutex<Connection>,
22 location: Option<String>,
23}
24
25impl SqliteStateStore {
26 pub fn new(db_path: impl AsRef<Path>) -> Result<Self> {
27 let raw_path = db_path.as_ref();
28 let (open_path, location) = normalize_path(raw_path)?;
29 let connection = Connection::open(&open_path).map_err(sqlite_to_io)?;
30 connection
31 .busy_timeout(Duration::from_secs(5))
32 .map_err(sqlite_to_io)?;
33 connection
34 .execute_batch(
35 r#"
36 PRAGMA journal_mode=WAL;
37 CREATE TABLE IF NOT EXISTS checkpoints (
38 task_id TEXT PRIMARY KEY,
39 cycle_index INTEGER NOT NULL,
40 status TEXT NOT NULL,
41 messages TEXT NOT NULL,
42 cycles TEXT NOT NULL,
43 shared_state TEXT NOT NULL,
44 revision INTEGER NOT NULL DEFAULT 0,
45 claim_token TEXT,
46 claimed_cycle INTEGER,
47 lease_expires_at_ms INTEGER,
48 terminal_result TEXT,
49 budget_usage TEXT
50 );
51 "#,
52 )
53 .map_err(sqlite_to_io)?;
54 migrate_control_columns(&connection)?;
55 Ok(Self {
56 connection: Mutex::new(connection),
57 location,
58 })
59 }
60
61 pub fn close(self) -> Result<()> {
62 let connection = self
63 .connection
64 .into_inner()
65 .map_err(|_| Error::other("sqlite state store lock is poisoned"))?;
66 connection.close().map_err(|(_, error)| sqlite_to_io(error))
67 }
68}
69
70impl StateStore for SqliteStateStore {
71 fn create_checkpoint(&self, checkpoint: Checkpoint) -> Result<bool> {
72 let values = checkpoint_values(&checkpoint)?;
73 let changed = self
74 .connection
75 .lock()
76 .map_err(|_| poisoned())?
77 .execute(
78 r#"
79 INSERT OR IGNORE INTO checkpoints
80 (task_id, cycle_index, status, messages, cycles, shared_state,
81 revision, claim_token, claimed_cycle, lease_expires_at_ms, terminal_result,
82 budget_usage)
83 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
84 "#,
85 params![
86 values.task_id,
87 values.cycle_index,
88 values.status,
89 values.messages,
90 values.cycles,
91 values.shared_state,
92 values.revision,
93 values.claim_token,
94 values.claimed_cycle,
95 values.lease_expires_at_ms,
96 values.terminal_result,
97 values.budget_usage,
98 ],
99 )
100 .map_err(sqlite_to_io)?;
101 Ok(changed == 1)
102 }
103
104 fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
105 let values = checkpoint_values(&checkpoint)?;
106 self.connection
107 .lock()
108 .map_err(|_| poisoned())?
109 .execute(
110 r#"
111 INSERT OR REPLACE INTO checkpoints
112 (task_id, cycle_index, status, messages, cycles, shared_state,
113 revision, claim_token, claimed_cycle, lease_expires_at_ms, terminal_result,
114 budget_usage)
115 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
116 "#,
117 params![
118 values.task_id,
119 values.cycle_index,
120 values.status,
121 values.messages,
122 values.cycles,
123 values.shared_state,
124 values.revision,
125 values.claim_token,
126 values.claimed_cycle,
127 values.lease_expires_at_ms,
128 values.terminal_result,
129 values.budget_usage,
130 ],
131 )
132 .map_err(sqlite_to_io)?;
133 Ok(())
134 }
135
136 fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
137 let connection = self.connection.lock().map_err(|_| poisoned())?;
138 load_checkpoint_row(&connection, task_id)
139 }
140
141 fn claim_checkpoint(
142 &self,
143 task_id: &str,
144 cycle_index: u32,
145 claim_token: &str,
146 lease_expires_at_ms: u64,
147 now_ms: u64,
148 ) -> Result<Option<Checkpoint>> {
149 validate_claim(cycle_index, claim_token, lease_expires_at_ms, now_ms)?;
150 let mut connection = self.connection.lock().map_err(|_| poisoned())?;
151 let transaction = connection
152 .transaction_with_behavior(TransactionBehavior::Immediate)
153 .map_err(sqlite_to_io)?;
154 let Some(mut checkpoint) = load_checkpoint_row(&transaction, task_id)? else {
155 transaction.commit().map_err(sqlite_to_io)?;
156 return Ok(None);
157 };
158 check_claim(&checkpoint, cycle_index, now_ms)?;
159 let next_revision = checkpoint
160 .revision
161 .checked_add(1)
162 .ok_or_else(|| Error::new(ErrorKind::InvalidData, "checkpoint revision overflow"))?;
163 let changed = transaction
164 .execute(
165 r#"
166 UPDATE checkpoints
167 SET revision = ?1, claim_token = ?2, claimed_cycle = ?3, lease_expires_at_ms = ?4
168 WHERE task_id = ?5 AND revision = ?6
169 AND (claim_token IS NULL OR lease_expires_at_ms <= ?7)
170 "#,
171 params![
172 to_sql_u64(next_revision, "revision")?,
173 claim_token,
174 cycle_index,
175 to_sql_u64(lease_expires_at_ms, "lease_expires_at_ms")?,
176 task_id,
177 to_sql_u64(checkpoint.revision, "revision")?,
178 to_sql_u64(now_ms, "now_ms")?,
179 ],
180 )
181 .map_err(sqlite_to_io)?;
182 if changed != 1 {
183 return Err(Error::new(
184 ErrorKind::AlreadyExists,
185 format!("checkpoint cycle {cycle_index} for task {task_id} is already claimed"),
186 ));
187 }
188 transaction.commit().map_err(sqlite_to_io)?;
189 checkpoint.revision = next_revision;
190 checkpoint.claim_token = Some(claim_token.to_string());
191 checkpoint.claimed_cycle = Some(cycle_index);
192 checkpoint.lease_expires_at_ms = Some(lease_expires_at_ms);
193 Ok(Some(checkpoint))
194 }
195
196 fn commit_checkpoint(
197 &self,
198 mut checkpoint: Checkpoint,
199 claim_token: &str,
200 expected_revision: u64,
201 ) -> Result<bool> {
202 let claimed_cycle = checkpoint.cycle_index;
203 let next_revision = expected_revision
204 .checked_add(1)
205 .ok_or_else(|| Error::new(ErrorKind::InvalidData, "checkpoint revision overflow"))?;
206 checkpoint.revision = next_revision;
207 clear_claim(&mut checkpoint);
208 let values = checkpoint_values(&checkpoint)?;
209 let changed = self
210 .connection
211 .lock()
212 .map_err(|_| poisoned())?
213 .execute(
214 r#"
215 UPDATE checkpoints
216 SET cycle_index = ?1, status = ?2, messages = ?3, cycles = ?4,
217 shared_state = ?5, revision = ?6, claim_token = NULL,
218 claimed_cycle = NULL, lease_expires_at_ms = NULL, terminal_result = ?7,
219 budget_usage = ?8
220 WHERE task_id = ?9 AND revision = ?10 AND claim_token = ?11 AND claimed_cycle = ?12
221 "#,
222 params![
223 values.cycle_index,
224 values.status,
225 values.messages,
226 values.cycles,
227 values.shared_state,
228 values.revision,
229 values.terminal_result,
230 values.budget_usage,
231 values.task_id,
232 to_sql_u64(expected_revision, "revision")?,
233 claim_token,
234 claimed_cycle,
235 ],
236 )
237 .map_err(sqlite_to_io)?;
238 Ok(changed == 1)
239 }
240
241 fn renew_checkpoint_claim(
242 &self,
243 task_id: &str,
244 claim_token: &str,
245 expected_revision: u64,
246 lease_expires_at_ms: u64,
247 now_ms: u64,
248 ) -> Result<bool> {
249 validate_renew(claim_token, expected_revision, lease_expires_at_ms, now_ms)?;
250 let clock = LeaseOperationClock::new(now_ms);
251 let mut connection = self.connection.lock().map_err(|_| poisoned())?;
252 let transaction = connection
253 .transaction_with_behavior(TransactionBehavior::Immediate)
254 .map_err(sqlite_to_io)?;
255 let current_now_ms = clock.now_ms();
256 if lease_expires_at_ms <= current_now_ms {
257 transaction.commit().map_err(sqlite_to_io)?;
258 return Ok(false);
259 }
260 let changed = transaction
261 .execute(
262 r#"
263 UPDATE checkpoints
264 SET lease_expires_at_ms = ?1
265 WHERE task_id = ?2 AND revision = ?3 AND claim_token = ?4
266 AND lease_expires_at_ms > ?5
267 "#,
268 params![
269 to_sql_u64(lease_expires_at_ms, "lease_expires_at_ms")?,
270 task_id,
271 to_sql_u64(expected_revision, "revision")?,
272 claim_token,
273 to_sql_u64(current_now_ms, "now_ms")?,
274 ],
275 )
276 .map_err(sqlite_to_io)?;
277 transaction.commit().map_err(sqlite_to_io)?;
278 Ok(changed == 1)
279 }
280
281 fn finalize_checkpoint(
282 &self,
283 mut checkpoint: Checkpoint,
284 expected_revision: u64,
285 ) -> Result<bool> {
286 if checkpoint.terminal_result.is_none() {
287 return Err(Error::new(
288 ErrorKind::InvalidInput,
289 "finalized checkpoint must include terminal_result",
290 ));
291 }
292 checkpoint.revision = expected_revision
293 .checked_add(1)
294 .ok_or_else(|| Error::new(ErrorKind::InvalidData, "checkpoint revision overflow"))?;
295 clear_claim(&mut checkpoint);
296 let values = checkpoint_values(&checkpoint)?;
297 let changed = self
298 .connection
299 .lock()
300 .map_err(|_| poisoned())?
301 .execute(
302 r#"
303 UPDATE checkpoints
304 SET cycle_index = ?1, status = ?2, messages = ?3, cycles = ?4,
305 shared_state = ?5, revision = ?6, claim_token = NULL,
306 claimed_cycle = NULL, lease_expires_at_ms = NULL, terminal_result = ?7,
307 budget_usage = ?8
308 WHERE task_id = ?9 AND revision = ?10 AND claim_token IS NULL
309 AND terminal_result IS NULL
310 "#,
311 params![
312 values.cycle_index,
313 values.status,
314 values.messages,
315 values.cycles,
316 values.shared_state,
317 values.revision,
318 values.terminal_result,
319 values.budget_usage,
320 values.task_id,
321 to_sql_u64(expected_revision, "revision")?,
322 ],
323 )
324 .map_err(sqlite_to_io)?;
325 Ok(changed == 1)
326 }
327
328 fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
329 self.connection
330 .lock()
331 .map_err(|_| poisoned())?
332 .execute(
333 "DELETE FROM checkpoints WHERE task_id = ?1",
334 params![task_id],
335 )
336 .map_err(sqlite_to_io)?;
337 Ok(())
338 }
339
340 fn acknowledge_terminal(&self, task_id: &str, expected_revision: u64) -> Result<bool> {
341 let changed = self
342 .connection
343 .lock()
344 .map_err(|_| poisoned())?
345 .execute(
346 "DELETE FROM checkpoints WHERE task_id = ?1 AND revision = ?2 AND terminal_result IS NOT NULL",
347 params![task_id, to_sql_u64(expected_revision, "revision")?],
348 )
349 .map_err(sqlite_to_io)?;
350 Ok(changed == 1)
351 }
352
353 fn list_checkpoints(&self) -> Result<Vec<String>> {
354 let connection = self.connection.lock().map_err(|_| poisoned())?;
355 let mut statement = connection
356 .prepare("SELECT task_id FROM checkpoints ORDER BY task_id")
357 .map_err(sqlite_to_io)?;
358 let rows = statement
359 .query_map([], |row| row.get::<_, String>(0))
360 .map_err(sqlite_to_io)?;
361 rows.collect::<rusqlite::Result<Vec<_>>>()
362 .map_err(sqlite_to_io)
363 }
364
365 fn state_store_spec(&self) -> Option<StateStoreSpec> {
366 self.location
367 .as_ref()
368 .and_then(|location| StateStoreSpec::sqlite(location).ok())
369 }
370}
371
372#[derive(Debug)]
373struct CheckpointValues {
374 task_id: String,
375 cycle_index: u32,
376 status: String,
377 messages: String,
378 cycles: String,
379 shared_state: String,
380 revision: i64,
381 claim_token: Option<String>,
382 claimed_cycle: Option<u32>,
383 lease_expires_at_ms: Option<i64>,
384 terminal_result: Option<String>,
385 budget_usage: Option<String>,
386}
387
388fn checkpoint_values(checkpoint: &Checkpoint) -> Result<CheckpointValues> {
389 validate_checkpoint(checkpoint)?;
390 Ok(CheckpointValues {
391 task_id: checkpoint.task_id.clone(),
392 cycle_index: checkpoint.cycle_index,
393 status: checkpoint_status_value(checkpoint.status).to_string(),
394 messages: serde_json::to_string(
395 &checkpoint
396 .messages
397 .iter()
398 .map(crate::types::Message::to_dict)
399 .collect::<Vec<_>>(),
400 )
401 .map_err(json_to_io)?,
402 cycles: serde_json::to_string(
403 &checkpoint
404 .cycles
405 .iter()
406 .map(crate::types::CycleRecord::to_dict)
407 .collect::<Vec<_>>(),
408 )
409 .map_err(json_to_io)?,
410 shared_state: serde_json::to_string(&checkpoint.shared_state).map_err(json_to_io)?,
411 revision: to_sql_u64(checkpoint.revision, "revision")?,
412 claim_token: checkpoint.claim_token.clone(),
413 claimed_cycle: checkpoint.claimed_cycle,
414 lease_expires_at_ms: checkpoint
415 .lease_expires_at_ms
416 .map(|value| to_sql_u64(value, "lease_expires_at_ms"))
417 .transpose()?,
418 terminal_result: checkpoint
419 .terminal_result
420 .as_ref()
421 .map(|result| serde_json::to_string(&result.to_dict()).map_err(json_to_io))
422 .transpose()?,
423 budget_usage: checkpoint
424 .budget_usage
425 .as_ref()
426 .map(|usage| serde_json::to_string(usage).map_err(json_to_io))
427 .transpose()?,
428 })
429}
430
431type CheckpointRow = (
432 String,
433 u32,
434 String,
435 String,
436 String,
437 String,
438 i64,
439 Option<String>,
440 Option<u32>,
441 Option<i64>,
442 Option<String>,
443 Option<String>,
444);
445
446fn load_checkpoint_row(connection: &Connection, task_id: &str) -> Result<Option<Checkpoint>> {
447 let row = connection
448 .query_row(
449 &format!("{SELECT_CHECKPOINT} WHERE task_id = ?1"),
450 params![task_id],
451 |row| {
452 Ok((
453 row.get(0)?,
454 row.get(1)?,
455 row.get(2)?,
456 row.get(3)?,
457 row.get(4)?,
458 row.get(5)?,
459 row.get(6)?,
460 row.get(7)?,
461 row.get(8)?,
462 row.get(9)?,
463 row.get(10)?,
464 row.get(11)?,
465 ))
466 },
467 )
468 .optional()
469 .map_err(sqlite_to_io)?;
470 row.map(checkpoint_from_row).transpose()
471}
472
473fn checkpoint_from_row(row: CheckpointRow) -> Result<Checkpoint> {
474 let revision = u64::try_from(row.6).map_err(|_| {
475 Error::new(
476 ErrorKind::InvalidData,
477 "checkpoint revision must be non-negative",
478 )
479 })?;
480 let lease_expires_at_ms = row
481 .9
482 .map(|value| {
483 u64::try_from(value).map_err(|_| {
484 Error::new(
485 ErrorKind::InvalidData,
486 "checkpoint lease_expires_at_ms must be non-negative",
487 )
488 })
489 })
490 .transpose()?;
491 let payload = json!({
492 "task_id": row.0,
493 "cycle_index": row.1,
494 "status": row.2,
495 "messages": parse_json(&row.3, "messages")?,
496 "cycles": parse_json(&row.4, "cycles")?,
497 "shared_state": parse_json(&row.5, "shared_state")?,
498 "revision": revision,
499 "claim_token": row.7,
500 "claimed_cycle": row.8,
501 "lease_expires_at_ms": lease_expires_at_ms,
502 "terminal_result": row.10.as_deref().map(|value| parse_json(value, "terminal_result")).transpose()?,
503 "budget_usage": row.11.as_deref().map(|value| parse_json(value, "budget_usage")).transpose()?,
504 });
505 checkpoint_from_json(&payload.to_string())
506}
507
508fn normalize_path(path: &Path) -> Result<(PathBuf, Option<String>)> {
509 if path == Path::new(":memory:") {
510 return Ok((PathBuf::from(":memory:"), None));
511 }
512 let absolute = if path.is_absolute() {
513 path.to_path_buf()
514 } else {
515 std::env::current_dir()?.join(path)
516 };
517 let location = absolute.to_string_lossy().to_string();
518 Ok((absolute, Some(location)))
519}
520
521fn migrate_control_columns(connection: &Connection) -> Result<()> {
522 let mut statement = connection
523 .prepare("PRAGMA table_info(checkpoints)")
524 .map_err(sqlite_to_io)?;
525 let columns = statement
526 .query_map([], |row| row.get::<_, String>(1))
527 .map_err(sqlite_to_io)?
528 .collect::<rusqlite::Result<std::collections::BTreeSet<_>>>()
529 .map_err(sqlite_to_io)?;
530 drop(statement);
531 for (name, declaration) in [
532 ("revision", "INTEGER NOT NULL DEFAULT 0"),
533 ("claim_token", "TEXT"),
534 ("claimed_cycle", "INTEGER"),
535 ("lease_expires_at_ms", "INTEGER"),
536 ("terminal_result", "TEXT"),
537 ("budget_usage", "TEXT"),
538 ] {
539 if !columns.contains(name) {
540 connection
541 .execute(
542 &format!("ALTER TABLE checkpoints ADD COLUMN {name} {declaration}"),
543 [],
544 )
545 .map_err(sqlite_to_io)?;
546 }
547 }
548 Ok(())
549}
550
551fn parse_json(raw: &str, field: &str) -> Result<Value> {
552 serde_json::from_str(raw).map_err(|error| {
553 Error::new(
554 ErrorKind::InvalidData,
555 format!("invalid checkpoint {field} JSON: {error}"),
556 )
557 })
558}
559
560fn to_sql_u64(value: u64, field: &str) -> Result<i64> {
561 i64::try_from(value).map_err(|_| {
562 Error::new(
563 ErrorKind::InvalidData,
564 format!("checkpoint {field} exceeds SQLite INTEGER range"),
565 )
566 })
567}
568
569fn poisoned() -> Error {
570 Error::other("sqlite state store lock is poisoned")
571}
572
573fn json_to_io(error: serde_json::Error) -> Error {
574 Error::new(ErrorKind::InvalidData, error)
575}
576
577fn sqlite_to_io(error: rusqlite::Error) -> Error {
578 Error::other(error.to_string())
579}