(*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

exception Worker_should_cancel

let () = Callback.register_exception "worker_should_cancel" Worker_should_cancel

external stop_workers : unit -> unit = "hh_stop_workers"

external resume_workers : unit -> unit = "hh_resume_workers"

external check_should_cancel : unit -> unit = "hh_check_should_cancel"

external set_can_worker_stop : bool -> unit = "hh_set_can_worker_stop"

external get_can_worker_stop : unit -> bool = "hh_get_can_worker_stop"

let with_no_cancellations f =
  if get_can_worker_stop () then
    Exception.protect
      ~f:
        begin
          fun () ->
            set_can_worker_stop false;
            f ()
        end
      ~finally:(fun () ->
        set_can_worker_stop true;
        check_should_cancel ())
  else
    (* If can_worker_stop is set to false, then that means we are already within an enclosing
     * call to with_no_cancellations. So, we can just call `f` directly. Note that a more
     * sophisticated approach will be necessary if we involve Lwt here, since then calls to
     * with_no_cancellations could be interleaved without being nested. *)
    f ()
