use rusqlite::Connection; use crate::core::backoff::compute_next_attempt_at; use crate::core::error::Result; use crate::core::time::now_ms; use crate::documents::SourceType; const DIRTY_SOURCES_BATCH_SIZE: usize = 500; pub fn mark_dirty_tx( tx: &rusqlite::Transaction<'_>, source_type: SourceType, source_id: i64, ) -> Result<()> { tx.execute( "INSERT INTO dirty_sources (source_type, source_id, queued_at) VALUES (?1, ?2, ?3) ON CONFLICT(source_type, source_id) DO UPDATE SET queued_at = excluded.queued_at, attempt_count = 0, last_attempt_at = NULL, last_error = NULL, next_attempt_at = NULL", rusqlite::params![source_type.as_str(), source_id, now_ms()], )?; Ok(()) } pub fn mark_dirty(conn: &Connection, source_type: SourceType, source_id: i64) -> Result<()> { conn.execute( "INSERT INTO dirty_sources (source_type, source_id, queued_at) VALUES (?1, ?2, ?3) ON CONFLICT(source_type, source_id) DO UPDATE SET queued_at = excluded.queued_at, attempt_count = 0, last_attempt_at = NULL, last_error = NULL, next_attempt_at = NULL", rusqlite::params![source_type.as_str(), source_id, now_ms()], )?; Ok(()) } pub fn get_dirty_sources(conn: &Connection) -> Result> { let now = now_ms(); let mut stmt = conn.prepare( "SELECT source_type, source_id FROM dirty_sources WHERE next_attempt_at IS NULL OR next_attempt_at <= ?1 ORDER BY attempt_count ASC, queued_at ASC LIMIT ?2", )?; let rows = stmt .query_map( rusqlite::params![now, DIRTY_SOURCES_BATCH_SIZE as i64], |row| { let st_str: String = row.get(0)?; let source_id: i64 = row.get(1)?; Ok((st_str, source_id)) }, )? .collect::, _>>()?; let mut results = Vec::with_capacity(rows.len()); for (st_str, source_id) in rows { let source_type = SourceType::parse(&st_str).ok_or_else(|| { crate::core::error::LoreError::Other(format!( "Invalid source_type in dirty_sources: {}", st_str )) })?; results.push((source_type, source_id)); } Ok(results) } pub fn clear_dirty(conn: &Connection, source_type: SourceType, source_id: i64) -> Result<()> { conn.execute( "DELETE FROM dirty_sources WHERE source_type = ?1 AND source_id = ?2", rusqlite::params![source_type.as_str(), source_id], )?; Ok(()) } pub fn record_dirty_error( conn: &Connection, source_type: SourceType, source_id: i64, error: &str, ) -> Result<()> { let now = now_ms(); let attempt_count: i64 = conn.query_row( "SELECT attempt_count FROM dirty_sources WHERE source_type = ?1 AND source_id = ?2", rusqlite::params![source_type.as_str(), source_id], |row| row.get(0), )?; let new_attempt = attempt_count + 1; let next_at = compute_next_attempt_at(now, new_attempt); conn.execute( "UPDATE dirty_sources SET attempt_count = ?1, last_attempt_at = ?2, last_error = ?3, next_attempt_at = ?4 WHERE source_type = ?5 AND source_id = ?6", rusqlite::params![ new_attempt, now, error, next_at, source_type.as_str(), source_id ], )?; Ok(()) } #[cfg(test)] mod tests { use super::*; fn setup_db() -> Connection { let conn = Connection::open_in_memory().unwrap(); conn.execute_batch(" CREATE TABLE dirty_sources ( source_type TEXT NOT NULL CHECK (source_type IN ('issue','merge_request','discussion')), source_id INTEGER NOT NULL, queued_at INTEGER NOT NULL, attempt_count INTEGER NOT NULL DEFAULT 0, last_attempt_at INTEGER, last_error TEXT, next_attempt_at INTEGER, PRIMARY KEY(source_type, source_id) ); CREATE INDEX idx_dirty_sources_next_attempt ON dirty_sources(next_attempt_at); ").unwrap(); conn } #[test] fn test_mark_dirty_inserts() { let conn = setup_db(); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); let count: i64 = conn .query_row("SELECT COUNT(*) FROM dirty_sources", [], |r| r.get(0)) .unwrap(); assert_eq!(count, 1); } #[test] fn test_mark_dirty_tx_inserts() { let mut conn = setup_db(); { let tx = conn.transaction().unwrap(); mark_dirty_tx(&tx, SourceType::Issue, 1).unwrap(); tx.commit().unwrap(); } let count: i64 = conn .query_row("SELECT COUNT(*) FROM dirty_sources", [], |r| r.get(0)) .unwrap(); assert_eq!(count, 1); } #[test] fn test_requeue_resets_backoff() { let conn = setup_db(); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); record_dirty_error(&conn, SourceType::Issue, 1, "test error").unwrap(); let attempt: i64 = conn .query_row( "SELECT attempt_count FROM dirty_sources WHERE source_id = 1", [], |r| r.get(0), ) .unwrap(); assert_eq!(attempt, 1); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); let attempt: i64 = conn .query_row( "SELECT attempt_count FROM dirty_sources WHERE source_id = 1", [], |r| r.get(0), ) .unwrap(); assert_eq!(attempt, 0); let next_at: Option = conn .query_row( "SELECT next_attempt_at FROM dirty_sources WHERE source_id = 1", [], |r| r.get(0), ) .unwrap(); assert!(next_at.is_none()); } #[test] fn test_get_respects_backoff() { let conn = setup_db(); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); conn.execute( "UPDATE dirty_sources SET next_attempt_at = 9999999999999 WHERE source_id = 1", [], ) .unwrap(); let results = get_dirty_sources(&conn).unwrap(); assert!(results.is_empty()); } #[test] fn test_get_orders_by_attempt_count() { let conn = setup_db(); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); conn.execute( "UPDATE dirty_sources SET attempt_count = 2 WHERE source_id = 1", [], ) .unwrap(); mark_dirty(&conn, SourceType::Issue, 2).unwrap(); let results = get_dirty_sources(&conn).unwrap(); assert_eq!(results.len(), 2); assert_eq!(results[0].1, 2); assert_eq!(results[1].1, 1); } #[test] fn test_batch_size_500() { let conn = setup_db(); for i in 0..600 { mark_dirty(&conn, SourceType::Issue, i).unwrap(); } let results = get_dirty_sources(&conn).unwrap(); assert_eq!(results.len(), 500); } #[test] fn test_clear_removes() { let conn = setup_db(); mark_dirty(&conn, SourceType::Issue, 1).unwrap(); clear_dirty(&conn, SourceType::Issue, 1).unwrap(); let count: i64 = conn .query_row("SELECT COUNT(*) FROM dirty_sources", [], |r| r.get(0)) .unwrap(); assert_eq!(count, 0); } #[test] fn test_drain_loop() { let conn = setup_db(); for i in 0..1200 { mark_dirty(&conn, SourceType::Issue, i).unwrap(); } let mut total = 0; loop { let batch = get_dirty_sources(&conn).unwrap(); if batch.is_empty() { break; } for (st, id) in &batch { clear_dirty(&conn, *st, *id).unwrap(); } total += batch.len(); } assert_eq!(total, 1200); } }