[sql] reuse prepared statements

Summary:
Preparing statements allocates memory and is generally best done once and for all.

We need to re-prepare statements when opening a new DB connexion (after each fork).

Reviewed By: mbouaziz

Differential Revision: D5824157

fbshipit-source-id: 4d239ac

infer/src/base/KeyValue.ml

@@ -7,6 +7,7 @@
  * of patent rights can be found in the PATENTS file in the same directory.
  *)
 open! IStd
+module L = Logging
 
 module type Table = sig
   type key
@@ -65,35 +66,53 @@ module Make (Table : Table) : S with module Table = Table = struct
   (* cannot mix, e.g., blob_key and blob_value now *)
   include Unsafe
 
-  let replace =
-    let replace_statement =
-      Printf.sprintf "REPLACE INTO %s(key, value) VALUES(:k, :v)" Table.table
+  let register_statement stmt_fmt =
+    let k stmt0 =
+      let stmt_ref = ref None in
+      let new_statement db =
+        let stmt = Sqlite3.prepare db stmt0 in
+        ResultsDir.on_close_database ~f:(fun _ ->
+            Option.iter !stmt_ref ~f:(SqliteUtils.finalize ~log:"db close callback") ) ;
+        stmt_ref := Some stmt
       in
-    fun key value ->
-      let db = ResultsDir.get_database () in
-      let replace_stmt = Sqlite3.prepare db replace_statement in
-      SqliteUtils.check_sqlite_error ~log:"replace bind key"
-        (Sqlite3.bind replace_stmt 1 (blob_of_key key)) ;
-      SqliteUtils.check_sqlite_error ~log:"replace bind value"
-        (Sqlite3.bind replace_stmt 2 (blob_of_value value)) ;
-      SqliteUtils.sqlite_unit_step ~log:"KeyValue.replace" replace_stmt
-
-  let find =
-    let select_statement = Printf.sprintf "SELECT value FROM %s WHERE key = :k" Table.table in
-    fun key ->
-      let db = ResultsDir.get_database () in
-      let select_stmt = Sqlite3.prepare db select_statement in
-      SqliteUtils.check_sqlite_error ~log:"insert bind key"
-        (Sqlite3.bind select_stmt 1 (blob_of_key key)) ;
-      Option.bind ~f:value_of_blob
-        (SqliteUtils.sqlite_result_step ~log:"KeyValue.find" select_stmt)
-
-  let delete =
-    let delete_statement = Printf.sprintf "DELETE FROM %s WHERE key = :k" Table.table in
-    fun key ->
-      let db = ResultsDir.get_database () in
-      let delete_stmt = Sqlite3.prepare db delete_statement in
-      SqliteUtils.check_sqlite_error ~log:"delete bind key"
-        (Sqlite3.bind delete_stmt 1 (blob_of_key key)) ;
-      SqliteUtils.sqlite_unit_step ~log:"KeyValue.delete" delete_stmt
+      ResultsDir.on_new_database_connection ~f:new_statement ;
+      fun () ->
+        match !stmt_ref with
+        | None
+         -> L.(die InternalError) "database not initialized"
+        | Some stmt
+         -> Sqlite3.reset stmt |> SqliteUtils.check_sqlite_error ~log:"reset prepared statement" ;
+            Sqlite3.clear_bindings stmt
+            |> SqliteUtils.check_sqlite_error ~log:"clear bindings of prepared statement" ;
+            stmt
+    in
+    Printf.ksprintf k stmt_fmt
+
+  let get_replace_statement =
+    register_statement "REPLACE INTO %s(key, value) VALUES(:k, :v)" Table.table
+
+  let replace key value =
+    let replace_stmt = get_replace_statement () in
+    Sqlite3.bind replace_stmt 1 (blob_of_key key)
+    |> SqliteUtils.check_sqlite_error ~log:"replace bind key" ;
+    Sqlite3.bind replace_stmt 2 (blob_of_value value)
+    |> SqliteUtils.check_sqlite_error ~log:"replace bind value" ;
+    SqliteUtils.sqlite_unit_step ~finalize:false ~log:"KeyValue.replace" replace_stmt
+
+  let get_select_statement = register_statement "SELECT value FROM %s WHERE key = :k" Table.table
+
+  let find key =
+    let select_stmt = get_select_statement () in
+    Sqlite3.bind select_stmt 1 (blob_of_key key)
+    |> SqliteUtils.check_sqlite_error ~log:"insert bind key" ;
+    SqliteUtils.sqlite_result_step ~finalize:false ~log:"KeyValue.find" select_stmt
+    |> Option.bind ~f:value_of_blob
+
+  let get_delete_statement = register_statement "DELETE FROM %s WHERE key = :k" Table.table
+
+  let delete key =
+    let delete_stmt = get_delete_statement () in
+    Sqlite3.bind delete_stmt 1 (blob_of_key key)
+    |> SqliteUtils.check_sqlite_error ~log:"delete bind key" ;
+    SqliteUtils.sqlite_unit_step ~finalize:false ~log:"KeyValue.delete" delete_stmt
 end

infer/src/base/ResultsDir.ml

@@ -12,10 +12,6 @@ module L = Logging
 
 let database : Sqlite3.db option ref = ref None
 
