talircd/lib/logging/logging.ml

type level =
  | TRACE (* 0 *)
  | DEBUG (* 1 *)
  | INFO  (* 2 *)
  | WARN  (* 3 *)
  | ERROR (* 4 *)

external int_of_level : level -> int = "%identity"

module Time = struct
  (* TODO: abstract this? *)
  type t = Ptime.t
  let stamp = Ptime_clock.now
  let pp = Ptime.pp_rfc3339 ()
             ?tz_offset_s:(Ptime_clock.current_tz_offset_s ())
  let to_string x =
    Format.asprintf "%a" pp x
end

type writer =
  ts:Time.t ->
  ns:string ->
  lvl:level ->
  string ->
  unit

let ( +++ ) w1 w2 ~ts ~ns ~lvl msg =
  w1 ~ts ~ns ~lvl msg;
  w2 ~ts ~ns ~lvl msg

type logger = {
  namespace : string;
  mutable child : logger;
  sibling : logger;
  mutable min_level : int;
  mutable writers : writer option;
}

let rec root_logger = {
  namespace = "";
  min_level = 0;
  child = root_logger;
  sibling = root_logger;
  writers = None;
}

let make_logger parent_logger ns =
  let namespace =
    if parent_logger == root_logger then ns
    else String.concat "." [parent_logger.namespace; ns];
  in
  let rec logger = {
    namespace;
    min_level = Int.max_int;
    child = logger;
    sibling = parent_logger.child;
    writers = None;
  } in
  parent_logger.child <- logger;
  logger

type 'a log_function =
  ((('a, Format.formatter, unit) format -> 'a) -> unit) -> unit

let write logger lvl msg =
  let ts = Time.stamp () in
  let ns = logger.namespace in
  Option.iter (fun w -> w ~ts ~ns ~lvl msg) logger.writers

let[@inline] logf logger lvl : _ log_function = fun logk ->
  if int_of_level lvl >= logger.min_level then
    logk (fun fmt ->
           Format.kasprintf (write logger lvl)
             ("@[<hov>" ^^ fmt))


module type Logs = sig
  val logger : logger
  val trace : _ log_function
  val debug : _ log_function
  val info  : _ log_function
  val warn  : _ log_function
  val error : _ log_function
  val log   : level -> _ log_function
end

let make_logs parent_logger ns =
  (module struct
    let logger = make_logger parent_logger ns
    let[@inline] trace k = logf logger TRACE k
    let[@inline] debug k = logf logger DEBUG k
    let[@inline] info  k = logf logger INFO  k
    let[@inline] warn  k = logf logger WARN  k
    let[@inline] error k = logf logger ERROR k
    let[@inline] log level k = logf logger level k
  end : Logs)

let logs ns = make_logs root_logger ns
let sublogs logger ns = make_logs logger ns

(* TODO: filters: namespace[*] => min_level *)

let rec add_writer_rec min_level writer parent_logger logger =
  if parent_logger != logger then
    begin
      logger.min_level <- min min_level logger.min_level;
      logger.writers <-
        (match logger.writers with
         | None -> Some writer
         | Some writer' -> Some (writer' +++ writer));
      add_writer_rec min_level writer logger logger.child;
      add_writer_rec min_level writer parent_logger logger.sibling
    end

let[@inline] add_writer ?(min_level = WARN) writer =
  add_writer_rec (int_of_level min_level) writer root_logger root_logger.child

module Pretty = struct
  type t = {
    mutex : Mutex.t;
    bp : Buffer.t;
    mutable align_to : int;
    out : out_channel;
    flags : int;
  }

  type config = {
    color : bool;
    timestamp : bool;
    namespace : bool;
    level : bool;
  }

  let _fl_c = 1 (* color *)
  let _fl_t = 2 (* timestamp *)
  let _fl_n = 4 (* namespace *)
  let _fl_l = 8 (* level *)
  let _fl_ct = _fl_c + _fl_t
  let _fl_cl = _fl_c + _fl_l

  let header = function
    | TRACE -> "TRACE"
    | DEBUG -> "DEBUG"
    | INFO  -> "INFO"
    | WARN  -> "WARN"
    | ERROR -> "ERROR"

  let ansi_header = function
    | TRACE -> "\x1b[34m"
    | DEBUG -> "\x1b[36m"
    | INFO  -> "\x1b[32m"
    | WARN  -> "\x1b[33m"
    | ERROR -> "\x1b[31m"

  let ansi_dim  = "\x1b[2m"
  let ansi_bold = "\x1b[1m"
  let ansi_off  = "\x1b[0m"

  let pr_timestamp bp ts =
    Buffer.add_string bp (Time.to_string ts)

  let pr_spaces bp n =
    for _ = 1 to n do
      Buffer.add_char bp ' '
    done

  let rec pr_lines_rec indent bp s i =
    match String.index_from s i '\n' with
    | exception Not_found ->
      Buffer.add_substring bp s i (String.length s - i)
    | k ->
      Buffer.add_substring bp s i (k + 1 - i);
      pr_spaces bp indent;
      pr_lines_rec indent bp s (k + 1)

  let pr_lines ~indent bp s =
    pr_lines_rec indent bp s 0

  let make (out : out_channel) (cfg : config) = {
    mutex = Mutex.create ();
    bp = Buffer.create 512;
    align_to = 0;
    out;
    flags = (if cfg.color     then _fl_c else 0) +
            (if cfg.timestamp then _fl_t else 0) +
            (if cfg.namespace then _fl_n else 0) +
            (if cfg.level     then _fl_l else 0)
  }

  let pr_msg bp msg ts ns lvl ~align ~indent ~f =
    begin
      if f _fl_ct then Buffer.add_string bp ansi_dim;
      if f _fl_t  then Printf.bprintf bp "%a " pr_timestamp ts;
      if f _fl_ct then Buffer.add_string bp ansi_off;

      if f _fl_n then Printf.bprintf bp "%s " ns;
      if f _fl_n then pr_spaces bp (align bp);

      if f _fl_cl then Buffer.add_string bp (ansi_header lvl);
      if f _fl_l then Printf.bprintf bp "%-5s " (header lvl);
      if f _fl_cl then Buffer.add_string bp ansi_off;

      if f _fl_c then Buffer.add_string bp ansi_bold;
      pr_lines bp msg ~indent:(indent bp);
      if f _fl_c then Buffer.add_string bp ansi_off;

      Buffer.add_string bp "\n";
    end

  let writer t ~ts ~ns ~lvl msg =
    let f mask = t.flags land mask = mask in
    let align bp =
      let n = t.align_to - Buffer.length bp in
      t.align_to <- Buffer.length bp;
      n
    in
    let indent bp =
      Buffer.length bp
      - (if f _fl_c  then 4 else 0)
      - (if f _fl_ct then 8 else 0)
      - (if f _fl_cl then 9 else 0)
    in
    begin
      Mutex.lock t.mutex;
      pr_msg t.bp msg ts ns lvl ~align ~indent ~f;
      Buffer.output_buffer t.out t.bp;
      flush t.out;
      Buffer.clear t.bp;
      Mutex.unlock t.mutex
    end
end

let init_pretty_writer
      ?min_level
      ?(color = true)
      ?(timestamp = true)
      ?(namespace = true)
      ?(level = true)
      out
  =
  Pretty.make out { color; timestamp; namespace; level } |>
  Pretty.writer |>
  add_writer ?min_level