fathomdb-engine 0.4.1

Storage engine and write coordinator for the fathomdb agent datastore
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
/// Background actor that serializes async property-FTS rebuild tasks.
///
/// Modeled exactly on [`crate::writer::WriterActor`]: one OS thread,
/// `std::sync::mpsc`, `JoinHandle` for shutdown.  No tokio.
use std::path::Path;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};

use fathomdb_schema::SchemaManager;
use rusqlite::OptionalExtension;

use crate::{EngineError, sqlite};

/// Mode passed to `register_fts_property_schema_with_entries`.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RebuildMode {
    /// Legacy behavior: full rebuild runs inside the register transaction.
    Eager,
    /// 0.4.1+: schema is persisted synchronously; rebuild runs in background.
    #[default]
    Async,
}

/// A request to rebuild property-FTS for a single kind.
#[derive(Debug)]
pub struct RebuildRequest {
    pub kind: String,
    pub schema_id: i64,
}

/// Single-threaded actor that processes property-FTS rebuild requests one at
/// a time.  Shutdown is cooperative: drop the sender side to close the channel,
/// then join the thread.
///
/// The `RebuildActor` owns the `JoinHandle` only. The `SyncSender` lives in
/// [`crate::admin::AdminService`] so the service can enqueue rebuild requests
/// directly without going through the runtime.  The channel is created by
/// [`RebuildActor::create_channel`] and the two halves are distributed by
/// [`crate::runtime::EngineRuntime::open`].
#[derive(Debug)]
pub struct RebuildActor {
    thread_handle: Option<thread::JoinHandle<()>>,
}

impl RebuildActor {
    /// Create the mpsc channel used to communicate with the rebuild thread.
    ///
    /// Returns `(sender, actor)`.  The sender is given to
    /// [`crate::admin::AdminService`]; the actor is kept in
    /// [`crate::runtime::EngineRuntime`] for lifecycle management.
    ///
    /// # Errors
    /// Returns [`EngineError::Io`] if the thread cannot be spawned.
    pub fn start(
        path: impl AsRef<Path>,
        schema_manager: Arc<SchemaManager>,
        receiver: mpsc::Receiver<RebuildRequest>,
    ) -> Result<Self, EngineError> {
        let database_path = path.as_ref().to_path_buf();

        let handle = thread::Builder::new()
            .name("fathomdb-rebuild".to_owned())
            .spawn(move || {
                rebuild_loop(&database_path, &schema_manager, receiver);
            })
            .map_err(EngineError::Io)?;

        Ok(Self {
            thread_handle: Some(handle),
        })
    }
}

impl Drop for RebuildActor {
    fn drop(&mut self) {
        // The sender was already closed by AdminService (or dropped when the
        // engine closes).  Just join the thread.
        if let Some(handle) = self.thread_handle.take() {
            match handle.join() {
                Ok(()) => {}
                Err(payload) => {
                    if std::thread::panicking() {
                        trace_warn!(
                            "rebuild thread panicked during shutdown (suppressed: already panicking)"
                        );
                    } else {
                        std::panic::resume_unwind(payload);
                    }
                }
            }
        }
    }
}

// ── rebuild loop ────────────────────────────────────────────────────────────

/// Target wall-clock time for each batch transaction.
const BATCH_TARGET_MS: u128 = 1000;
/// Initial batch size.
const INITIAL_BATCH_SIZE: usize = 5000;

