perf(storage): batch index node writes into one storage_store call per flush

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))

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

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

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
 

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) ->

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