CP-53843: reusable SM session

ocaml/xapi-idl/lib/debug_info.ml

@@ -93,6 +93,8 @@ let with_dbg ?attributes ?(with_thread = false) ?(module_name = "") ~name ~dbg f
   | false ->
       f_with_trace ()
 
+let span_of di = di.tracing
+
 let traceparent_of_dbg dbg =
   match String.split_on_char separator dbg with
   | [_; trace_context] -> (

ocaml/xapi-idl/lib/debug_info.mli

@@ -31,4 +31,6 @@ val with_dbg :
   -> (t -> 'a)
   -> 'a
 
+val span_of : t -> Tracing.Span.t option
+
 val traceparent_of_dbg : string -> string option

ocaml/xapi/context.ml

@@ -18,7 +18,10 @@ module D = Debug.Make (struct let name = "dummytaskhelper" end)
 (** Every operation has an origin: either the HTTP connection it came from or
     an internal subsystem (eg synchroniser thread / event handler
     thread) *)
-type origin = Http of Http.Request.t * Unix.file_descr | Internal
+type origin =
+  | Http of Http.Request.t * Unix.file_descr
+  | Internal
+  | Internal_Traced of Tracing.Span.t option
 
 let string_of_origin = function
   | Http (req, fd) ->
@@ -32,7 +35,7 @@ let string_of_origin = function
       (* unfortunately all connections come from stunnel on localhost *)
       Printf.sprintf "HTTP request from %s with User-Agent: %s" peer
         (Option.value ~default:"unknown" req.Http.Request.user_agent)
-  | Internal ->
+  | Internal | Internal_Traced _ ->
       "Internal"
 
 (** A Context is used to represent every API invocation. It may be extended
@@ -105,7 +108,7 @@ let default_database () =
 
 let preauth ~__context =
   match __context.origin with
-  | Internal ->
+  | Internal | Internal_Traced _ ->
       None
   | Http (_, s) -> (
     match Unix.getsockname s with
@@ -203,7 +206,7 @@ let trackid ?(with_brackets = false) ?(prefix = "") __context =
   trackid_of_session ~with_brackets ~prefix __context.session_id
 
 let _client_of_origin = function
-  | Internal ->
+  | Internal | Internal_Traced _ ->
       None
   | Http (req, fd) ->
       Http_svr.client_of_req_and_fd req fd
@@ -233,7 +236,9 @@ let parent_of_origin (origin : origin) span_name =
       let* span_context = SpanContext.of_trace_context context in
       let span = Tracer.span_of_span_context span_context span_name in
       Some span
-  | _ ->
+  | Internal_Traced span ->
+      span
+  | Internal ->
       None
 
 let attribute_helper_fn f v = Option.fold ~none:[] ~some:f v
@@ -312,7 +317,7 @@ let make_attributes ?task_name ?task_id ?task_uuid ?session_id ?origin () =
   ; attribute_helper_fn
       (fun origin ->
         match origin with
-        | Internal ->
+        | Internal | Internal_Traced _ ->
             [("xs.xapi.task.origin", "internal")]
         | Http (req, s) ->
             [attr_of_req req; attr_of_fd s] |> List.concat
@@ -519,7 +524,11 @@ let get_client_ip context =
   context.client |> Option.map (fun (_, ip) -> Ipaddr.to_string ip)
 
 let get_user_agent context =
-  match context.origin with Internal -> None | Http (rq, _) -> rq.user_agent
+  match context.origin with
+  | Internal | Internal_Traced _ ->
+      None
+  | Http (rq, _) ->
+      rq.user_agent
 
 let finally_destroy_context ~__context f =
   let tracing = __context.tracing in

ocaml/xapi/context.mli

@@ -16,7 +16,10 @@
     to include extra data without changing all the autogenerated signatures *)
 type t
 
-type origin = Http of Http.Request.t * Unix.file_descr | Internal
+type origin =
+  | Http of Http.Request.t * Unix.file_descr
+  | Internal
+  | Internal_Traced of Tracing.Span.t option
 
 (** {6 Constructors} *)
 

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,31 +318,6 @@ let methodResponse xml =
 (****************************************************************************************)
 (* Functions that actually execute the python backends *)
 
-let with_session sr f =
-  Server_helpers.exec_with_new_task "sm_exec" (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 ->
@@ -466,7 +439,8 @@ let exec_xmlrpc ~dbg ?context:_ ?(needs_session = true) (driver : string)
           )
   in
   if needs_session then
-    with_session call.sr_ref (fun 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

ocaml/xapi/xapi_session.ml

@@ -1589,3 +1589,71 @@ 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 finally = Xapi_stdext_pervasives.Pervasiveext.finally
+
+  let sm_username = "__sm__backend"
+
+  let reusable_session = Atomic.make Ref.null
+
+  let is_valid_session ~__context session_id =
+    if Atomic.get reusable_session = Ref.null then
+      false
+    else
+      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 =
+    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
+
+  let rec get_session ~__context sr =
+    let session = Atomic.get reusable_session in
+    if is_valid_session ~__context session then
+      session
+    else
+      let new_session = create_session ~__context in
+      if Atomic.compare_and_set reusable_session session new_session then (
+        session_access ~__context session sr ;
+        new_session
+      ) else (
+        destroy_db_session ~__context ~self:new_session ;
+        (get_session [@tailcall]) ~__context sr
+      )
+
+  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 in
+          let () = session_access ~__context session_id 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