fn rebuild_loop(
    database_path: &Path,
    schema_manager: &Arc<SchemaManager>,
    receiver: mpsc::Receiver<RebuildRequest>,
) {
    trace_info!("rebuild thread started");

    let mut conn = match sqlite::open_connection(database_path) {
        Ok(conn) => conn,
        #[allow(clippy::used_underscore_binding)]
        Err(_error) => {
            trace_error!(error = %_error, "rebuild thread: database connection failed");
            return;
        }
    };

    #[allow(clippy::used_underscore_binding)]
    if let Err(_error) = schema_manager.bootstrap(&conn) {
        trace_error!(error = %_error, "rebuild thread: schema bootstrap failed");
        return;
    }

    for req in receiver {
        trace_info!(kind = %req.kind, schema_id = req.schema_id, "rebuild task started");
        match run_rebuild(&mut conn, &req) {
            Ok(()) => {
                trace_info!(kind = %req.kind, "rebuild task COMPLETE");
            }
            Err(error) => {
                trace_error!(kind = %req.kind, error = %error, "rebuild task failed");
                let _ = mark_failed(&conn, &req.kind, &error.to_string());
            }
        }
    }

    trace_info!("rebuild thread exiting");
}

#[allow(clippy::too_many_lines)]
fn run_rebuild(conn: &mut rusqlite::Connection, req: &RebuildRequest) -> Result<(), EngineError> {
    // Step 1: mark BUILDING.
    {
        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
        tx.execute(
            "UPDATE fts_property_rebuild_state SET state = 'BUILDING' \
             WHERE kind = ?1 AND schema_id = ?2",
            rusqlite::params![req.kind, req.schema_id],
        )?;
        tx.commit()?;
    }

    // Step 2: count nodes for this kind (plain SELECT, no tx needed).
    let rows_total: i64 = conn.query_row(
        "SELECT count(*) FROM nodes WHERE kind = ?1 AND superseded_at IS NULL",
        rusqlite::params![req.kind],
        |r| r.get(0),
    )?;

    {
        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
        tx.execute(
            "UPDATE fts_property_rebuild_state SET rows_total = ?1 WHERE kind = ?2",
            rusqlite::params![rows_total, req.kind],
        )?;
        tx.commit()?;
    }

    // Load the schema for this kind (plain SELECT).
    let (paths_json, separator): (String, String) = conn
        .query_row(
            "SELECT property_paths_json, separator FROM fts_property_schemas WHERE kind = ?1",
            rusqlite::params![req.kind],
            |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)),
        )
        .optional()?
        .ok_or_else(|| {
            EngineError::Bridge(format!("rebuild: schema for kind '{}' missing", req.kind))
        })?;
    let schema = crate::writer::parse_property_schema_json(&paths_json, &separator);

    // Step 3: batch-iterate nodes, insert into staging.
    let mut offset: i64 = 0;
    let mut batch_size = INITIAL_BATCH_SIZE;
    let mut rows_done: i64 = 0;

    loop {
        // Fetch a batch of node logical_ids + properties (plain SELECT — no tx needed for reads).
        let batch: Vec<(String, String)> = {
            let mut stmt = conn.prepare(
                "SELECT logical_id, properties FROM nodes \
                 WHERE kind = ?1 AND superseded_at IS NULL \
                 ORDER BY logical_id \
                 LIMIT ?2 OFFSET ?3",
            )?;
            stmt.query_map(
                rusqlite::params![
                    req.kind,
                    i64::try_from(batch_size).unwrap_or(i64::MAX),
                    offset
                ],
                |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)),
            )?
            .collect::<Result<Vec<_>, _>>()?
        };

        if batch.is_empty() {
            break;
        }

        let batch_len = batch.len();
        let batch_start = Instant::now();

        // Insert staging rows in a single short transaction.
        {
            let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;

            for (logical_id, properties_str) in &batch {
                let props: serde_json::Value =
                    serde_json::from_str(properties_str).unwrap_or_default();
                let (text, positions, _stats) =
                    crate::writer::extract_property_fts(&props, &schema);

                // Serialize positions to a compact JSON blob for later use at swap time.
                let positions_blob: Option<Vec<u8>> = if positions.is_empty() {
                    None
                } else {
                    let v: Vec<(usize, usize, &str)> = positions
                        .iter()
                        .map(|p| (p.start_offset, p.end_offset, p.leaf_path.as_str()))
                        .collect();
                    serde_json::to_vec(&v).ok()
                };

                let text_content = text.unwrap_or_default();

                tx.execute(
                    "INSERT INTO fts_property_rebuild_staging \
                     (kind, node_logical_id, text_content, positions_blob) \
                     VALUES (?1, ?2, ?3, ?4) \
                     ON CONFLICT(kind, node_logical_id) DO UPDATE \
                     SET text_content = excluded.text_content, \
                         positions_blob = excluded.positions_blob",
                    rusqlite::params![req.kind, logical_id, text_content, positions_blob],
                )?;
            }

            rows_done += i64::try_from(batch_len).unwrap_or(i64::MAX);
            let now_ms = now_unix_ms();
            tx.execute(
                "UPDATE fts_property_rebuild_state \
                 SET rows_done = ?1, last_progress_at = ?2 \
                 WHERE kind = ?3",
                rusqlite::params![rows_done, now_ms, req.kind],
            )?;
            tx.commit()?;
        }

        let elapsed_ms = batch_start.elapsed().as_millis();
        // Save the limit used for THIS batch before adjusting for the next one.
        let limit_used = batch_size;
        // Dynamically adjust batch size to target ~1s per batch.
        if elapsed_ms > 0 {
            let new_size = (batch_size as u128 * BATCH_TARGET_MS / elapsed_ms).clamp(100, 50_000);
            batch_size = usize::try_from(new_size).unwrap_or(50_000);
        }

        offset += i64::try_from(batch_len).unwrap_or(i64::MAX);

        // If the batch was smaller than the limit used for THIS query, we've reached the end.
        if batch_len < limit_used {
            break;
        }
    }

    // Step 4: mark SWAPPING.
    {
        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
        let now_ms = now_unix_ms();
        tx.execute(
            "UPDATE fts_property_rebuild_state \
             SET state = 'SWAPPING', last_progress_at = ?1 \
             WHERE kind = ?2",
            rusqlite::params![now_ms, req.kind],
        )?;
        tx.commit()?;
    }

    // Step 5: Final swap — atomic IMMEDIATE transaction replacing live FTS rows.
    {
        let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;

        // 5a. Delete old live FTS rows for this kind.
        tx.execute(
            "DELETE FROM fts_node_properties WHERE kind = ?1",
            rusqlite::params![req.kind],
        )?;

        // 5b. Insert new rows from staging into the live FTS table.
        tx.execute(
            "INSERT INTO fts_node_properties(node_logical_id, kind, text_content) \
             SELECT node_logical_id, kind, text_content \
             FROM fts_property_rebuild_staging WHERE kind = ?1",
            rusqlite::params![req.kind],
        )?;

        // 5c. Delete old position rows for this kind.
        tx.execute(
            "DELETE FROM fts_node_property_positions WHERE kind = ?1",
            rusqlite::params![req.kind],
        )?;

        // 5d. Re-populate fts_node_property_positions from positions_blob in staging.
        {
            let mut stmt = tx.prepare(
                "SELECT node_logical_id, positions_blob \
                 FROM fts_property_rebuild_staging \
                 WHERE kind = ?1 AND positions_blob IS NOT NULL",
            )?;
            let mut ins_pos = tx.prepare(
                "INSERT INTO fts_node_property_positions \
                 (node_logical_id, kind, start_offset, end_offset, leaf_path) \
                 VALUES (?1, ?2, ?3, ?4, ?5)",
            )?;

            let rows: Vec<(String, Vec<u8>)> = stmt
                .query_map(rusqlite::params![req.kind], |r| {
                    Ok((r.get::<_, String>(0)?, r.get::<_, Vec<u8>>(1)?))
                })?
                .collect::<Result<Vec<_>, _>>()?;

            for (node_logical_id, blob) in &rows {
                // positions_blob is JSON: Vec<(start, end, leaf_path)>
                let positions: Vec<(usize, usize, String)> =
                    serde_json::from_slice(blob).unwrap_or_default();
                for (start, end, leaf_path) in positions {
                    ins_pos.execute(rusqlite::params![
                        node_logical_id,
                        req.kind,
                        i64::try_from(start).unwrap_or(i64::MAX),
                        i64::try_from(end).unwrap_or(i64::MAX),
                        leaf_path,
                    ])?;
                }
            }
        }

        // 5e. Delete staging rows for this kind.
        tx.execute(
            "DELETE FROM fts_property_rebuild_staging WHERE kind = ?1",
            rusqlite::params![req.kind],
        )?;

        // 5f. Mark state COMPLETE.
        let now_ms = now_unix_ms();
        tx.execute(
            "UPDATE fts_property_rebuild_state \
             SET state = 'COMPLETE', last_progress_at = ?1 \
             WHERE kind = ?2",
            rusqlite::params![now_ms, req.kind],
        )?;

        tx.commit()?;
    }

    Ok(())
}

