Source file countdown_latch.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
open! Base
open! Import

(** See [Adaptive_backoff.once] *)
let log_scale = 10

module State : sig
  type t : immediate

  val max_count : int

  (* *)

  val from_count : int -> t

  (* *)

  val count : t -> int
  val is_poisoned : t -> bool
  val is_0_or_less : t -> bool
  val is_1_or_less : t -> bool

  (* *)

  val and_poisoned : t -> t
  val and_incr : t -> t

  (* *)

  val fetch_and_add : t Awaitable.t @ local -> int -> t
end = struct
  type t = int

  let max_count = 1 lsl (Int.num_bits - 2)
  let poisoned_bit = 1
  let count_shift = 1
  let from_count n = n lsl count_shift
  let count t = if t < 0 then 0 else t lsr count_shift
  let is_poisoned t = t land poisoned_bit <> 0
  let and_poisoned t = t lor poisoned_bit
  let is_0_or_less t = t <= (0 lsl count_shift) lor poisoned_bit
  let is_1_or_less t = t <= (1 lsl count_shift) lor poisoned_bit
  let and_incr t = t + (1 lsl count_shift)
  let fetch_and_add t d = Awaitable.fetch_and_add t (d lsl count_shift)
end

type t = State.t Awaitable.t

let sexp_of_t t =
  let s = Awaitable.get t in
  let c = State.count s in
  [%sexp
    (if State.is_poisoned s then `poisoned c else `count c
     : [ `poisoned of int | `count of int ])]
;;

let max_count = State.max_count

let create ?padded count =
  if count < 0 || max_count < count
  then invalid_arg "Countdown_latch.create: invalid initial count"
  else Awaitable.make ?padded (State.from_count count)
;;

let rec poison t =
  let before = Awaitable.get t in
  if not (State.is_poisoned before)
  then (
    let after = State.and_poisoned before in
    match Awaitable.compare_and_set t ~if_phys_equal_to:before ~replace_with:after with
    | Set_here -> Awaitable.broadcast t
    | Compare_failed -> poison t)
;;

let is_poisoned t = State.is_poisoned (Awaitable.get t)

let[@inline never] already_reached_0 t d =
  let _ : _ = State.fetch_and_add t d in
  invalid_arg "Countdown_latch.decr: already reached zero"
;;

let incr t =
  let[@inline] rec loop () =
    let before = Awaitable.get t in
    if State.is_poisoned before
    then raise Poisoned
    else if State.is_0_or_less before
    then already_reached_0 t 0
    else (
      let after = State.and_incr before in
      match Awaitable.compare_and_set t ~if_phys_equal_to:before ~replace_with:after with
      | Set_here -> ()
      | Compare_failed ->
        Adaptive_backoff.once ~random_key:(Awaitable.random_key t) ~log_scale;
        loop ())
  in
  loop () [@nontail]
;;

let decr t =
  let before_decr = State.fetch_and_add t (-1) in
  if State.is_1_or_less before_decr
  then
    if State.is_0_or_less before_decr
    then already_reached_0 t 1
    else Awaitable.broadcast t
;;

let rec await w t =
  let state = Awaitable.get t in
  if State.is_poisoned state
  then raise Poisoned
  else if not (State.is_0_or_less state)
  then (
    match Awaitable.await w t ~until_phys_unequal_to:state with
    | Signaled -> await w t
    | Terminated -> raise Await.Terminated)
;;

let rec await_or_cancel w c t : _ Or_canceled.t =
  let state = Awaitable.get t in
  if State.is_poisoned state
  then raise Poisoned
  else if State.is_0_or_less state
  then Completed ()
  else (
    match Awaitable.await_or_cancel w c t ~until_phys_unequal_to:state with
    | Signaled -> await_or_cancel w c t
    | Terminated -> raise Await.Terminated
    | Canceled -> Canceled)
;;

let count t = State.count (Awaitable.get t)