-let () =
-  Epilogues.register "closing results database" ~f:(fun () ->
-      Option.iter !database ~f:SqliteUtils.db_close )
-
 let database_filename = "results.db"
 
 let database_fullpath = Config.results_dir ^/ database_filename
@@ -69,13 +65,33 @@ let create_db () =
   try Sys.rename temp_db database_fullpath
   with Sys_error _ -> (* lost the race, doesn't matter *) ()
 
+let new_db_callbacks = ref []
+
+let on_new_database_connection ~f = new_db_callbacks := f :: !new_db_callbacks
+
+let close_db_callbacks = ref []
+
+let on_close_database ~f = close_db_callbacks := f :: !close_db_callbacks
+
+let do_db_close db =
+  List.iter ~f:(fun callback -> callback db) !close_db_callbacks ;
+  close_db_callbacks := [] ;
+  SqliteUtils.db_close db
+
+let db_close () =
+  Option.iter !database ~f:do_db_close ;
+  database := None
+
 let new_database_connection () =
-  Option.iter !database ~f:SqliteUtils.db_close ;
+  db_close () ;
   let db = Sqlite3.db_open ~mode:`NO_CREATE ~cache:`PRIVATE ~mutex:`FULL database_fullpath in
   Sqlite3.busy_timeout db 1000 ;
   (* Higher level of "synchronous" are only useful to guarantee that the db will not be corrupted if the machine crashes for some reason before the data has been actually written to disk. We do not need this kind of guarantee for infer results as one can always rerun infer if interrupted. *)
   SqliteUtils.exec db ~log:"synchronous=OFF" ~stmt:"PRAGMA synchronous=OFF" ;
-  database := Some db
+  database := Some db ;
+  List.iter ~f:(fun callback -> callback db) !new_db_callbacks
+
+let () = Epilogues.register "closing results database" ~f:db_close
 
 let create_results_dir () =
   Unix.mkdir_p Config.results_dir ;

infer/src/base/ResultsDir.mli

@@ -38,3 +38,7 @@ val delete_capture_and_analysis_data : unit -> unit
 
 val canonicalize_db : unit -> unit
 (** put the database on disk in deterministic form *)
+
+val on_new_database_connection : f:(Sqlite3.db -> unit) -> unit
+
+val on_close_database : f:(Sqlite3.db -> unit) -> unit

infer/src/base/SqliteUtils.ml

@@ -33,7 +33,7 @@ let finalize ~log stmt =
   | Sqlite3.Error err
    -> error ~fatal:true "finalize: %s: %s" log err
 
-let sqlite_result_rev_list_step ~log stmt =
+let sqlite_result_rev_list_step ?finalize:(do_finalize = true) ~log stmt =
   let rec aux rev_results =
     match Sqlite3.step stmt with
     | Sqlite3.Rc.ROW
@@ -45,10 +45,11 @@ let sqlite_result_rev_list_step ~log stmt =
     | err
      -> L.die InternalError "%s: %s" log (Sqlite3.Rc.to_string err)
   in
-  protect ~finally:(fun () -> finalize ~log stmt) ~f:(fun () -> aux [])
+  if do_finalize then protect ~finally:(fun () -> finalize ~log stmt) ~f:(fun () -> aux [])
+  else aux []
 
-let sqlite_result_step ~log stmt =
-  match sqlite_result_rev_list_step ~log stmt with
+let sqlite_result_step ?finalize ~log stmt =
+  match sqlite_result_rev_list_step ?finalize ~log stmt with
   | []
    -> None
   | [x]
@@ -56,8 +57,8 @@ let sqlite_result_step ~log stmt =
   | l
    -> L.die InternalError "%s: zero or one result expected, got %d instead" log (List.length l)
 
-let sqlite_unit_step ~log stmt =
-  match sqlite_result_rev_list_step ~log stmt with
+let sqlite_unit_step ?finalize ~log stmt =
+  match sqlite_result_rev_list_step ?finalize ~log stmt with
   | []
    -> ()
   | l

infer/src/base/SqliteUtils.mli

@@ -22,13 +22,14 @@ val exec : Sqlite3.db -> log:string -> stmt:string -> unit
 val finalize : log:string -> Sqlite3.stmt -> unit
 (** Finalize the given [stmt]. Raises [Error] on failure. *)
 
-val sqlite_result_rev_list_step : log:string -> Sqlite3.stmt -> Sqlite3.Data.t option list
+val sqlite_result_rev_list_step :
+  ?finalize:bool -> log:string -> Sqlite3.stmt -> Sqlite3.Data.t option list
 (** Return a reversed list of results obtained by repeatedly stepping through [stmt] and saving only column 0 of each returned row (all that's been needed so far). *)
 
-val sqlite_result_step : log:string -> Sqlite3.stmt -> Sqlite3.Data.t option
+val sqlite_result_step : ?finalize:bool -> log:string -> Sqlite3.stmt -> Sqlite3.Data.t option
 (** Same as [sqlite_result_rev_list_step] but asserts that exactly one result is returned. *)
 
-val sqlite_unit_step : log:string -> Sqlite3.stmt -> unit
+val sqlite_unit_step : ?finalize:bool -> log:string -> Sqlite3.stmt -> unit
 (** Same as [sqlite_result_rev_list_step] but asserts that no result is returned. *)
 
 val db_close : Sqlite3.db -> unit