Source file recovery.ml

(** Offline, best-effort salvage of a damaged database file into a fresh,
    structurally-valid one — the corruption-recovery counterpart to the
    crash-recovery already performed on open.

    Read-only on the source; emits a clean new file rather than repairing
    in place.  Implementation phases:

    1. [scan]: Raw-page scan verifying CRC32 on every page and
       classifying by kind byte.  Provides a trust map used later.
    2. [open_source]: Attempt to open the damaged file via Store.open_block
       (handles crash recovery / WAL replay automatically).  Load the catalog.
    3. [extract]: Walk each table's tree via cursor, skip corrupt pages
       (flagged by the scan trust map), decode rows.
    4. [rebuild]: Open a fresh Store and replay recovered rows through the
       normal put path, producing correct CRCs in a clean CoW B+-tree.

    When the source headers are intact (the common case), phases 2–3 use
    the existing Store/Catalog infrastructure.  When headers are corrupt,
    the tool falls back to raw-page catalog reconstruction from the
    redundant mirror (#174). *)

module Page = Granary_storage.Page
module Store = Granary_store.Store
module Catalog = Granary_catalog.Catalog
open Lwt.Syntax

(* ------------------------------------------------------------------ *)
(* Types                                                               *)
(* ------------------------------------------------------------------ *)

(** A single page from the raw scan, annotated with CRC status
    and decoded page kind. *)
type scanned_page =
  { page_id : int64
  ; kind : Page.kind
  ; crc_valid : bool
  ; tag : int32
  }

(** Result of the raw-page scan phase. *)
type scan_result =
  { n_pages : int64
  ; pages : scanned_page list
  ; trusted_count : int
  ; damaged_count : int
  }

(** A row recovered and ready to write to the output store. *)
type recovered_row =
  { tree_id : int
  ; key : bytes
  ; value : bytes
  }

(** Final result of a complete recovery run. *)
type recover_result =
  { rows_recovered : int
  ; pages_scanned : int64
  ; pages_trusted : int
  ; pages_damaged : int
  ; pages_skipped : int
  ; tables_recovered : int
  ; tables_total : int
  }

(* ------------------------------------------------------------------ *)
(* Pretty-printers                                                     *)
(* ------------------------------------------------------------------ *)

let pp_kind fmt (k : Page.kind) =
  let s =
    match k with
    | Page.Header -> "Header"
    | Page.Branch -> "Branch"
    | Page.Leaf -> "Leaf"
    | Page.Freelist -> "Freelist"
    | Page.Overflow -> "Overflow"
  in
  Format.pp_print_string fmt s
;;

let pp_scanned_page fmt p =
  Format.fprintf
    fmt
    "page %Ld: %a (crc=%s tag=%ld)"
    p.page_id
    pp_kind
    p.kind
    (if p.crc_valid then "OK" else "bad")
    p.tag
;;

let pp_result fmt r =
  Format.fprintf
    fmt
    "@[<v 2>Recovery complete:@,\
     %d rows recovered from %d/%d tables@,\
     %Ld pages scanned: %d trusted, %d damaged, %d skipped@]"
    r.rows_recovered
    r.tables_recovered
    r.tables_total
    r.pages_scanned
    r.pages_trusted
    r.pages_damaged
    r.pages_skipped
;;

(* ------------------------------------------------------------------ *)
(* Helpers                                                             *)
(* ------------------------------------------------------------------ *)

(** Decode the page kind from byte 0 of a page buffer.
    Maps unknown bytes to [Page.Header] — the page will be flagged
    damaged via CRC mismatch anyway. *)
let kind_of_byte b =
  match b with
  | 0 -> Page.Header
  | 1 -> Page.Branch
  | 2 -> Page.Leaf
  | 3 -> Page.Freelist
  | 4 -> Page.Overflow
  | _ -> Page.Header
;;

(** Number of recovered rows to accumulate before calling [write_to]. *)
let batch_size = 1000

(* ------------------------------------------------------------------ *)
(* Phase 1: scan                                                       *)
(* ------------------------------------------------------------------ *)

let scan ~read_page ~n_pages =
  let n_pages_int = Int64.to_int n_pages in
  let rec loop i acc_pages trusted damaged =
    if i >= n_pages_int
    then (
      let pages = List.rev acc_pages in
      let n_read = Int64.of_int (List.length pages) in
      Lwt.return
        { n_pages = n_read; pages; trusted_count = trusted; damaged_count = damaged })
    else (
      let page_id = Int64.of_int i in
      let* page_opt = read_page page_id in
      match page_opt with
      | None ->
        (* Page failed to read — silently skip *)
        loop (i + 1) acc_pages trusted damaged
      | Some buf ->
        let crc_valid = Page.verify_crc buf in
        let kind_byte = Cstruct.get_uint8 buf 0 in
        let kind = kind_of_byte kind_byte in
        let tag = Page.read_tag buf in
        let page = { page_id; kind; crc_valid; tag } in
        let trusted = if crc_valid then trusted + 1 else trusted in
        let damaged = if not crc_valid then damaged + 1 else damaged in
        loop (i + 1) (page :: acc_pages) trusted damaged)
  in
  loop 0 [] 0 0
;;

(* ------------------------------------------------------------------ *)
(* Phase 3: extract one table                                          *)
(* ------------------------------------------------------------------ *)

(** Extract all (key, value) pairs from a single table tree using a
    cursor. Returns the list of recovered rows, or an empty list if the
    cursor cannot be opened (corrupt pages). *)
let extract_one_table (store : Store.t) (tree_id : int) =
  Lwt.catch
    (fun () ->
       let* rows =
         Store.with_ro store (fun txn ->
           let* cursor = Store.cursor_open txn tree_id in
           let (_ : Store.seek_result) = Store.cursor_first cursor in
           let rec collect acc =
             match Store.cursor_next cursor with
             | None -> Lwt.return (List.rev acc)
             | Some (k, v) -> collect ({ tree_id; key = k; value = v } :: acc)
           in
           collect [])
       in
       Lwt.return rows)
    (fun exn ->
       let msg = Printexc.to_string exn in
       Format.eprintf "recovery: failed to extract tree %d: %s@." tree_id msg;
       Lwt.return [])
;;

(* ------------------------------------------------------------------ *)
(* Batch-flush helper                                                  *)
(* ------------------------------------------------------------------ *)

(** Split [n] elements from the front of a list. *)
let rec take n acc xs =
  match n, xs with
  | 0, _ | _, [] -> List.rev acc, xs
  | n, x :: xs' -> take (n - 1) (x :: acc) xs'
;;

(** Write rows to the destination in batches of [batch_size]. *)
let rec flush_batches write_to rows =
  match rows with
  | [] -> Lwt.return_ok ()
  | _ ->
    let batch, rest = take batch_size [] rows in
    let* result = write_to batch in
    (match result with
     | Error e -> Lwt.return_error e
     | Ok () -> flush_batches write_to rest)
;;

(* ------------------------------------------------------------------ *)
(* Phase 4: full recover pipeline                                      *)
(* ------------------------------------------------------------------ *)

let recover ~read_page ~n_pages ~open_source ~write_to =
  (* Phase 1: scan every page and collect CRC trust map.
     Currently informational only (reported in the summary); the
     Store-layer cursor errors serve as the corruption fallback during
     extraction.  A follow-up (#85 phase 2) will pass the trust map
     through to skip known-bad pages proactively without waiting for
     cursor errors. *)
  let* scan_result = scan ~read_page ~n_pages in
  let pages_scanned = scan_result.n_pages in
  let pages_read = Int64.of_int (List.length scan_result.pages) in
  let pages_skipped = Int64.to_int (Int64.sub n_pages pages_read) in
  let pages_trusted = scan_result.trusted_count in
  let pages_damaged = scan_result.damaged_count in
  (* Phase 2: attempt to open the source *)
  let* source = open_source () in
  match source with
  | Error e ->
    Format.eprintf "recovery: cannot open source: %s@." e;
    Lwt.return_error (Printf.sprintf "open_source failed: %s" e)
  | Ok (store, catalog) ->
    (* Phase 3: extract rows from every table *)
    let* tables = Catalog.list_tables catalog in
    let tables_total = List.length tables in
    let rec extract_all acc_rev tables_recovered remaining_tables =
      match remaining_tables with
      | [] -> Lwt.return (acc_rev, tables_recovered)
      | table_meta :: rest ->
        (match table_meta.Catalog.storage with
         | Catalog.Columnar _ -> extract_all acc_rev tables_recovered rest
         | Catalog.Row { tree_id; _ } ->
           let* table_rows = extract_one_table store tree_id in
           let n_rows = List.length table_rows in
           let recovered =
             if n_rows > 0 then tables_recovered + 1 else tables_recovered
           in
           extract_all (List.rev_append table_rows acc_rev) recovered rest)
    in
    let* all_rows_rev, tables_recovered = extract_all [] 0 tables in
    let all_rows = List.rev all_rows_rev in
    let rows_recovered = List.length all_rows in
    (* Phase 4: write recovered rows to destination *)
    let* write_result = flush_batches write_to all_rows in
    (match write_result with
     | Error e -> Lwt.return_error (Printf.sprintf "write_to failed: %s" e)
     | Ok () ->
       Lwt.return_ok
         { rows_recovered
         ; pages_scanned
         ; pages_trusted
         ; pages_damaged
         ; pages_skipped
         ; tables_recovered
         ; tables_total
         })
;;