wip

lib/datastore.ml

@@ -1,4 +1,13 @@
+open Lwt.Infix
+
 type connection = Sqlite3.db
+
+type connection_pool = {
+  mutex : Lwt_mutex.t;
+  condition : unit Lwt_condition.t;
+  mutable available : connection list;
+}
+
 type prepared_stmt = { sql : string; stmt : Sqlite3.stmt }
 type value = Int of int | Text of string | Null
 
@@ -13,24 +22,52 @@ let sqlite3_data_of_value = function
   | Text s -> Sqlite3.Data.TEXT s
   | Null -> Sqlite3.Data.NULL
 
-let in_memory_database = ":memory:"
-
 let execute_no_prepare con sql =
   let rc = Sqlite3.exec con sql in
   if Sqlite3.Rc.is_success rc then Ok () else Error (Sqlite3.Rc.to_string rc)
 
-let connect path =
-  let con = Sqlite3.db_open path in
+let open_db path =
+  let num_connections = 10 in
+  let mutex = Lwt_mutex.create () in
+  let condition = Lwt_condition.create () in
+  let available =
+    List.init num_connections @@ fun _ ->
+    let con = Sqlite3.db_open path in
+    execute_no_prepare con "pragma journal_mode = WAL" |> Result.get_ok;
+    execute_no_prepare con "pragma synchronous = normal" |> Result.get_ok;
+    execute_no_prepare con "pragma temp_store = memory" |> Result.get_ok;
+    execute_no_prepare con "pragma mmap_size = 30000000000" |> Result.get_ok;
+    con
+  in
+  { mutex; condition; available }
+
+let close_db _pool = (* FIXME *) Lwt.return_unit
+
+let acquire pool =
+  Lwt_mutex.with_lock pool.mutex @@ fun () ->
+  let rec loop () =
+    match pool.available with
+    | [] ->
+        Lwt_condition.wait pool.condition ~mutex:pool.mutex >>= fun () ->
+        loop ()
+    | con :: rest ->
+        pool.available <- rest;
+        Lwt.return con
+  in
+  loop ()
 
-  (* Thanks to: https://phiresky.github.io/blog/2020/sqlite-performance-tuning/ *)
-  execute_no_prepare con "pragma journal_mode = WAL" |> Result.get_ok;
-  execute_no_prepare con "pragma synchronous = normal" |> Result.get_ok;
-  execute_no_prepare con "pragma temp_store = memory" |> Result.get_ok;
-  execute_no_prepare con "pragma mmap_size = 30000000000" |> Result.get_ok;
+let release pool con =
+  Lwt_mutex.with_lock pool.mutex @@ fun () ->
+  pool.available <- con :: pool.available;
+  Lwt_condition.signal pool.condition ();
+  Lwt.return_unit
 
-  con
+let use pool f =
+  acquire pool >>= fun con ->
+  Lwt.finalize
+    (fun () -> Lwt_preemptive.detach (fun () -> f con) ())
+    (fun () -> release pool con)
 
-let disconnect con = Sqlite3.db_close con |> ignore
 let prepare conn sql = { sql; stmt = Sqlite3.prepare conn sql }
 
 exception Sqlite_error of string

lib/datastore.mli

@@ -1,10 +1,11 @@
 type connection
+type connection_pool
 type prepared_stmt
 type value = Int of int | Text of string | Null
 
-val in_memory_database : string
-val connect : string -> connection
-val disconnect : connection -> unit
+val open_db : string -> connection_pool
+val close_db : connection_pool -> unit Lwt.t
+val use : connection_pool -> (connection -> 'a) -> 'a Lwt.t
 val prepare : connection -> string -> prepared_stmt
 val execute : prepared_stmt -> value list -> (unit, string) result
 val query : prepared_stmt -> value list -> (value list list, string) result

lib/sql_writer.ml

@@ -1,3 +1,5 @@
+open Util
+
 let escape_single_quote s =
   let buf = Buffer.create (String.length s) in
   String.iter
@@ -313,9 +315,8 @@ let dump_transaction_tag_records (model : Model.t) tag_id_map =
          t.tags |> List.map (fun tag -> (i, tag_id_map |> StringMap.find tag)))
   |> List.flatten |> List.sort_uniq compare
 
-let dump path (model : Model.t) =
-  let con = Datastore.(connect path) in
-
+let dump' model pool =
+  Datastore.use pool @@ fun con ->
   Store.create_accounts con;
   Store.create_transactions con;
   Store.create_postings con;
@@ -344,11 +345,14 @@ let dump path (model : Model.t) =
 
   Store.create_full_accounts_view con;
 
-  (model.accounts
+  model.accounts
   |> List.iter @@ fun (acc : Model.open_account) ->
      Store.create_account_transactions_view con
-       ~account:(String.concat ":" acc.account));
+       ~account:(String.concat ":" acc.account)
 
-  con
+let dump path (model : Model.t) =
+  let pool = Datastore.open_db path in
+  Lwt_main.run (dump' model pool);
+  pool
 
-let dump_on_memory = dump Datastore.in_memory_database
+let dump_on_memory = with_temp_file dump

lib/util.ml

@@ -15,3 +15,7 @@ let iota n =
   aux [] n
 
 let failwithf fmt = Printf.ksprintf failwith fmt
+
+let with_temp_file f =
+  let filepath = Filename.temp_file "qash" "test" in
+  Fun.protect ~finally:(fun () -> Sys.remove filepath) (fun () -> f filepath)

lib/verifier.ml

@@ -24,7 +24,9 @@ HAVING sum(p.amount) <> 0
          | _ -> assert false)
 end
 
-let verify_balanced_transactions con =
+let verify_balanced_transactions pool =
+  Lwt_main.run @@ Datastore.use pool
+  @@ fun con ->
   let unbalanced = Store.select_unbalance_transactions con in
   if unbalanced <> [] then
     Error
@@ -35,7 +37,9 @@ let verify_balanced_transactions con =
       |> Printf.sprintf "Unbalanced transactions:\n%s")
   else Ok ()
 
