add after_shutdown to ocurl-lwt client

This commit is contained in:
Simon Cruanes 2025-12-10 11:53:31 -05:00
parent a6bf8171bb
commit b55598685f
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
3 changed files with 27 additions and 8 deletions

View file

@ -90,12 +90,13 @@ let create_exporter ?(config = Config.make ()) () =
let create_backend = create_exporter
let setup_ ?config () : unit =
let setup_ ?config () : Exporter.t =
let exp = create_backend ?config () in
Main_exporter.set exp;
()
exp
let setup ?config ?(enable = true) () = if enable then setup_ ?config ()
let setup ?config ?(enable = true) () =
if enable then ignore (setup_ ?config () : Exporter.t)
let remove_exporter () : unit Lwt.t =
let done_fut, done_u = Lwt.wait () in
@ -104,18 +105,21 @@ let remove_exporter () : unit Lwt.t =
let remove_backend = remove_exporter
let with_setup ?(config = Config.make ()) ?(enable = true) () f : _ Lwt.t =
if enable then (
let with_setup ?(after_shutdown = ignore) ?(config = Config.make ())
?(enable = true) () f : _ Lwt.t =
if enable then
let open Lwt.Syntax in
setup_ ~config ();
let exp = setup_ ~config () in
Lwt.catch
(fun () ->
let* res = f () in
let+ () = remove_backend () in
after_shutdown exp;
res)
(fun exn ->
let* () = remove_backend () in
after_shutdown exp;
Lwt.reraise exn)
) else
else
f ()

View file

@ -32,6 +32,11 @@ val remove_backend : unit -> unit Lwt.t
@since NEXT_RELEASE *)
val with_setup :
?config:Config.t -> ?enable:bool -> unit -> (unit -> 'a Lwt.t) -> 'a Lwt.t
?after_shutdown:(Opentelemetry.Exporter.t -> unit) ->
?config:Config.t ->
?enable:bool ->
unit ->
(unit -> 'a Lwt.t) ->
'a Lwt.t
(** [with_setup () f] is like [setup(); f()] but takes care of cleaning up after
[f()] returns See {!setup} for more details. *)

View file

@ -170,8 +170,18 @@ let () =
Printf.printf "\ndone. %d spans in %.4fs (%.4f/s)\n%!" (Atomic.get num_tr)
elapsed n_per_sec
in
let after_exp_shutdown exp =
(* print some stats *)
if !final_stats then (
let ms = OT.Exporter.self_metrics exp in
Format.eprintf "@[exporter metrics:@ %a@]@."
(Format.pp_print_list Opentelemetry.Metrics.pp)
ms
)
in
Lwt_main.run
@@
let@ () = Fun.protect ~finally in
Opentelemetry_client_ocurl_lwt.with_setup ~config () run
~after_shutdown:after_exp_shutdown