From 519891c06eca554e9a44984b5c5ffa30bf35794d Mon Sep 17 00:00:00 2001 From: rcmerci Date: Thu, 1 Oct 2026 15:29:43 +0800 Subject: [PATCH 1/2] fix(storage): restore index counts from snapshot metadata --- datascript-ocaml-melange.opam | 2 +- datascript-ocaml-native.opam | 2 +- datascript_ocaml.opam | 2 +- impl/storage.ml | 48 +++++++------ test/js_smoke.ml | 16 +++++ test/melange_smoke.ml | 16 +++++ test/test_storage.ml | 124 ++++++++++++++++++++++++++++++++++ 7 files changed, 188 insertions(+), 22 deletions(-) diff --git a/datascript-ocaml-melange.opam b/datascript-ocaml-melange.opam index de6eeb1..3ffd545 100644 --- a/datascript-ocaml-melange.opam +++ b/datascript-ocaml-melange.opam @@ -13,7 +13,7 @@ depends: [ "melange-transit-melange" {= "0.1.2"} ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] ["melange-edn-core.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-edn-melange.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-transit-core.0.1.2" "git+https://github.com/logseq/melange-transit.git#main"] diff --git a/datascript-ocaml-native.opam b/datascript-ocaml-native.opam index 9c32710..458c6f2 100644 --- a/datascript-ocaml-native.opam +++ b/datascript-ocaml-native.opam @@ -14,7 +14,7 @@ depends: [ "yojson" ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] ["melange-edn-core.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-edn-native.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-transit-core.0.1.2" "git+https://github.com/logseq/melange-transit.git#main"] diff --git a/datascript_ocaml.opam b/datascript_ocaml.opam index 72cf185..7380e1e 100644 --- a/datascript_ocaml.opam +++ b/datascript_ocaml.opam @@ -11,7 +11,7 @@ depends: [ "melange" ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] ] build: [ ["dune" "build" "-p" name "-j" jobs] diff --git a/impl/storage.ml b/impl/storage.ml index 34ce6b9..96f6ec5 100644 --- a/impl/storage.ml +++ b/impl/storage.ml @@ -298,15 +298,17 @@ let verify_stored_index node_storage cmp address = let datoms = ref [] in let last = ref None in let ordered = ref true in + let count = ref 0 in fold_stored_index node_storage address (fun () datom -> + incr count; datoms := datom :: !datoms; (match !ordered, !last with | true, Some prev when cmp prev datom > 0 -> ordered := false | _ -> ()); last := Some datom) (); - if !ordered then `Clean + if !ordered then `Clean !count else let datoms = Array.of_list (List.rev !datoms) in Array.sort cmp datoms; @@ -320,22 +322,30 @@ let restore context storage = let schema = Schema.validate_schema root.storage_schema in let settings = settings_of_root root in let node_storage = restoring_node_storage ~schema storage in - let restore_index index address = + let restore_index index address metadata = let cmp = Util.compare_datom index in - match - if root.storage_index_order_version <> index_order_version - then verify_stored_index node_storage cmp address - else `Clean - with - | `Dirty datoms -> - (* The stored tree routes seeks by the old ordering: rebuild - the index in memory in the correct order. The next store - rewrites its nodes, healing storage. *) - PSet.of_sorted_array_by ~settings ~cmp datoms - | `Clean -> - (match PSet.restore ~cmp ~settings node_storage address with - | Some index -> index - | None -> invalid_arg ("storage root points at a missing index: " ^ address)) + let restore_lazy ?count () = + match PSet.restore ?count ~cmp ~settings node_storage address with + | Some index -> index + | None -> invalid_arg ("storage root points at a missing index: " ^ address) + in + if root.storage_index_order_version = index_order_version then + (* Counts and index addresses come from the same stored root snapshot. + Seed before replaying its separate tail, whose edits update counts. + Older roots without metadata retain PSS's lazy counting fallback. *) + let count = Option.map (fun metadata -> metadata.storage_index_count) metadata in + restore_lazy ?count () + else + match verify_stored_index node_storage cmp address with + | `Clean count -> + (* Verification already visited every datom. Use that measured + count rather than trusting metadata from an older format. *) + restore_lazy ~count () + | `Dirty datoms -> + (* The stored tree routes seeks by the old ordering: rebuild + the index in memory in the correct order. The next store + rewrites its nodes, healing storage. *) + PSet.of_sorted_array_by ~settings ~cmp datoms in let duplicate_datoms = normalize_stored_datoms schema root.storage_duplicate_datoms in let duplicate_eavt_by_entity = @@ -364,12 +374,12 @@ let restore context storage = |> List.filter (fun datom -> Schema.schema_attr_is_avet_accessible schema datom.a) |> List.sort (Util.compare_datom Avet) in - let aevt_index = restore_index Aevt root.storage_aevt in - let avet_index = restore_index Avet root.storage_avet in + let aevt_index = restore_index Aevt root.storage_aevt root.storage_aevt_metadata in + let avet_index = restore_index Avet root.storage_avet root.storage_avet_metadata in let db = { db_uid = context.next_db_uid () ; schema - ; eavt_index = restore_index Eavt root.storage_eavt + ; eavt_index = restore_index Eavt root.storage_eavt root.storage_eavt_metadata ; aevt_index ; avet_index ; aevt_by_attr = Hashtbl.create 0 diff --git a/test/js_smoke.ml b/test/js_smoke.ml index 2b00151..d675a8e 100644 --- a/test/js_smoke.ml +++ b/test/js_smoke.ml @@ -32,3 +32,19 @@ let () = match datoms db Eavt ~e:1 ~a:"touched" () with | [ { v = Bool true; _ } ] -> () | _ -> failwith "JavaScript smoke transaction function returned an unexpected result" + +let () = + let base = memory_storage () in + let reads = ref 0 in + let storage = { base with storage_restore = (fun address -> + incr reads; base.storage_restore address) } in + let schema = schema_of_edn_string "{:name {:db/index true}}" in + let db = init_db ~schema [datom ~e:1 ~a:"name" ~v:(String "v") ()] in + store ~storage db; + let restored = Option.get (restore storage) in + reads := 0; + List.iter (fun set -> + if Persistent_sorted_set.count set <> 1 || Persistent_sorted_set.count set <> 1 then + failwith "restored smoke count differs") + [restored.eavt_index;restored.aevt_index;restored.avet_index]; + if !reads <> 0 then failwith "restored smoke count loaded storage" diff --git a/test/melange_smoke.ml b/test/melange_smoke.ml index 7352f22..29024ab 100644 --- a/test/melange_smoke.ml +++ b/test/melange_smoke.ml @@ -25,3 +25,19 @@ let () = "Melange regex seq" [ "123"; "456" ] (Built_ins.regex_seq "[0-9]+" "a123b456") + +let () = + let base = memory_storage () in + let reads = ref 0 in + let storage = { base with storage_restore = (fun address -> + incr reads; base.storage_restore address) } in + let schema = schema_of_edn_string "{:name {:db/index true}}" in + let db = init_db ~schema [datom ~e:1 ~a:"name" ~v:(String "v") ()] in + store ~storage db; + let restored = Option.get (restore storage) in + reads := 0; + List.iter (fun set -> + if Persistent_sorted_set.count set <> 1 || Persistent_sorted_set.count set <> 1 then + failwith "restored smoke count differs") + [restored.eavt_index;restored.aevt_index;restored.avet_index]; + if !reads <> 0 then failwith "restored smoke count loaded storage" diff --git a/test/test_storage.ml b/test/test_storage.ml index ab6e6d2..5abbbdd 100644 --- a/test/test_storage.ml +++ b/test/test_storage.ml @@ -462,7 +462,131 @@ let test_storage__test_restored_misordered_index () = [ 1, "a/b", String "one"; 1, "a/c", String "two" ] (datoms restored Eavt ~e:1 ())) +module PSet = Persistent_sorted_set + +let count_fixture () = + let storage = memory_storage () in + let schema = [ "probe/indexed", indexed; "probe/plain", { indexed with indexed = false } ] in + let datoms = List.init 4096 (fun i -> + datom ~e:(i / 2 + 1) + ~a:(if i mod 2 = 0 then "probe/indexed" else "probe/plain") + ~v:(String (Printf.sprintf "v%04d" i)) ()) in + store ~storage (init_db ~schema datoms); + storage + +let index_counts db = + List.map PSet.count [ db.eavt_index; db.aevt_index; db.avet_index ] + +let assert_index_counts label expected db = + let actual = index_counts db in + if actual <> expected then + failf "%s: expected [%s], got [%s]" label + (String.concat ";" (List.map string_of_int expected)) + (String.concat ";" (List.map string_of_int actual)) + +let check_cached_counts label expected db reads = + reads := []; + assert_index_counts label expected db; + assert_index_counts (label ^ " repeated") expected db; + assert_equal_int (label ^ " does not read storage") 0 (List.length !reads) + +let map_storage_root storage f = + match storage.storage_restore "0" with + | Some (Storage_root root) -> storage.storage_store [ "0", Storage_root (f root) ] + | _ -> failwith "expected stored root" + +let test_storage__test_restore_keeps_snapshot_counts () = + let storage = count_fixture () in + (match storage.storage_restore "0" with + | Some (Storage_root root) -> + let counts = List.map (fun metadata -> (Option.get metadata).storage_index_count) + [ root.storage_eavt_metadata; root.storage_aevt_metadata; root.storage_avet_metadata ] in + if counts <> [4096;4096;2048] then failwith "root metadata count differs" + | _ -> failwith "expected root metadata"); + let measured, reads = restore_counting_storage storage in + let restored = Option.get (restore measured) in + assert_equal_int "restore reads only root and tail" 2 (List.length !reads); + check_cached_counts "cold snapshot counts" [4096;4096;2048] restored reads; + let changed = db_with [ Add (Entity_id 2049, "probe/indexed", String "new") ] restored in + check_cached_counts "successful add count" [4097;4097;2049] changed reads; + let unchanged = db_with + [ Add (Entity_id 2049, "probe/indexed", String "new") + ; Retract (Entity_id 9999, "probe/indexed", Some (String "missing")) ] changed in + check_cached_counts "no-op counts" [4097;4097;2049] unchanged reads; + let removed = db_with + [ Retract (Entity_id 1, "probe/indexed", Some (String "v0000")) ] unchanged in + check_cached_counts "successful removal count" [4096;4096;2048] removed reads; + check_cached_counts "persistent original count" [4096;4096;2048] restored reads + +let test_storage__test_restore_counts_after_tail () = + let storage = count_fixture () in + store_tail storage + [ [ datom ~e:2049 ~a:"probe/indexed" ~v:(String "new") ~tx:(tx0 + 1) () ] + ; [ datom ~e:1 ~a:"probe/indexed" ~v:(String "v0000") ~added:false ~tx:(tx0 + 2) () ] + ]; + let measured, reads = restore_counting_storage storage in + let restored = Option.get (restore measured) in + assert_int_at_most "tail replay does not scan all index nodes" 60 (List.length !reads); + check_cached_counts "tail applies to snapshot count exactly once" [4096;4096;2048] restored reads; + List.iter2 (fun index count -> + assert_equal_int "tail cached count agrees with datoms" count + (List.length (datoms restored index ()))) + [Eavt;Aevt;Avet] [4096;4096;2048]; + let conn = Option.get (restore_conn measured) in + check_cached_counts "restore_conn counts after tail" [4096;4096;2048] (conn_db conn) reads; + let tx = transact_conn conn [ Add (Entity_id 2050, "probe/indexed", String "later") ] in + check_cached_counts "incremental conn transaction count" [4097;4097;2049] tx.db_after reads; + store tx.db_after; + let restored_again = Option.get (restore measured) in + check_cached_counts "new snapshot count excludes cleared tail" [4097;4097;2049] restored_again reads + +let test_storage__test_empty_and_missing_count_metadata () = + let storage = memory_storage () in + store ~storage (empty_db ()); + let measured, reads = restore_counting_storage storage in + let empty = Option.get (restore measured) in + check_cached_counts "empty snapshot count" [0;0;0] empty reads; + let changed = db_with [ Add (Entity_id 1, "plain", String "v") ] empty in + check_cached_counts "empty snapshot edited count" [1;1;0] changed reads; + let storage = count_fixture () in + map_storage_root storage (fun root -> {root with + storage_eavt_metadata = None; storage_aevt_metadata = None; storage_avet_metadata = None}); + let measured, reads = restore_counting_storage storage in + let restored = Option.get (restore measured) in + assert_equal_int "missing metadata restore stays lazy" 2 (List.length !reads); + reads := []; + assert_index_counts "missing metadata falls back to actual count" [4096;4096;2048] restored; + if !reads = [] then failwith "missing metadata should load nodes for count"; + let storage = count_fixture () in + map_storage_root storage (fun root -> {root with + storage_eavt_metadata = Some {storage_index_count = -1; storage_index_shift = 2}}); + (match restore storage with + | exception Invalid_argument _ -> () + | _ -> failwith "negative current snapshot count must be rejected") + +let test_storage__test_legacy_verification_ignores_stale_counts () = + let storage = memory_storage () in + let d1 = datom ~e:1 ~a:"name" ~v:(String "one") () in + let d2 = datom ~e:2 ~a:"name" ~v:(String "two") () in + let stale = Some { storage_index_count = 999; storage_index_shift = 0 } in + let root = { (legacy_storage_root ~version:0 ~eavt:"20" ~aevt:"20" ~avet:"20") with + storage_eavt_metadata = stale; storage_aevt_metadata = stale; storage_avet_metadata = stale } in + storage.storage_store ["20", Storage_node (PSet.Leaf [d1;d2]); "0", Storage_root root]; + let measured, reads = restore_counting_storage storage in + let restored = Option.get (restore measured) in + check_cached_counts "verified legacy count replaces stale metadata" [2;2;2] restored reads; + let d3 = datom ~e:1 ~a:"a.x/b" ~v:(String "B") () in + let d4 = datom ~e:1 ~a:"a/c" ~v:(String "C") () in + storage.storage_store ["20", Storage_node (PSet.Leaf [d3;d4])]; + let rebuilt = Option.get (restore measured) in + check_cached_counts "rebuilt index uses actual count" [2;2;2] rebuilt reads; + assert_equal_int "rebuilt count matches actual contents" 2 (List.length (datoms rebuilt Eavt ())) + let () = + test_storage__test_restore_keeps_snapshot_counts (); + test_storage__test_restore_counts_after_tail (); + test_storage__test_empty_and_missing_count_metadata (); + test_storage__test_legacy_verification_ignores_stale_counts (); test_storage__test_basics (); test_storage__test_upstream_wire_addresses (); test_storage__test_file_storage (); From 4fc92bbecc0b34b989ba2c82cfed53a32d268439 Mon Sep 17 00:00:00 2001 From: rcmerci Date: Thu, 1 Oct 2026 20:08:01 +0800 Subject: [PATCH 2/2] build: restore PSS main dependency after upstream merge --- datascript-ocaml-melange.opam | 2 +- datascript-ocaml-native.opam | 2 +- datascript_ocaml.opam | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datascript-ocaml-melange.opam b/datascript-ocaml-melange.opam index 3ffd545..de6eeb1 100644 --- a/datascript-ocaml-melange.opam +++ b/datascript-ocaml-melange.opam @@ -13,7 +13,7 @@ depends: [ "melange-transit-melange" {= "0.1.2"} ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] ["melange-edn-core.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-edn-melange.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-transit-core.0.1.2" "git+https://github.com/logseq/melange-transit.git#main"] diff --git a/datascript-ocaml-native.opam b/datascript-ocaml-native.opam index 458c6f2..9c32710 100644 --- a/datascript-ocaml-native.opam +++ b/datascript-ocaml-native.opam @@ -14,7 +14,7 @@ depends: [ "yojson" ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] ["melange-edn-core.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-edn-native.0.5.0" "git+https://github.com/logseq/melange-edn.git#main"] ["melange-transit-core.0.1.2" "git+https://github.com/logseq/melange-transit.git#main"] diff --git a/datascript_ocaml.opam b/datascript_ocaml.opam index 7380e1e..72cf185 100644 --- a/datascript_ocaml.opam +++ b/datascript_ocaml.opam @@ -11,7 +11,7 @@ depends: [ "melange" ] pin-depends: [ - ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#9c99ee9760978654540d0482cbbc3d8ed97aa9c2"] + ["persistent_sorted_set_ocaml.dev" "git+https://github.com/logseq/persistent-sorted-set-ocaml.git#main"] ] build: [ ["dune" "build" "-p" name "-j" jobs]