CP-53843: reusable SM session (#7002)

Builds on top of #6991

The session reuse logic in `xapi_session.ml` does work in this usecase,
sm calls are made with `pool=false`. So reimplement the logic so that we
use a single session for `sm_exec` per sr.

This is achieved by having a hashtable that maps SRs to a corresponding
session. The session is created the first time a particular SR needs it
and then get reused afterwards. The session gets recreated only if it
becomes invalid.

This should help the number of database calls during congestion times.

Passes BVT and BST.

Author: GabrielBuica

ocaml/xapi/sm_exec.ml

@@ -30,8 +30,6 @@ module E = Debug.Make (struct let name = "mscgen" end)
 
 let cmd_name driver = sprintf "%s/%sSR" !Xapi_globs.sm_dir driver
 
-let sm_username = "__sm__backend"
-
 let with_dbg ~name ~dbg f =
   Debug_info.with_dbg ~module_name:"Sm_exec" ~name ~dbg f
 
@@ -320,32 +318,6 @@ let methodResponse xml =
 (****************************************************************************************)
 (* Functions that actually execute the python backends *)
 
-let with_session ~traceparent sr f =
-  Server_helpers.exec_with_new_task "sm_exec"
-    ~origin:(Internal_Traced traceparent) (fun __context ->
-      let create_session () =
-        let host = !Xapi_globs.localhost_ref in
-        let session =
-          Xapi_session.login_no_password ~__context ~uname:None ~host
-            ~pool:false ~is_local_superuser:true ~subject:Ref.null
-            ~auth_user_sid:"" ~auth_user_name:sm_username ~rbac_permissions:[]
-        in
-        (* Give this session access to this particular SR *)
-        Option.iter
-          (fun sr ->
-            Db.Session.add_to_other_config ~__context ~self:session
-              ~key:Xapi_globs._sm_session ~value:(Ref.string_of sr)
-          )
-          sr ;
-        session
-      in
-      let destroy_session session_id =
-        Xapi_session.destroy_db_session ~__context ~self:session_id
-      in
-      let session_id = create_session () in
-      finally (fun () -> f session_id) (fun () -> destroy_session session_id)
-  )
-
 let exec_xmlrpc ~dbg ?context:_ ?(needs_session = true) (driver : string)
     (call : call) =
   with_dbg ~name:call.cmd ~dbg @@ fun di ->
@@ -467,8 +439,9 @@ let exec_xmlrpc ~dbg ?context:_ ?(needs_session = true) (driver : string)
           )
   in
   if needs_session then
-    with_session ~traceparent:(Debug_info.span_of di) call.sr_ref
-      (fun session_id -> do_call {call with session_ref= Some session_id}
+    Xapi_session.SM.with_session ~traceparent:(Debug_info.span_of di)
+      call.sr_ref (fun session_id ->
+        do_call {call with session_ref= Some session_id}
     )
   else
     do_call call

ocaml/xapi/xapi_session.ml

@@ -1589,3 +1589,77 @@ let create_from_db_file ~__context ~filename =
   in
   let db_ref = Some (Xapi_database.Db_ref.in_memory (Atomic.make db)) in
   create_readonly_session ~__context ~uname:"db-from-file" ~db_ref
+
+module SM = struct
+  let reusable_sessions : (string option, API.ref_session) Hashtbl.t =
+    Hashtbl.create 5
+
+  let sm_sessions_m = Mutex.create ()
+
+  let with_sm_sessions_lock f =
+    Xapi_stdext_threads.Threadext.Mutex.execute sm_sessions_m f
+
+  let finally = Xapi_stdext_pervasives.Pervasiveext.finally
+
+  let sm_username = "__sm__backend"
+
+  let is_valid_session ~__context session_id =
+    try
+      (* Call an API function to check the session is still valid *)
+      let rpc = Helpers.make_rpc ~__context in
+      ignore (Client.Pool.get_all ~rpc ~session_id) ;
+      true
+    with Api_errors.Server_error (err, _) ->
+      debug "%s: Invalid session: %s" __FUNCTION__ err ;
+      false
+
+  let session_access ~__context session sr =
+    (* Give this session access to this particular SR *)
+    Option.iter
+      (fun sr ->
+        Db.Session.add_to_other_config ~__context ~self:session
+          ~key:Xapi_globs._sm_session ~value:(Ref.string_of sr)
+      )
+      sr
+
+  let create_session ~__context sr =
+    let host = !Xapi_globs.localhost_ref in
+    let session =
+      login_no_password ~__context ~uname:None ~host ~pool:false
+        ~is_local_superuser:true ~subject:Ref.null ~auth_user_sid:""
+        ~auth_user_name:sm_username ~rbac_permissions:[]
+    in
+    session_access ~__context session sr ;
+    session
+
+  let get_session ~__context sr =
+    let sr_key = Option.map Ref.string_of sr in
+    with_sm_sessions_lock (fun () ->
+        match Hashtbl.find_opt reusable_sessions sr_key with
+        | Some session when is_valid_session ~__context session ->
+            session
+        | Some _ ->
+            Hashtbl.remove reusable_sessions sr_key ;
+            let new_session = create_session ~__context sr in
+            Hashtbl.add reusable_sessions sr_key new_session ;
+            new_session
+        | None ->
+            let new_session = create_session ~__context sr in
+            Hashtbl.add reusable_sessions sr_key new_session ;
+            new_session
+    )
+
+  let with_session ~traceparent sr f =
+    Server_helpers.exec_with_new_task "sm_exec"
+      ~origin:(Internal_Traced traceparent) (fun __context ->
+        if !Xapi_globs.reuse_pool_sessions then
+          let session_id = get_session ~__context sr in
+          f session_id
+        else
+          let session_id = create_session ~__context sr in
+          let destroy_session () =
+            destroy_db_session ~__context ~self:session_id
+          in
+          finally (fun () -> f session_id) destroy_session
+    )
+end

ocaml/xapi/xapi_session.mli

@@ -112,3 +112,10 @@ val set_local_auth_max_threads : int64 -> unit
 val set_ext_auth_max_threads : int64 -> unit
 
 val clear_external_auth_cache : unit -> unit
+
+module SM : sig
+  val with_session :
+       traceparent:Tracing.Span.t option
+    -> [< Uuidx.all] Ref.t option
+    -> (Uuidx.secret Ref.t -> 'a)
+    -> 'a end