fn mark_failed(
    conn: &rusqlite::Connection,
    kind: &str,
    error_message: &str,
) -> Result<(), EngineError> {
    let now_ms = now_unix_ms();
    conn.execute(
        "UPDATE fts_property_rebuild_state \
         SET state = 'FAILED', error_message = ?1, last_progress_at = ?2 \
         WHERE kind = ?3",
        rusqlite::params![error_message, now_ms, kind],
    )?;
    Ok(())
}

fn now_unix_ms() -> i64 {
    now_unix_ms_pub()
}

/// Public-in-crate version of `now_unix_ms` so `admin.rs` can use it.
pub(crate) fn now_unix_ms_pub() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or(Duration::ZERO)
        .as_millis()
        .try_into()
        .unwrap_or(i64::MAX)
}

/// Rebuild progress row returned from `AdminService::get_property_fts_rebuild_state`.
#[derive(Debug)]
pub struct RebuildStateRow {
    pub kind: String,
    pub schema_id: i64,
    pub state: String,
    pub rows_total: Option<i64>,
    pub rows_done: i64,
    pub started_at: i64,
    pub is_first_registration: bool,
    pub error_message: Option<String>,
}

/// Public progress snapshot returned from
/// [`crate::coordinator::ExecutionCoordinator::get_property_fts_rebuild_progress`].
#[derive(Debug, Clone, serde::Serialize)]
pub struct RebuildProgress {
    /// Current state: `"PENDING"`, `"BUILDING"`, `"SWAPPING"`, `"COMPLETE"`, or `"FAILED"`.
    pub state: String,
    /// Total rows to process. `None` until the actor has counted the nodes.
    pub rows_total: Option<i64>,
    /// Rows processed so far.
    pub rows_done: i64,
    /// Unix milliseconds when the rebuild was registered.
    pub started_at: i64,
    /// Unix milliseconds of the last progress update, if any.
    pub last_progress_at: Option<i64>,
    /// Error message if `state == "FAILED"`.
    pub error_message: Option<String>,
}

/// Run crash recovery: mark any in-progress rebuilds as FAILED and clear their
/// staging rows.  Called by `EngineRuntime::open` before spawning the actor.
///
/// # Errors
/// Returns [`crate::EngineError`] if database access fails.
pub(crate) fn recover_interrupted_rebuilds(
    conn: &rusqlite::Connection,
) -> Result<(), crate::EngineError> {
    // Collect kinds that are in a non-terminal state.
    let kinds: Vec<String> = {
        let mut stmt = conn.prepare(
            "SELECT kind FROM fts_property_rebuild_state \
             WHERE state IN ('PENDING', 'BUILDING', 'SWAPPING')",
        )?;
        stmt.query_map([], |r| r.get::<_, String>(0))?
            .collect::<Result<Vec<_>, _>>()?
    };

    for kind in &kinds {
        conn.execute(
            "DELETE FROM fts_property_rebuild_staging WHERE kind = ?1",
            rusqlite::params![kind],
        )?;
        conn.execute(
            "UPDATE fts_property_rebuild_state \
             SET state = 'FAILED', error_message = 'interrupted by engine restart' \
             WHERE kind = ?1",
            rusqlite::params![kind],
        )?;
    }

    Ok(())
}