Add safeguards against ML database corruption and streamline recovery process - #22478
Conversation
There was a problem hiding this comment.
Pull request overview
Adds corruption detection and safe recovery mechanics for the ML SQLite database (ml.db) so the agent can stop using a damaged DB, quarantine it, and recreate a fresh one on the next start (retraining models instead of repeatedly failing DB operations).
Changes:
- Introduces an in-process “poison” flag (
ml_db_unusable) plusml_db_mark_corrupt()to latch corruption and drop a quarantine sentinel. - Updates ML DB operations (add/delete/prune/load + flush) to short-circuit after corruption is detected and to preserve useful logging context.
- Adds startup logic to consume the sentinel, quarantine
ml.dbasml.db.bad, delete-wal/-shm, and open a fresh DB.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/ml/ml.cc |
Adds corruption latching/sentinel creation and short-circuiting in several ML DB paths. |
src/ml/ml_public.cc |
Consumes the sentinel on startup to quarantine and recreate ml.db (with WAL/SHM cleanup). |
src/ml/ml_private.h |
Exposes the new corruption flag and ml_db_mark_corrupt() to ML code. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
No issues found across 3 files
Confidence score: 5/5
- Automated review surfaced no issues in the provided summaries.
- No files require special attention.
Architecture diagram
sequenceDiagram
participant Cache as Cache Dir
participant ML as ML Module
participant SQLite as SQLite ml.db
participant Worker as ML Worker
participant Dim as Dimension
Note over ML,SQLite: NEW: Corruption Detection & Recovery Flow
ML->>ML: Startup: check for .ml.db.delete sentinel
alt Sentinel exists
ML->>Cache: unlink(sentinel)
ML->>SQLite: rename(ml.db -> ml.db.bad.<usec>)
alt Rename succeeds or ml.db missing
ML->>Cache: Clean up -wal and -shm files
ML->>SQLite: sqlite3_open() creates fresh DB
else Rename fails (non-ENOENT)
ML->>Cache: Restore sentinel via O_CREAT|O_EXCL
ML->>SQLite: Still opens existing corrupt DB
end
else No sentinel
ML->>SQLite: Normal sqlite3_open()
end
Note over ML,Worker: Runtime: Corruption Detection
Worker->>Worker: Process pending models
Worker->>SQLite: BEGIN TRANSACTION
SQLite-->>Worker: sqlite_rc
Worker->>ML: ml_db_mark_if_corrupt(sqlite_rc)
alt Corruption detected (SQLITE_CORRUPT/SQLITE_NOTADB)
ML->>ML: ml_db_mark_corrupt(rc)
ML->>ML: Set ml_db_unusable = true
ML->>Cache: Create .ml.db.delete sentinel
Note over ML,Worker: All subsequent operations skip
end
Worker->>Dim: Add/delete/prune model
Dim->>SQLite: Execute SQL statement
SQLite-->>Dim: SQL error (corruption)
Dim->>ML: ml_db_mark_if_corrupt(rc)
alt Corruption detected
ML->>ML: Same corruption handling
end
Note over Worker,Dim: On load models with corruption
Dim->>SQLite: Prepare & step through models
SQLite-->>Dim: SQLITE_CORRUPT step result
Dim->>ML: ml_db_mark_if_corrupt(step_rc)
alt Step rc is corruption
Dim->>Dim: Clear km_contexts
Dim->>Dim: Set ts = UNTRAINED
end
Note over Worker: Transaction finalization
Worker->>Worker: Check ml_db_unusable
alt Became unusable during transaction
Worker->>Worker: Skip rollback
Worker->>Worker: Skip vacuum
Worker->>Worker: Clear pending model info
else Still usable
alt Transaction failed
Worker->>SQLite: ROLLBACK
SQLite-->>Worker: sqlite_rc
else Successful
Worker->>SQLite: COMMIT
SQLite-->>Worker: sqlite_rc
end
Worker->>SQLite: VACUUM (periodic)
end
Note over ML,Dim: On next ml_db operation
ML->>ML: Check ml_db_unusable flag
alt Flag is true
ML->>ML: Short-circuit, return early
else Flag is false
ML->>SQLite: Continue normal operation
end
Note over ML,Dim: On agent restart after corruption
ML->>Cache: Sentinel present from prior session
ML->>SQLite: Rename corrupt DB -> ml.db.bad.<usec>
ML->>SQLite: Create fresh ml.db with sqlite3_open()
ML->>Dim: Load models (none exist)
Dim->>Dim: All dimensions remain UNTRAINED
ML->>ML: Trigger retraining immediately
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (3)
src/ml/ml_public.cc:577
- The new quarantine mechanism relies on
ml_db_mark_corrupt()being called, but init-time failures (migration/config PRAGMAs/table creation) can hitSQLITE_CORRUPT/SQLITE_NOTADBbefore any of the runtime operations that now callml_db_mark_if_corrupt(). In that caseml_dbis closed and ML stays disabled on every restart, but no.ml.db.deletesentinel is written so quarantine never happens. Consider latching corruption duringml_init()setup failures too (e.g., inspectsqlite3_errcode()/sqlite3_extended_errcode()afterconfigure_sqlite_database()/sqlite3_exec()failures and callml_db_mark_corrupt()when appropriate).
int rc = sqlite3_open(path, &ml_db);
if (rc != SQLITE_OK) {
error_report("Failed to initialize database at %s, due to \"%s\"", path, sqlite3_errstr(rc));
sqlite3_close(ml_db);
ml_db = NULL;
}
src/ml/ml.cc:398
- This function marks corruption when
execute_insert()fails, butsqlite3_reset(res)(immediately after this block) may also returnSQLITE_CORRUPT/SQLITE_NOTADBeven if the delete step succeeded. Consider callingml_db_mark_if_corrupt(rc)on the reset failure path too, so cleanup-time corruption reliably latchesml_db_unusableand triggers quarantine behavior.
rc = execute_insert(res);
if (unlikely(rc != SQLITE_DONE)) {
error_report("Failed to delete models, rc = %d", rc);
ml_db_mark_if_corrupt(rc);
return rc;
}
src/ml/ml.cc:456
- Similar to the other ML DB operations: corruption is latched when
execute_insert()fails, butsqlite3_reset(res)after this block can also returnSQLITE_CORRUPT/SQLITE_NOTADB. Consider addingml_db_mark_if_corrupt(rc)to the reset failure path too so corruption detected during cleanup still poisons the session and writes the sentinel.
rc = execute_insert(res);
if (unlikely(rc != SQLITE_DONE)) {
error_report("Failed to prune old models, rc = %d", rc);
ml_db_mark_if_corrupt(rc);
return rc;
}
There was a problem hiding this comment.
1 issue found across 2 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/ml/ml_public.cc">
<violation number="1" location="src/ml/ml_public.cc:522">
P2: Quarantine is triggered on any sentinel unlink error, which can repeatedly rotate a healthy `ml.db` on every startup if the sentinel path cannot be unlinked.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
There was a problem hiding this comment.
1 issue found across 1 file (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
… process - Introduced `ml_db_unusable` flag to prevent further operations on a corrupted database. - Added `ml_db_mark_corrupt()` to handle SQLite corruption errors (`SQLITE_CORRUPT`, `SQLITE_NOTADB`). - Implemented logic to quarantine and replace corrupt ML databases, including WAL file cleanup. - Updated database operation functions to detect corruption and bail gracefully, preserving log details for diagnostics. - Enhanced session behavior by ensuring a fresh ML database is recreated upon restart after corruption detection.
…upt()` helper - Consolidated database corruption checks (`SQLITE_CORRUPT`, `SQLITE_NOTADB`) into a reusable helper function `ml_db_mark_if_corrupt()`. - Replaced repetitive corruption handling code in ML database operations with the new helper for improved readability and maintainability.
- Improved sentinel-based handling to ensure retries on startup after failures. - Added logic to handle existing `ml.db.bad` files by appending a timestamp to avoid overwrite issues. - Introduced safeguards to clear WAL/SHM siblings reliably during corruption quarantine. - Enhanced dimension reset behavior to prevent stale models after corruption detection. - Optimized operations to skip rollback and vacuum when the database is marked unusable.
- Replace overwrite-prone `ml.db.bad` renames with timestamped destinations to ensure unique handling across Windows and POSIX. - Simplify quarantine and sentinel restoration flow, avoiding TOCTOU races by using `unlink()` for existence checks. - Ensure safe retries on startup by restoring sentinel after failed quarantines.
- Use `O_CREAT|O_EXCL` for atomic create-or-fail behavior in sentinel restoration. - Improve safety by avoiding symlink swaps during sentinel creation attempts.
- Use `O_CREAT|O_EXCL` to mitigate symlink TOCTOU attacks during sentinel creation. - Enhance SQLite corruption recovery by logging and marking corruption (`ml_db_mark_if_corrupt`). - Refactor transaction handling to incorporate detailed error checks and proper rollback behavior.
- Add microsecond-resolution timestamps to `ml.db.bad` renames for collision prevention across consecutive restarts. - Update logging to reflect timestamped quarantine behavior clearly. - Improve cross-platform support for rename operations on corrupted ML databases.
- Refine sentinel creation to distinguish between occupied path retries and actual failures. - Enhance corruption detection by separately handling step- and cleanup-time errors. - Adjust dimension state reset logic for consistent handling of partial results.
- Refine sentinel unlink logic to handle edge cases (e.g., permission errors) with detailed logs. - Update dimension reset logic to ensure consistency by rolling back partial results on non-`SQLITE_DONE` errors. - Add `ml_db_mark_if_corrupt()` to additional SQLite operations for improved corruption tracking.
- Introduce `ml_db_is_unusable()` for consistent access of unusable flag. - Replace direct `ml_db_unusable` operations with accessor for atomic contract enforcement. - Refine sentinel unlink logic to prevent re-quarantining on every restart. - Ensure robust handling of SQLite primary and extended error codes during corruption detection. - Replace hardcoded buffer sizes in `snprintfz()` calls with `sizeof` for improved safety.
…perations - Update `ml_db_mark_if_corrupt()` to handle both primary and extended SQLite error codes. - Integrate corruption checks into additional ML database operations to ensure consistent handling. - Refactor and expose `ml_db_mark_if_corrupt()` for broader usage across public and private APIs.
- Improve logging for deferred quarantine attempts during sentinel creation failures (e.g., permission errors, read-only mounts). - Add safety check for `ml_db` null state to prevent redundant operations and log spam on initialization failures. - Update sentinel and database handling workflows for enhanced robustness against edge cases.
- Introduce `execute_and_reset_model_stmt()` for consistent execution and reset of model-table statements. - Add `handle_model_bind_fail()` for centralized bind-failure handling with improved logging. - Replace duplicate logic in "store model," "delete models," and "prune old models" operations with these helpers.
…tabases - Add `ml_db_force_unusable()` to mark the database unusable without disk writes, enabling safe bypass of poisoned databases at startup. - Improve sentinel and transaction error handling to avoid retriggering corruption scenarios and enhance error recovery mechanisms. - Refine SQLite statement execution and reset logic for consistent corruption tracking.
|
thiagoftsm
left a comment
There was a problem hiding this comment.
PR is working as expected after few hours, LGTM!
… process (#22478) * Add safeguards against ML database corruption and streamline recovery process - Introduced `ml_db_unusable` flag to prevent further operations on a corrupted database. - Added `ml_db_mark_corrupt()` to handle SQLite corruption errors (`SQLITE_CORRUPT`, `SQLITE_NOTADB`). - Implemented logic to quarantine and replace corrupt ML databases, including WAL file cleanup. - Updated database operation functions to detect corruption and bail gracefully, preserving log details for diagnostics. - Enhanced session behavior by ensuring a fresh ML database is recreated upon restart after corruption detection. * Refactor corruption handling logic by introducing `ml_db_mark_if_corrupt()` helper - Consolidated database corruption checks (`SQLITE_CORRUPT`, `SQLITE_NOTADB`) into a reusable helper function `ml_db_mark_if_corrupt()`. - Replaced repetitive corruption handling code in ML database operations with the new helper for improved readability and maintainability. * Enhance ML database corruption recovery logic - Improved sentinel-based handling to ensure retries on startup after failures. - Added logic to handle existing `ml.db.bad` files by appending a timestamp to avoid overwrite issues. - Introduced safeguards to clear WAL/SHM siblings reliably during corruption quarantine. - Enhanced dimension reset behavior to prevent stale models after corruption detection. - Optimized operations to skip rollback and vacuum when the database is marked unusable. * Refactor ML database quarantine logic - Replace overwrite-prone `ml.db.bad` renames with timestamped destinations to ensure unique handling across Windows and POSIX. - Simplify quarantine and sentinel restoration flow, avoiding TOCTOU races by using `unlink()` for existence checks. - Ensure safe retries on startup by restoring sentinel after failed quarantines. * Enhance sentinel restoration logic to prevent TOCTOU vulnerability - Use `O_CREAT|O_EXCL` for atomic create-or-fail behavior in sentinel restoration. - Improve safety by avoiding symlink swaps during sentinel creation attempts. * Harden ML database corruption handling and sentinel creation logic - Use `O_CREAT|O_EXCL` to mitigate symlink TOCTOU attacks during sentinel creation. - Enhance SQLite corruption recovery by logging and marking corruption (`ml_db_mark_if_corrupt`). - Refactor transaction handling to incorporate detailed error checks and proper rollback behavior. * Refactor ML database corruption quarantine logic - Add microsecond-resolution timestamps to `ml.db.bad` renames for collision prevention across consecutive restarts. - Update logging to reflect timestamped quarantine behavior clearly. - Improve cross-platform support for rename operations on corrupted ML databases. * Improve ML database corruption handling and logging - Refine sentinel creation to distinguish between occupied path retries and actual failures. - Enhance corruption detection by separately handling step- and cleanup-time errors. - Adjust dimension state reset logic for consistent handling of partial results. * Fix compilation error * Improve ML database corruption handling and sentinel unlink behavior - Refine sentinel unlink logic to handle edge cases (e.g., permission errors) with detailed logs. - Update dimension reset logic to ensure consistency by rolling back partial results on non-`SQLITE_DONE` errors. - Add `ml_db_mark_if_corrupt()` to additional SQLite operations for improved corruption tracking. * Refactor ML database corruption handling - Introduce `ml_db_is_unusable()` for consistent access of unusable flag. - Replace direct `ml_db_unusable` operations with accessor for atomic contract enforcement. - Refine sentinel unlink logic to prevent re-quarantining on every restart. - Ensure robust handling of SQLite primary and extended error codes during corruption detection. - Replace hardcoded buffer sizes in `snprintfz()` calls with `sizeof` for improved safety. * Extend corruption handling with `ml_db_mark_if_corrupt()` in SQLite operations - Update `ml_db_mark_if_corrupt()` to handle both primary and extended SQLite error codes. - Integrate corruption checks into additional ML database operations to ensure consistent handling. - Refactor and expose `ml_db_mark_if_corrupt()` for broader usage across public and private APIs. * Refine ML database corruption handling and sentinel logic - Improve logging for deferred quarantine attempts during sentinel creation failures (e.g., permission errors, read-only mounts). - Add safety check for `ml_db` null state to prevent redundant operations and log spam on initialization failures. - Update sentinel and database handling workflows for enhanced robustness against edge cases. * Refactor inline variable declarations in sentinel and corruption handling logic * Refactor and centralize model-table statement handling - Introduce `execute_and_reset_model_stmt()` for consistent execution and reset of model-table statements. - Add `handle_model_bind_fail()` for centralized bind-failure handling with improved logging. - Replace duplicate logic in "store model," "delete models," and "prune old models" operations with these helpers. * Introduce `ml_db_force_unusable()` for better handling of poisoned databases - Add `ml_db_force_unusable()` to mark the database unusable without disk writes, enabling safe bypass of poisoned databases at startup. - Improve sentinel and transaction error handling to avoid retriggering corruption scenarios and enhance error recovery mechanisms. - Refine SQLite statement execution and reset logic for consistent corruption tracking. (cherry picked from commit 69520c1)

Summary
ml_db_unusableflag to prevent further operations on a corrupted database.ml_db_mark_corrupt()to handle SQLite corruption errors (SQLITE_CORRUPT,SQLITE_NOTADB).Summary by cubic
Adds robust corruption detection and safe recovery for the ML SQLite database so the agent keeps running and retrains instead of crashing. Startup now consumes a sentinel, quarantines
ml.dbtoml.db.bad.<usec>, cleans WAL/SHM, and avoids re-quarantine loops by marking the DB unusable in-memory when needed.New Features
ml_db_mark_if_corrupt(rc)handles primary and extended SQLite errors; all open/migrate/prepare/step/reset/finalize and transaction paths use it and gate onml_db_is_unusable().ml.db.bad.<usec>, clean WAL/SHM, restore sentinel withO_CREAT|O_EXCLon rename failure;ml_db_force_unusable()skips opening poisoned DBs when the sentinel can’t be removed.SQLITE_DONEstep to discard partial results.Refactors
execute_and_reset_model_stmt()andhandle_model_bind_fail().ml_db == NULLor unusable, pending work is cleared and rollback/vacuum are skipped.snprintfz()sizes withsizeofand improved logs for deferred quarantine and retry paths.Written for commit 9023c68. Summary will update on new commits. Review in cubic