Merge pull request #17 from imandra-ai/simon/synchronous-logger

sync logger

Author: c-cube

src/log/dune

@@ -3,6 +3,7 @@
  (name imandrakit_log)
  (preprocess
   (pps ppx_deriving.std imandrakit.ser-pack.ppx imandrakit.twine.ppx))
+ (private_modules lock_)
  (flags :standard -open Imandrakit)
  (libraries
   imandrakit

src/log/lock_.ml

@@ -0,0 +1,20 @@
+open struct
+  external reraise : exn -> 'a = "%reraise"
+end
+
+type 'a t = {
+  v: 'a;
+  m: Mutex.t;
+}
+
+let create v = { v; m = Mutex.create () }
+
+let with_ (self : _ t) f =
+  Mutex.lock self.m;
+  match f self.v with
+  | x ->
+    Mutex.unlock self.m;
+    x
+  | exception e ->
+    Mutex.unlock self.m;
+    reraise e

src/log/log_google.ml

@@ -22,19 +22,23 @@ let event_to_json ?(other_fields = []) (ev : Log_event.t) : json =
   `Assoc (List.rev_append other_fields fields)
 
 let logger ?other_fields oc : Logger.Output.t =
-  Logger.Output.to_event () ~emit_ev:(fun ev ->
+  let oc = Lock_.create oc in
+  Logger.Output.to_event ()
+    ~flush:(fun () -> Lock_.with_ oc Stdlib.flush)
+    ~emit_ev:(fun ev ->
       let json = event_to_json ?other_fields ev in
 
       try
-        ((* use a local buffer *)
-         let@ buf = Apool.with_resource Logger.Output.buf_pool in
-         assert (Buffer.length buf = 0);
-         Yojson.Safe.to_buffer buf json;
-         (* write buffer's content immediately *)
-         Buffer.output_buffer oc buf);
+        (* use a local buffer *)
+        let@ buf = Apool.with_resource Logger.Output.buf_pool in
+        assert (Buffer.length buf = 0);
+        Yojson.Safe.to_buffer buf json;
 
+        (* write buffer's content immediately *)
+        let@ oc = Lock_.with_ oc in
+        Buffer.output_buffer oc buf;
         output_char oc '\n';
         flush oc
       with exn ->
-        Printf.eprintf "log to json chan: failed with %s\n%!"
+        Printf.eprintf "log to json in google format: failed with\n%s\n%!"
           (Printexc.to_string exn))

src/log/logger.ml

@@ -1,5 +1,4 @@
-open Log_level
-module Sync_queue = Moonpool.Blocking_queue
+open! Log_level
 
 type level = Log_level.t [@@deriving show, eq]
 
@@ -51,9 +50,12 @@ open struct
 end
 
 module Output = struct
-  type t = { emit: Log_event.t -> unit } [@@unboxed]
+  type t = {
+    emit: Log_event.t -> unit;
+    flush: unit -> unit;
+  }
 
-  let to_event ~emit_ev () : t = { emit = emit_ev }
+  let to_event ~emit_ev ~flush () : t = { emit = emit_ev; flush }
   let () = Fmt.set_color_default true
 
   let to_str_ (ev : Log_event.t) : string =
@@ -92,41 +94,49 @@ module Output = struct
 
     Buf_fmt.get_contents buf_fmt
 
-  let to_str ~(emit_str : string -> unit) () : t =
+  let to_str ~(emit_str : string -> unit) ~flush () : t =
     {
       emit =
         (fun ev ->
           let s = to_str_ ev in
           emit_str s);
+      flush;
     }
 
-  let to_chan (oc : out_channel) : t =
+  let to_chan ?(autoflush = true) (oc : out_channel) : t =
+    let oc = Lock_.create oc in
     to_str
+      ~flush:(fun () -> Lock_.with_ oc Stdlib.flush)
       ~emit_str:(fun s ->
         try
+          let@ oc = Lock_.with_ oc in
           output_string oc s;
           output_char oc '\n';
-          flush oc
+          if autoflush then Stdlib.flush oc
         with _ -> Printf.eprintf "logger: failed to log to chan\n%!")
       ()
 
   let stdout () = to_chan stdout
   let stderr () = to_chan stderr
 
   let filter_level pred (self : t) : t =
-    { emit = (fun ev -> if pred ev.lvl then self.emit ev) }
+    { emit = (fun ev -> if pred ev.lvl then self.emit ev); flush = self.flush }
 
   let buf_pool : Buffer.t Apool.t =
     Apool.create ~clear:Buffer.reset
       ~mk_item:(fun () -> Buffer.create 256)
       ~max_size:8 ()
 
   (** Logger that writes events, one per line, on the given channel. *)
-  let to_chan_jsonl (oc : out_channel) : t =
-    to_event () ~emit_ev:(fun ev ->
+  let to_chan_jsonl ?(autoflush = true) (oc : out_channel) : t =
+    let oc = Lock_.create oc in
+    to_event ()
+      ~flush:(fun () -> Lock_.with_ oc flush)
+      ~emit_ev:(fun ev ->
         let json = Log_event.to_yojson ev in
 
         try
+          let@ oc = Lock_.with_ oc in
           ((* use a local buffer *)
            let@ buf = Apool.with_resource buf_pool in
            Buffer.clear buf;
@@ -135,16 +145,12 @@ module Output = struct
            Buffer.output_buffer oc buf);
 
           output_char oc '\n';
-          flush oc
+          if autoflush then flush oc
         with exn ->
           Printf.eprintf "log to json chan: failed with %s\n%!"
             (Printexc.to_string exn))
 end
 
-type task =
-  | T_fence of { wakeup: unit Moonpool.Fut.promise }
-  | T_emit of Log_event.t
-
 type capture_meta_hook = unit -> (string * Log_meta.t) list
 
 open struct
@@ -183,19 +189,23 @@ end
 let add_rich_tag = add_rich_tag
 
 type t = {
-  q: task Sync_queue.t;
-  events: Log_event.t Observer.t;
+  active: bool Atomic.t;
+  events: Log_event.t Observer.t;  (** Additional subscriptions on events *)
   outputs: Output.t list Atomic.t;
-  reporter: Logs.reporter;
-  mutable bg_thread: Thread.t option;
+  emit_ev: Log_event.t -> unit;
+  emit_fence: unit -> unit;
+  reporter: Logs.reporter;  (** As reporter *)
 }
 
 let[@inline] as_reporter self = self.reporter
 let[@inline] events self = self.events
+let[@inline] shutdown self = Atomic.set self.active false
 
-let shutdown self =
-  Sync_queue.close self.q;
-  Option.iter Thread.join self.bg_thread
+let trace_level_of_level : level -> Trace.Level.t = function
+  | Info | App -> Trace.Level.Info
+  | Error -> Trace.Level.Error
+  | Warning -> Trace.Level.Warning
+  | Debug -> Trace.Level.Debug3
 
 let add_output self out : unit =
   while
@@ -267,7 +277,8 @@ let to_event_if_ (p : level -> bool) ~emit_ev : Logs.reporter =
            at least the TEF collector needs to know on which thread we are running. *)
         if Trace_core.enabled () then (
           let msg = Ansi_clean.remove_escape_codes msg in
-          Trace_core.message msg ~data:(fun () ->
+          Trace_core.message ~level:(trace_level_of_level level) msg
+            ~data:(fun () ->
               let meta =
                 List.map (fun (k, v) -> k, Log_meta.to_trace_data v) meta
               in
@@ -290,6 +301,7 @@ let to_event_if_ (p : level -> bool) ~emit_ev : Logs.reporter =
   in
   { Logs.report }
 
+(*
 let bg_thread_loop_ (self : t) : unit =
   Trace_core.set_thread_name "logger.bg";
   ignore
@@ -304,7 +316,7 @@ let bg_thread_loop_ (self : t) : unit =
          Sys.sigusr2;
          Sys.sigvtalrm;
        ]
-    : _ list);
+      : _ list);
 
   let local_q = Queue.create () in
   try
@@ -323,55 +335,71 @@ let bg_thread_loop_ (self : t) : unit =
       Queue.clear local_q
     done
   with Sync_queue.Closed -> ()
+  *)
 
-let fence_ : (unit -> unit Moonpool.Fut.t) ref =
-  ref (fun () -> Moonpool.Fut.return ())
-
-let[@inline] emit_ev (self : t) ev : unit =
-  try Sync_queue.push self.q (T_emit ev) with Sync_queue.Closed -> ()
+let fence_ : (unit -> unit) ref = ref ignore
+let set_as_fence_ (self : t) : unit = fence_ := self.emit_fence
 
 let to_outputs (outs : Output.t list) : t =
+  let active = Atomic.make true in
   let outputs = Atomic.make outs in
   let events = Observer.create () in
-  let q = Sync_queue.create () in
+  let emit_ev ev =
+    (try Observer.emit events ev
+     with exn ->
+       Printf.eprintf "logger observer failed with %s\n%!"
+         (Printexc.to_string exn));
+    match Atomic.get outputs with
+    | [] -> ()
+    | outs ->
+      List.iter
+        (fun (out : Output.t) ->
+          try out.emit ev
+          with exn ->
+            Printf.eprintf "logger output failed with %s\n%!"
+              (Printexc.to_string exn))
+        outs
+  in
   let reporter =
     to_event_if_
       (fun _ ->
         (* emit event only if we have some outputs or event subscribers *)
-        Observer.has_subscribers events
-        || not (CCList.is_empty (Atomic.get outputs)))
-      ~emit_ev:(fun ev ->
-        try Sync_queue.push q (T_emit ev) with Sync_queue.Closed -> ())
+        Atomic.get active
+        && (Observer.has_subscribers events
+           || not (CCList.is_empty (Atomic.get outputs))))
+      ~emit_ev
   in
 
-  let fence () =
-    let fut, prom = Moonpool.Fut.make () in
-    (try Sync_queue.push q (T_fence { wakeup = prom })
-     with Sync_queue.Closed -> Moonpool.Fut.fulfill prom @@ Ok ());
-    fut
+  let emit_fence () =
+    let outs = Atomic.get outputs in
+    List.iter (fun (out : Output.t) -> out.flush ()) outs
   in
 
-  fence_ := fence;
-  let self = { q; outputs; reporter; events; bg_thread = None } in
-  self.bg_thread <- Some (Thread.create bg_thread_loop_ self);
+  let self = { active; emit_ev; emit_fence; outputs; reporter; events } in
   self
 
+let[@inline] emit_ev (self : t) ev : unit = self.emit_ev ev
 let null () : t = to_outputs []
 
+let setup_ (self : t) : unit =
+  Logs.set_reporter self.reporter;
+  set_as_fence_ self;
+  ()
+
 let with_no_logger () f =
   let old = Logs.reporter () in
   Logs.set_reporter Logs.nop_reporter;
   Fun.protect ~finally:(fun () -> Logs.set_reporter old) f
 
 let setup_logger_to_stdout () =
   let outs = [ Output.stdout () ] in
-  let reporter = to_outputs outs in
-  Logs.set_reporter (as_reporter reporter)
+  let logger = to_outputs outs in
+  setup_ logger
 
 let setup_logger_to_stderr () =
   let outs = [ Output.stderr () ] in
-  let reporter = to_outputs outs in
-  Logs.set_reporter (as_reporter reporter)
+  let logger = to_outputs outs in
+  setup_ logger
 
 (** Setup a logger that emits into the file specified in ["LOG_FILE"] env,
     or no logger otherwise. *)
@@ -380,7 +408,8 @@ let setup_logger_to_LOG_FILE ?filename () k =
   | Some file, _ | None, Some file ->
     let@ oc = CCIO.with_out file in
     let outs = [ Output.to_chan oc ] in
-    Logs.set_reporter (to_outputs outs |> as_reporter);
+    let logger = to_outputs outs in
+    setup_ logger;
     k ()
   | _ -> k ()
 
@@ -390,9 +419,7 @@ module type LOG = sig
   val src : Logs.src
 end
 
-let fence () =
-  let fut = !fence_ () in
-  Moonpool.Fut.wait_block_exn fut
+let[@inline] fence () = !fence_ ()
 
 let mk_log_str s : (module LOG) =
   let src = Logs.Src.create s in

src/log/logger.mli

@@ -8,6 +8,9 @@ module Log_event = Log_event
 
     We can have multiple log outputs, they each get a {!Log_event.t}.
     See {!Reporter.to_outputs} to see how to use these.
+
+    A log output must be thread-safe, as it will be called from
+    multiple threads.
 *)
 module Output : sig
   type t
@@ -16,15 +19,20 @@ module Output : sig
 
   val stdout : unit -> t
   val stderr : unit -> t
-  val to_event : emit_ev:(Log_event.t -> unit) -> unit -> t
+
+  val to_event :
+    emit_ev:(Log_event.t -> unit) -> flush:(unit -> unit) -> unit -> t
+
   val filter_level : (level -> bool) -> t -> t
-  val to_str : emit_str:(string -> unit) -> unit -> t
+  val to_str : emit_str:(string -> unit) -> flush:(unit -> unit) -> unit -> t
 
-  val to_chan : out_channel -> t
-  (** Write into the channel, as text. *)
+  val to_chan : ?autoflush:bool -> out_channel -> t
+  (** Write into the channel, as text.
+      @param autoflush if true (default) channel is flushed after each event *)
 
-  val to_chan_jsonl : out_channel -> t
-  (** Write into the channel as jsonl. *)
+  val to_chan_jsonl : ?autoflush:bool -> out_channel -> t
+  (** Write into the channel as jsonl.
+      @param autoflush if true (default) channel is flushed after each event *)
 
   val buf_pool : Buffer.t Apool.t
   (** Buffer pool for loggers. Please hold onto buffers for only