From 603e7bde8b3f4b43842f8f07e5154b0e61e8456f Mon Sep 17 00:00:00 2001 From: Tienson Qin Date: Thu, 8 Oct 2026 01:02:22 -0700 Subject: [PATCH] perf(storage): batch index node writes into one storage_store call per flush cljs buffers every node write in *store-buffer* and commits a single -store batch per tx; the restoring node storages called storage_store once per stored node instead. Graph_store maps each call to one sqlite transaction, so each ~300B node row paid a full commit (~4KB+ WAL), producing ~50KB/block of write amplification in the db-worker. Wrap the backend storage in batch_node_writes: node-only batches accumulate in a pending buffer that flushes together with the first batch carrying the root/tail markers store_to_storage/store_tail append at the end of every write batch. The wrapped record is bound to the conn/db at every entry point (store, restore, restore_root_snapshot, empty_db/init_db/create_conn/conn_from_db) so the index node storages, db.storage_ref and conn.storage all share the same buffer; store now prefers the bound storage_ref over a passed-in storage, matching cljs's single-storage-per-conn invariant, and conn.storage derives from it. Measured on a restored conn (200 txs x 5 blocks, bf=32): storage_store calls drop 53,590 -> 200 (one per tx) and the sqlite file 58MB -> 6.4MB, matching the cljs probe's call count and live row count. --- bench/dune | 6 ++ bench/restore_probe.ml | 134 +++++++++++++++++++++++++++++++++++++++++ impl/conn.ml | 10 ++- impl/datascript.ml | 17 ++++-- impl/storage.ml | 68 +++++++++++++++++++-- impl/storage.mli | 1 + 6 files changed, 224 insertions(+), 12 deletions(-) create mode 100644 bench/restore_probe.ml diff --git a/bench/dune b/bench/dune index 1a95efd..1647e5a 100644 --- a/bench/dune +++ b/bench/dune @@ -88,3 +88,9 @@ (modules entity_view_bench) (modes exe) (libraries datascript-ocaml-native unix)) + +(executable + (name restore_probe) + (modules restore_probe) + (modes exe) + (libraries datascript-ocaml-native datascript_sqlite unix sqlite3)) diff --git a/bench/restore_probe.ml b/bench/restore_probe.ml new file mode 100644 index 0000000..2a0bfdb --- /dev/null +++ b/bench/restore_probe.ml @@ -0,0 +1,134 @@ +open Datascript + +let indexed = + { cardinality = One + ; unique = None + ; indexed = true + ; is_component = false + ; no_history = false + ; doc = None + ; value_type = None + ; tuple_attrs = None + ; tuple_types = None + } + +let unique_identity = { indexed with unique = Some Identity } + +let schema = + [ "block/id", unique_identity + ; "block/journal-day", indexed + ; "block/content", indexed + ; "block/order", indexed + ; "block/collapsed", indexed + ] + +let block_tx base count = + List.init count (fun index -> + let i = base + index in + Entity + { db_id = Some (Temp_id (Printf.sprintf "b-%d" i)) + ; attrs = + [ "block/id", One_value (String (Printf.sprintf "block-%07d" i)) + ; "block/journal-day", One_value (String "2026-06-27") + ; "block/content", One_value (String (Printf.sprintf "Block content %07d lorem ipsum dolor" i)) + ; "block/order", One_value (Float (Float.of_int i)) + ; "block/collapsed", One_value (Bool false) + ] + }) + +let kvs_rows db = + let stmt = Sqlite3.prepare db "select count(*) from kvs" in + let n = + match Sqlite3.step stmt with + | Sqlite3.Rc.ROW -> Sqlite3.column_int stmt 0 + | rc -> failwith (Sqlite3.Rc.to_string rc) + in + ignore (Sqlite3.finalize stmt); + n + +let stores = ref 0 +let restores = ref 0 +let leaf_writes = ref 0 +let branch_writes = ref 0 +let node_write_bytes = ref 0 +let delete_calls = ref 0 +let call_sizes = ref [] + +let counting_storage inner = + { storage_store = + (fun entries -> + incr stores; + call_sizes := List.length entries :: !call_sizes; + List.iter + (fun (_addr, payload) -> + match payload with + | Storage_node (Persistent_sorted_set.Leaf _) -> + incr leaf_writes; + node_write_bytes := + !node_write_bytes + String.length (Datascript_sqlite_codec.encode payload) + | Storage_node (Persistent_sorted_set.Branch _) -> + incr branch_writes; + node_write_bytes := + !node_write_bytes + String.length (Datascript_sqlite_codec.encode payload) + | _ -> ()) + entries; + inner.storage_store entries) + ; storage_restore = + (fun addr -> + incr restores; + inner.storage_restore addr) + ; storage_list_addresses = inner.storage_list_addresses + ; storage_delete = + (fun addrs -> + incr delete_calls; + inner.storage_delete addrs) + } + +let () = + let seed = int_of_string Sys.argv.(1) in + let txs = int_of_string Sys.argv.(2) in + let per_tx = int_of_string Sys.argv.(3) in + let db_path = Filename.temp_file "restore-probe" ".sqlite3" in + let db = Sqlite3.db_open db_path in + (* phase 1: seed via a fresh conn *) + let session = Datascript_sqlite.open_session db_path in + let storage = counting_storage (Datascript_sqlite.storage session) in + let conn = create_conn ~schema ~storage () in + ignore (transact_conn conn (block_tx 0 seed)); + Printf.printf "seeded %d blocks: stores=%d leaf=%d branch=%d rows=%d\n%!" seed + !stores !leaf_writes !branch_writes (kvs_rows db); + Datascript_sqlite.close session; + (* phase 2: reopen via restore_conn — backed deferred indexes *) + stores := 0; leaf_writes := 0; branch_writes := 0; node_write_bytes := 0; + restores := 0; delete_calls := 0; + let session = Datascript_sqlite.open_session db_path in + let storage = counting_storage (Datascript_sqlite.storage session) in + let conn = + match restore_conn storage with + | Some conn -> conn + | None -> failwith "restore failed" + in + Printf.printf "restored conn\n%!"; + let rows0 = kvs_rows db in + for b = 1 to txs do + ignore (transact_conn conn (block_tx (1000000 + b * 1000) per_tx)) + done; + let rows1 = kvs_rows db in + let fsz p = if Sys.file_exists p then (Unix.stat p).st_size else 0 in + Printf.printf "file=%d wal=%d bytes\n%!" (fsz db_path) (fsz (db_path ^ "-wal")); + Printf.printf + "after %d txs x %d blocks: stores=%d leaf_writes=%d branch_writes=%d node_bytes=%d restores=%d deletes=%d rows %d -> %d (+%d)\n%!" + txs per_tx !stores !leaf_writes !branch_writes !node_write_bytes !restores + !delete_calls rows0 rows1 (rows1 - rows0); + let sizes = List.rev !call_sizes in + let show label xs = + Printf.printf "%s store-call sizes: %s\n%!" label + (String.concat "," (List.map string_of_int xs)) + in + let rec take n xs = match n, xs with 0, _ | _, [] -> [] | n, x :: r -> x :: take (n - 1) r in + let rec drop n xs = match n, xs with 0, _ -> xs | _, [] -> [] | n, _ :: r -> drop (n - 1) r in + show "first-100" (take 100 sizes); + show "last-100" (drop (List.length sizes - 100) sizes); + ignore (Sqlite3.db_close db); + Datascript_sqlite.close session; + Sys.remove db_path diff --git a/impl/conn.ml b/impl/conn.ml index 125fb81..d9bee52 100644 --- a/impl/conn.ml +++ b/impl/conn.ml @@ -81,9 +81,12 @@ let create (context : creation_context) ?schema ?storage () = match storage with | None -> db | Some storage -> - { (context.store ~storage db) with storage_ref = Some storage } + (* store binds db.storage_ref to the effective (wrapped) storage; + carry that into the conn so node writes and flush share its + pending buffer *) + context.store ~storage db in - make ?storage db + make ?storage:db.storage_ref db let from_db (context : creation_context) db = match db.storage_ref with @@ -137,7 +140,8 @@ let reset_schema (context : schema_context) conn schema = let restore (context : restore_context) storage = match context.restore storage with | None -> None - | Some db -> Some (make ~storage ~storage_tail:(context.restore_tail_groups storage) db) + | Some db -> + Some (make ?storage:db.storage_ref ~storage_tail:(context.restore_tail_groups storage) db) let transact (context : transact_context) ?(tx_meta = []) conn tx_data = let skip_store = tx_meta_skips_store tx_meta in diff --git a/impl/datascript.ml b/impl/datascript.ml index f2b522a..1cd9aec 100644 --- a/impl/datascript.ml +++ b/impl/datascript.ml @@ -85,12 +85,14 @@ let refresh_db_indexes_with_tx_data db tx_data = db let empty_db ?(schema = []) ?storage () = - Db_impl.empty_db db_core_context ~schema ?storage () + Db_impl.empty_db db_core_context ~schema + ?storage:(Option.map Storage.batch_node_writes storage) () let empty db = Db_impl.empty db_core_context db let init_db ?(schema = []) ?storage datoms = - Db_impl.init_db db_core_context ~schema ?storage datoms + Db_impl.init_db db_core_context ~schema + ?storage:(Option.map Storage.batch_node_writes storage) datoms let is_filtered = Db_impl.is_filtered @@ -130,13 +132,18 @@ let conn_creation_context : Conn.creation_context = { empty_db; init_db; store } let create_conn ?schema ?storage () = - Conn.create conn_creation_context ?schema ?storage () + Conn.create conn_creation_context ?schema + ?storage:(Option.map Storage.batch_node_writes storage) () let conn_from_db db = - Conn.from_db conn_creation_context db + Conn.from_db conn_creation_context + { db with + storage_ref = Option.map Storage.batch_node_writes db.storage_ref + } let conn_from_datoms ?schema ?storage datoms = - Conn.from_datoms conn_creation_context ?schema ?storage datoms + Conn.from_datoms conn_creation_context ?schema + ?storage:(Option.map Storage.batch_node_writes storage) datoms let conn_db = Conn.db diff --git a/impl/storage.ml b/impl/storage.ml index 8d93bbc..a463053 100644 --- a/impl/storage.ml +++ b/impl/storage.ml @@ -223,6 +223,57 @@ let store_index node_storage index index_set = | exception Invalid_argument message when String.equal message "store requires a storage-backed set" -> fst (PSet.store (storage_backed_index node_storage index index_set)) +(* cljs store-impl! buffers every node write in *store-buffer* and commits a + single -store batch per tx. The index node storages above used to call + storage_store once per stored node — Graph_store maps each call to one + sqlite transaction, so a ~300B node row paid a full commit (~4KB+ WAL) + instead of ~300B inside a batch. + + Wrap the backend storage so node-only batches accumulate in pending and + flush together with the first batch carrying the root/tail — the markers + store_to_storage/store_tail append at the end of every write batch. The + wrapped record must be the one bound to the conn/db and to the index node + storages, so all per-node writes within one conn lifetime share the same + pending buffer; bind it at every entry point below (store, restore, + restore_root_snapshot, create_conn via empty_db/init_db). Wrapping an + already wrapped storage stays correct: the inner buffer still flushes on + the outer flush call. *) +let batch_node_writes (storage : storage) : storage = + let pending = ref [] in + { storage with + storage_store = + (fun entries -> + let nodes, rest = + List.partition + (fun (_, payload) -> match payload with Storage_node _ -> true | _ -> false) + entries + in + pending := List.rev_append nodes !pending; + match rest with + | [] -> () + | _ -> + let batch = List.rev !pending @ rest in + pending := []; + storage.storage_store batch) + ; storage_restore = + (fun address -> + match + List.find_map + (fun (addr, payload) -> + if String.equal addr address then Some payload else None) + !pending + with + | Some payload -> Some payload + | None -> storage.storage_restore address) + ; storage_delete = + (fun addresses -> + pending := + List.filter + (fun (addr, _) -> not (List.mem addr addresses)) + !pending; + storage.storage_delete addresses) + } + let store_to_storage db storage = let pending_entries = ref [] in let node_storage = buffered_node_storage pending_entries in @@ -265,10 +316,17 @@ let store_to_storage db storage = } let store ?storage db = - match storage, db.storage_ref with - | Some storage, _ -> store_to_storage db storage - | None, Some storage -> store_to_storage db storage - | None, None -> invalid_arg "db has no attached storage" + (* The bound storage_ref always wins: it carries the batch buffer the + index node storages write into. cljs throws when asked to store a db + to a different IStorage than it was created with; preferring the + bound one keeps that single-storage invariant. *) + let storage = + match db.storage_ref, storage with + | Some bound, _ -> bound + | None, Some storage -> batch_node_writes storage + | None, None -> invalid_arg "db has no attached storage" + in + { (store_to_storage db storage) with storage_ref = Some storage } let store_tail storage tail = storage.storage_store [ tail_address, Storage_tail tail ] @@ -283,6 +341,7 @@ let tail_datom_count tail = List.fold_left (fun count group -> count + List.length group) 0 tail let restore_root_snapshot storage = + let storage = batch_node_writes storage in match storage.storage_restore root_address with | Some (Storage_root root) -> note_storage_root root; @@ -328,6 +387,7 @@ let db_with_tail context db tail = tail let restore context storage = + let storage = batch_node_writes storage in match storage.storage_restore root_address with | None -> None | Some (Storage_root root) -> diff --git a/impl/storage.mli b/impl/storage.mli index afc094e..d065643 100644 --- a/impl/storage.mli +++ b/impl/storage.mli @@ -13,6 +13,7 @@ val root_address : storage_address val tail_address : storage_address val memory_storage : unit -> storage val file_storage : string -> storage +val batch_node_writes : storage -> storage val store : ?storage:storage -> db -> db val store_tail : storage -> datom list list -> unit val normalize_stored_datom : schema -> datom -> datom