Source file main.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2022 Nomadic Labs <contact@nomadic-labs.com>                *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Base

(* [Unix.error_message] is not available in JS, so we register the printer here
   instead of in [Base]. *)
let () =
  Printexc.register_printer @@ function
  | Unix.Unix_error (error, _, _) -> Some (Unix.error_message error)
  | _ -> None

let start_time = Unix.gettimeofday ()

module Memory_usage : sig
  val get_recursively_for_pid : int -> (int, string) result
end = struct
  let kb = 1000

  (* Get the memory usage of a single process (not recursively)
     by reading from /proc/PID/smaps.
     More specifically, we read the Pss field, which gives a more representative memory usage
     as it divides shared memory by the number of processes that share this memory. *)
  (* https://stackoverflow.com/questions/131303/how-can-i-measure-the-actual-memory-usage-of-an-application-or-process *)
  let get_total_pss pid =
    let ch = open_in ("/proc/" ^ string_of_int pid ^ "/smaps") in
    Fun.protect ~finally:(fun () -> close_in ch) @@ fun () ->
    let rec sum_pss acc =
      match input_line ch with
      | exception End_of_file -> acc
      | line ->
          let parse_error () =
            failwith (Printf.sprintf "failed to parse %S" line)
          in
          let line = String.trim line in
          let value =
            if not (String.starts_with ~prefix:"Pss:" line) then 0
            else if not (String.ends_with ~suffix:"kB" line) then parse_error ()
            else
              let value =
                String.sub
                  line
                  (String.length "Pss:")
                  (String.length line - String.length "Pss:"
                 - String.length "kB")
                |> String.trim
              in
              match int_of_string_opt value with
              | None -> parse_error ()
              | Some value -> value * kb
          in
          sum_pss (acc + value)
    in
    sum_pss 0

  (* Read all lines from an input channel and return them as a list. *)
  let rec input_lines ?(acc = []) ch =
    match input_line ch with
    | exception End_of_file -> List.rev acc
    | line -> input_lines ~acc:(line :: acc) ch

  (* Read /proc/PID/task/*/children to get the list of children of a process.
     Returns the list of PIDs of all children. *)
  let get_children pid =
    (* I think a task is the same as a thread. *)
    let task_path = "/proc/" ^ string_of_int pid ^ "/task" in
    let parse_tid line =
      match int_of_string_opt (String.trim line) with
      | None ->
          failwith (sf "failed to parse %S from %s as a TID" line task_path)
      | Some pid -> pid
    in
    let task_ids =
      Sys.readdir task_path |> Array.to_list |> List.map parse_tid
    in
    let get_task_children tid =
      let task_children_path =
        task_path ^ "/" ^ string_of_int tid ^ "/children"
      in
      let ch = open_in task_children_path in
      Fun.protect ~finally:(fun () -> close_in ch) @@ fun () ->
      let parse_pid line =
        match int_of_string_opt line with
        | None ->
            failwith
              (sf "failed to parse %S from %s as a PID" line task_children_path)
        | Some pid -> pid
      in
      input_lines ch
      |> List.map (String.split_on_char ' ')
      |> List.flatten |> List.map String.trim
      |> List.filter (( <> ) "")
      |> List.map parse_pid
    in
    List.flatten (List.map get_task_children task_ids)

  let get_recursively_for_pid pid =
    let rec get acc pid =
      let acc = acc + get_total_pss pid in
      let children = get_children pid in
      List.fold_left get acc children
    in
    try Ok (get 0 pid)
    with Sys_error message | Failure message -> Error message
end

module Scheduler : Test.SCHEDULER = struct
  (* Use the [Scheduler] module from the [tezt.scheduler] library,
     with an alias [S] so that we don't confuse it with the [Scheduler] module
     we are defining. *)
  module S = Scheduler

  type request = Run_test of {test_title : string}

  type response = Test_result of Test.test_result

  let internal_worker_error x =
    Printf.ksprintf
      (fun s ->
        Log.error "internal error in worker: %s" s ;
        exit 1)
      x

  let internal_scheduler_error x =
    Printf.ksprintf
      (fun s ->
        Log.error "internal error in scheduler: %s" s ;
        exit 1)
      x

  let encode_used_seed (seed : Test.used_seed) : S.Message.value =
    match seed with Used_fixed -> Unit | Used_random seed -> Int seed

  let decode_used_seed (value : S.Message.value) : Test.used_seed =
    match value with
    | Unit -> Used_fixed
    | Int seed -> Used_random seed
    | _ -> raise (S.Message.Failed_to_decode (value, "used_seed"))

  let test_result_type : Test.test_result S.Message.typ =
    let encode ({test_result; seed; peak_memory_usage} : Test.test_result) :
        S.Message.value =
      Block
        [|
          (match test_result with
          | Successful -> Int 0
          | Failed error_message -> String error_message
          | Aborted -> Int 1);
          encode_used_seed seed;
          (match peak_memory_usage with
          | None -> Int (-1)
          | Some peak_memory_usage -> Int peak_memory_usage);
        |]
    in
    let decode (value : S.Message.value) : Test.test_result =
      let fail () = raise (S.Message.Failed_to_decode (value, "test_result")) in
      match value with
      | Block [|test_result; seed; Int peak_memory_usage|] ->
          {
            test_result =
              (match test_result with
              | Int 0 -> Successful
              | String error_message -> Failed error_message
              | Int 1 -> Aborted
              | _ -> fail ());
            seed = decode_used_seed seed;
            peak_memory_usage =
              (if peak_memory_usage >= 0 then Some peak_memory_usage else None);
          }
      | _ -> fail ()
    in
    S.Message.typ ~encode ~decode

  let run_test ?(temp_start = Temp.start) (test_title : string) :
      Test.test_result =
    match Test.get_test_by_title test_title with
    | None ->
        internal_worker_error
          "scheduler requested to run test %S, but worker doesn't know about \
           this test"
          test_title
    | Some test ->
        let clean_up () =
          Lwt.catch Process.clean_up @@ fun exn ->
          Log.warn "Failed to clean up processes: %s" (Printexc.to_string exn) ;
          unit
        in
        (* If the following raises an exception, it will be caught by the scheduler
           library and sent as a string to the scheduler process.
           Any exception-specific handling is better done in the [Test] module
           so that other backends can also benefit. *)
        Lwt_main.run
        @@ Test.run_one
             ~sleep:Lwt_unix.sleep
             ~clean_up
             ~temp_start
             ~temp_stop:Temp.stop
             ~temp_clean_up:Temp.clean_up
             test

  let msg_seed = S.Message.(register int) "Seed"

  let add_task_for_request (Run_test {test_title})
      (on_response : response -> unit) =
    let term_timeout =
      let global_timeout =
        match Cli.Options.global_timeout with
        | None -> None
        | Some global_timeout ->
            Some (max 0. (start_time +. global_timeout -. Unix.gettimeofday ()))
      in
      match (global_timeout, Cli.Options.test_timeout) with
      | None, None -> None
      | (Some _ as x), None | None, (Some _ as x) -> x
      | Some a, Some b -> Some (min a b)
    in
    let kill_timeout =
      match term_timeout with
      | None -> None
      | Some _ -> Some Cli.Options.cleanup_timeout
    in
    let warn_after_timeout_timer : S.Timer.t option ref = ref None in
    let measure_memory_usage_timer : S.Timer.t option ref = ref None in
    let peak_memory_usage = ref None in
    let already_warned_about_memory_usage = ref false in
    let on_start ctx =
      let measure_memory_usage_once pid =
        match Memory_usage.get_recursively_for_pid pid with
        | Ok bytes -> (
            if
              (not !already_warned_about_memory_usage)
              &&
              match Cli.Options.mem_warn with
              | None -> false
              | Some threshold -> bytes >= threshold
            then (
              already_warned_about_memory_usage := true ;
              Log.warn
                "Test is using too much memory: %S (%d bytes)"
                test_title
                bytes) ;
            match !peak_memory_usage with
            | None -> peak_memory_usage := Some bytes
            | Some old -> peak_memory_usage := Some (max old bytes))
        | Error _ -> ()
      in
      let rec warn_after_timeout delay =
        warn_after_timeout_timer :=
          Some
            ( S.Timer.on_delay delay @@ fun () ->
              Log.warn "Test is still running: %S" test_title ;
              warn_after_timeout delay )
      in
      let rec measure_memory_usage delay =
        measure_memory_usage_timer :=
          Some
            ( S.Timer.on_delay delay @@ fun () ->
              match S.get_runner_pid ctx with
              | None -> ()
              | Some pid ->
                  measure_memory_usage_once pid ;
                  measure_memory_usage delay )
      in
      if Cli.Options.mem_warn <> None || Cli.Reports.record_mem_peak then
        measure_memory_usage (1. /. Cli.Options.mem_poll_frequency) ;
      if Cli.Options.warn_after_timeout > 0. then
        warn_after_timeout Cli.Options.warn_after_timeout
    in
    let random_seed = ref None in
    let on_message _ctx message =
      S.Message.match_with message msg_seed ~default:(fun () -> ())
      @@ fun seed -> random_seed := Some seed
    in
    let on_finish result =
      Option.iter S.Timer.cancel !warn_after_timeout_timer ;
      Option.iter S.Timer.cancel !measure_memory_usage_timer ;
      let test_result : Test.test_result =
        (match (Cli.Options.mem_warn, !peak_memory_usage) with
        | None, _ | _, None -> ()
        | Some threshold, Some peak ->
            if peak >= threshold then
              Log.warn
                "Test used too much memory: %S (peak: %d bytes)"
                test_title
                peak) ;
        match result with
        | Ok test_result ->
            {test_result with peak_memory_usage = !peak_memory_usage}
        | Error error_message ->
            Log.error "%s" error_message ;
            let seed : Test.used_seed =
              match !random_seed with
              | None -> Used_fixed
              | Some seed -> Used_random seed
            in
            {
              test_result = Failed error_message;
              seed;
              peak_memory_usage = !peak_memory_usage;
            }
      in
      on_response (Test_result test_result)
    in
    let on_term_timeout () =
      Log.error
        "Sending SIGTERM to test which is taking too long: %S"
        test_title
    in
    let on_kill_timeout () =
      Log.error
        "Sending SIGKILL to test which is taking too long to stop: %S"
        test_title
    in
    S.add_task
      ?term_timeout
      ?kill_timeout
      ~on_term_timeout
      ~on_kill_timeout
      ~on_start
      ~on_message
      ~on_finish
      test_result_type
    @@ fun ctx ->
    (* Abuse [~temp_start] to insert something to execute in the context of the test itself.
       This works because [temp_start] is executed just before the test function;
       more precisely, it is executed after the seed is chosen. *)
    let temp_start () =
      (* If the test is using a random seed, the main process will need the seed
         to tell the user how to reproduce in case of error. *)
      (match Test.current_test_seed_specification () with
      | Fixed _ -> ()
      | Random ->
          S.Message.send_to_scheduler ctx msg_seed (Test.current_test_seed ())) ;
      Temp.start ()
    in
    (* Find the test to run, run it and return with the test result. *)
    let test_result = run_test ~temp_start test_title in
    (* Don't leave logs, they could be printed in case of error in the next test
       (or in case of error in an [at_exit]). This is because it is the scheduler
       that prints the test report which normally clears logs in single process. *)
    Log.clear_error_context_queue () ;
    test_result

  let next_worker_id = ref 0

  let current_worker_id = ref None

  let run_multi_process
      ~(on_worker_available : unit -> (request * (response -> unit)) option)
      ~worker_count =
    (* Register signal handlers meant for the scheduler process.
       Doing that overrides existing handlers (i.e. those defined at toplevel
       in the [Test] module). The latter are meant for worker processes.
       So we store them in [old_sigint_behavior] and [old_sigterm_behavior]
       in order to restore them in worker processes. *)
    (* SIGINT is received in particular when the user presses Ctrl+C.
       In that case terminals tend to send SIGINT to the process group,
       so workers should receive SIGINT as well.
       The scheduler shall stop starting new tasks,
       but should not send SIGINT because that would be a duplicate
       and tests would think that the user pressed Ctrl+C twice to force a kill. *)
    let old_sigint_behavior =
      Sys.(signal sigint)
        (Signal_handle
           (fun _ ->
             (* If the user presses Ctrl+C again, let the program die immediately. *)
             Sys.(set_signal sigint) Signal_default ;
             S.stop ()))
    in
    (* SIGTERM is usually received programmatically.
       Contrary to SIGINT, we do not clear the signal handler:
       the equivalent of "a second SIGTERM to force quit" is SIGKILL. *)
    let old_sigterm_behavior =
      Sys.(signal sigterm) (Signal_handle (fun _ -> S.stop ()))
    in
    let on_empty_queue () =
      match on_worker_available () with
      | None -> ()
      | Some (request, on_response) -> add_task_for_request request on_response
    in
    let on_worker_kill_timeout () =
      Log.debug
        "Send SIGKILL to worker which is still running despite not running any \
         test"
    in
    let on_unexpected_worker_exit status =
      Log.warn
        "Worker %s while not running a test"
        (S.show_process_status status)
    in
    let fork () =
      let worker_id = !next_worker_id in
      incr next_worker_id ;
      let pid = Lwt_unix.fork () in
      if pid = 0 then (
        (* This is a child process.
           Reset stuff that only makes sense in the scheduler (signal handlers),
           and set stuff that is worker-specific. *)
        Sys.(set_signal sigint) old_sigint_behavior ;
        Sys.(set_signal sigterm) old_sigterm_behavior ;
        Temp.set_pid (Unix.getpid ()) ;
        current_worker_id := Some !next_worker_id ;
        Log.set_current_worker_id worker_id ;
        Log.clear_error_context_queue () ;
        Option.iter
          (fun filename ->
            (* Workers log to [BASENAME-WORKER_ID.EXT] *)
            let worker_filename =
              Filename.(
                concat
                  (dirname filename)
                  (remove_extension (basename filename)
                  ^ "-" ^ string_of_int worker_id ^ extension filename))
            in
            Log.set_file worker_filename)
          Cli.Logs.file) ;
      pid
    in
    S.run
      ~worker_idle_timeout:Cli.Options.cleanup_timeout
      ~worker_kill_timeout:Cli.Options.cleanup_timeout
      ~on_worker_kill_timeout
      ~on_empty_queue
      ~on_unexpected_worker_exit
      ~fork
      worker_count

  let rec run_single_process ~on_worker_available =
    Temp.set_pid (Unix.getpid ()) ;
    match on_worker_available () with
    | None -> ()
    | Some (Run_test {test_title}, on_response) ->
        let test_result = run_test test_title in
        on_response (Test_result test_result) ;
        run_single_process ~on_worker_available

  let run ~on_worker_available ~worker_count continue =
    (match worker_count with
    | None -> run_single_process ~on_worker_available
    | Some worker_count -> (
        try run_multi_process ~on_worker_available ~worker_count
        with exn -> internal_scheduler_error "%s" (Printexc.to_string exn))) ;
    continue ()

  let get_current_worker_id () = !current_worker_id
end

let run () = Test.run_with_scheduler (module Scheduler)