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 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} : 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;
        |]
    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|] ->
          {
            test_result =
              (match test_result with
              | Int 0 -> Successful
              | String error_message -> Failed error_message
              | Int 1 -> Aborted
              | _ -> fail ());
            seed = decode_used_seed seed;
          }
      | _ -> 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 on_start _ctx =
      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
      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 ;
      let test_result : Test.test_result =
        match result with
        | Ok test_result -> test_result
        | 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}
      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 =
    (if worker_count = 1 then run_single_process ~on_worker_available
     else
       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)