-let verify_notes (con : Datastore.connection) notes =
+let verify_notes pool notes =
+  Lwt_main.run @@ Datastore.use pool
+  @@ fun con ->
   try
     notes
     |> List.iter (fun note ->
@@ -54,8 +58,8 @@ let verify_notes (con : Datastore.connection) notes =
     Ok ()
   with Failure s -> Error s
 
-let verify con notes =
+let verify pool notes =
   let ( let* ) = Result.bind in
-  let* () = verify_balanced_transactions con in
-  let* () = verify_notes con notes in
+  let* () = verify_balanced_transactions pool in
+  let* () = verify_notes pool notes in
   Ok ()

lib/web_server.ml

@@ -459,23 +459,21 @@ let generate =
         Lwt_mutex.with_lock mtx (fun () ->
             Loader.load_file in_filename |> Lwt.return)
       in
-      Lwt_preemptive.detach
-        (fun () ->
-          let con = Sql_writer.dump_on_memory m in
-          match Verifier.verify con notes with
-          | Error s -> failwithf "Verification error: %s" s
-          | Ok () ->
-              Fun.protect
-                (fun () -> thn con)
-                ~finally:(fun () -> Datastore.disconnect con))
-        ()
+      let pool = Sql_writer.dump_on_memory m in
+      match Verifier.verify pool notes with
+      | Error s -> failwithf "Verification error: %s" s
+      | Ok () ->
+          Lwt.finalize (fun () -> thn pool) (fun () -> Datastore.close_db pool)
     with e ->
       let message = match e with Failure s -> s | _ -> Printexc.to_string e in
-      err message |> Lwt.return
+      err message
 
 let generate_json in_filename =
-  let aux_ok con = get_models ~year:2023 ~depth:1 con |> fun xs -> `Assoc xs in
-  let aux_err msg = `Assoc [ ("error", `String msg) ] in
+  let aux_ok pool =
+    Datastore.use pool @@ fun con ->
+    get_models ~year:2023 ~depth:1 con |> fun xs -> `Assoc xs
+  in
+  let aux_err msg = `Assoc [ ("error", `String msg) ] |> Lwt.return in
   generate in_filename aux_ok aux_err
 
 let handle_query ~in_filename ~query =
@@ -484,8 +482,8 @@ let handle_query ~in_filename ~query =
     | Int i -> `Int i
     | Null -> `Null
   in
-  let err msg = `Assoc [ ("error", `String msg) ] in
-  let thn con =
+  let err msg = `Assoc [ ("error", `String msg) ] |> Lwt.return in
+  let thn pool =
     let sql_queries =
       match query with
       | `List xs ->
@@ -495,6 +493,7 @@ let handle_query ~in_filename ~query =
                | _ -> failwith "Invalid query")
       | _ -> failwith "Invalid query"
     in
+    Datastore.use pool @@ fun con ->
     sql_queries
     |> List.map (fun q ->
            match Datastore.(query (prepare con q) []) with

test/dune

@@ -1,3 +1,5 @@
 (tests
  (names test_parser test_model test_loader test_datastore)
+ (preprocess
+  (pps lwt_ppx))
  (libraries qash alcotest fsnotify_stub))

test/test_datastore.ml

@@ -1,16 +1,25 @@
 open Qash.Datastore
 
+let with_temp_file f =
+  let filepath = Filename.temp_file "qash" "test" in
+  Fun.protect ~finally:(fun () -> Sys.remove filepath) (fun () -> f filepath)
+
 let test_query_case1 () =
-  let con = connect in_memory_database in
+  with_temp_file @@ fun db_path ->
+  let pool = open_db db_path in
+  Lwt_main.run @@ use pool
+  @@ fun con ->
   let stmt = prepare con "SELECT 1" in
   let res = query stmt [] in
   match res with
   | Ok x -> assert (x = [ [ Int 1 ] ])
   | Error s -> Alcotest.failf "query failed: %s" s
 
 let test_query_case2 () =
-  let con = connect in_memory_database in
-
+  with_temp_file @@ fun db_path ->
+  let pool = open_db db_path in
+  Lwt_main.run @@ use pool
+  @@ fun con ->
   let stmt = prepare con "CREATE TABLE t (id INTEGER)" in
   let res = execute stmt [] in
   assert (res = Ok ());
@@ -41,6 +50,26 @@ let test_query_case2 () =
 
   ()
 
+let test_query_case3 () =
+  with_temp_file @@ fun db_path ->
+  let pool = open_db db_path in
+  let got =
+    Lwt_main.run
+      (use pool (fun con ->
+           execute (prepare con "CREATE TABLE t (id INTEGER)") []
+           |> Result.get_ok);%lwt
+       List.init 10 Fun.id
+       |> Lwt_list.iter_p (fun i ->
+              use pool (fun con ->
+                  execute (prepare con "INSERT INTO t VALUES (?)") [ Int i ]
+                  |> Result.get_ok));%lwt
+       use pool (fun con ->
+           query (prepare con "SELECT id FROM t ORDER BY id") []
+           |> Result.get_ok))
+  in
+  assert (got = List.init 10 (fun i -> [ Int i ]));
+  ()
+
 let () =
   let open Alcotest in
   run "datastore"
@@ -49,5 +78,6 @@ let () =
         [
           test_case "case1" `Quick test_query_case1;
           test_case "case2" `Quick test_query_case2;
+          test_case "case3" `Quick test_query_case3;
         ] );
     ]