events: cosmetics and tweaks

raphael-proust · raphael-proust · commit 6011e300028f · 2025-10-14T11:19:32.000+08:00
diff --git a/src/unix/lwt_main.ml b/src/unix/lwt_main.ml
@@ -22,6 +22,7 @@ let abandon_yielded_and_paused () =
 
 type Runtime_events.User.tag += Scheduler_call
 let sch_call = Runtime_events.User.register "lwt-sch-call" Scheduler_call Runtime_events.Type.span
+
 type Runtime_events.User.tag += Scheduler_lap
 let sch_lap = Runtime_events.User.register "lwt-sch-lap" Scheduler_lap Runtime_events.Type.unit
 
@@ -37,14 +38,12 @@ let run p =
     (Lwt.Private.Multidomain_sync.register_notification[@alert "-trespassing"]) domain(fun () -> Lwt_unix.send_notification n)
   end
   in
-  Runtime_events.User.write sch_call Begin;
   let rec run_loop () =
     Runtime_events.User.write sch_lap ();
     Runtime_events.User.write Lwt_unix.unix_job_count (Domain.DLS.get Lwt_unix.job_count) ;
     Runtime_events.User.write (Lwt.Private.paused_count[@alert "-trespassing"]) (Lwt.paused_count ()) ;
     match Lwt.poll p with
     | Some x ->
-      Runtime_events.User.write sch_call End;
       x
     | None ->
       (* Call enter hooks. *)
@@ -64,7 +63,10 @@ let run p =
       run_loop ()
   in
 
-  run_loop ()
+  Runtime_events.User.write sch_call Begin;
+  Fun.protect
+    ~finally:(fun () -> Runtime_events.User.write sch_call End)
+    (fun () -> run_loop ())
 
 let run_already_called = Domain.DLS.new_key (fun () -> `No)
 let run_already_called_mutex = Domain.DLS.new_key (fun () -> Mutex.create ())
diff --git a/src/unix/lwt_unix.cppo.ml b/src/unix/lwt_unix.cppo.ml
@@ -14,6 +14,7 @@ open Lwt.Infix
 
 type Runtime_events.User.tag += Unix_job_count
 let unix_job_count = Runtime_events.User.register "lwt-unix-job-count" Unix_job_count Runtime_events.Type.int
+
 let job_count = Domain.DLS.new_key (fun () -> 0)
 
 (* +-----------------------------------------------------------------+
