(*
 * Copyright (c) Facebook, Inc. and its 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_exit

let () = Callback.register_exception "worker_should_exit" Worker_should_exit

external stop_workers : unit -> unit = "hh_stop_workers"

external resume_workers : unit -> unit = "hh_resume_workers"

external check_should_exit : unit -> unit = "hh_check_should_exit"

external set_can_worker_stop : bool -> unit = "hh_set_can_worker_stop"

let on_worker_cancelled = ref (fun () -> ())

let set_on_worker_cancelled f = on_worker_cancelled := f

let with_no_cancellations f =
  Utils.try_finally
    ~f:
      begin
        fun () ->
        set_can_worker_stop false;
        f ()
      end
    ~finally:(fun () -> set_can_worker_stop true)

let with_worker_exit f =
  try f ()
  with Worker_should_exit ->
    !on_worker_cancelled ();
    exit 0

(* Check if the workers are stopped and exit if they are *)
let check_should_exit () = with_worker_exit check_should_exit

(*
 * Copyright (c) Facebook, Inc. and its 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_exit

let () = Callback.register_exception "worker_should_exit" Worker_should_exit

external stop_workers : unit -> unit = "hh_stop_workers"

external resume_workers : unit -> unit = "hh_resume_workers"

external check_should_exit : unit -> unit = "hh_check_should_exit"

external set_can_worker_stop : bool -> unit = "hh_set_can_worker_stop"

let on_worker_cancelled = ref (fun () -> ())

let set_on_worker_cancelled f = on_worker_cancelled := f

let with_worker_exit f =
  try f ()
  with Worker_should_exit ->
    !on_worker_cancelled ();
    exit 0

(* Check if the workers are stopped and exit if they are *)
let check_should_exit () = with_worker_exit check_should_exit

let with_no_cancellations f =
  Utils.try_finally
    ~f:
      begin
        fun () ->
        set_can_worker_stop false;
        f ()
      end
    ~finally:(fun () ->
      set_can_worker_stop true;
      check_should_exit ())

Add check_should_exit to the end of with_no_cancellations
Add check_should_exit to the end of with_no_cancellations

Summary:
If you want a group of operations that can touch the shared heaps to be atomic, you can wrap them with `with_no_cancellations` to prevent cancellations from taking place during that operation. Unfortunately, if many of our shared heap operations are wrapped in this utility, there's the potential for cancellations to be delayed for a long time, with potential negative performance ramifications.

So, let's check if we need to cancel after each block of atomic operations.

Note that I had to rearrange the functions a bit so that the version of `check_should_exit` that exits cleanly is in scope.

Reviewed By: samwgoldman

Differential Revision: D24095226

fbshipit-source-id: 59e340023d8ffe324ef1bb2319cfda2b24eb6b76
