Source file mpmc_queue.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
open! Base
open! Import

type 'a t =
  { queue : 'a aliased_many Portable_mpmc_queue.t
  ; awaitable : int Awaitable.t
  }

let create ?padded () =
  let awaitable = Awaitable.make ?padded 0 in
  let queue = Portable_mpmc_queue.create ?padded () in
  { queue; awaitable }
;;

let push t a =
  Portable_mpmc_queue.push t.queue { aliased_many = a };
  Awaitable.incr t.awaitable;
  Awaitable.signal t.awaitable
;;

type state =
  | Never_awaited
  | Signaled

let pop await (t @ local) =
  let[@inline] rec loop state =
    match Portable_mpmc_queue.pop t.queue with
    | This v ->
      (match state with
       | Signaled -> Awaitable.signal t.awaitable
       | Never_awaited -> ());
      v
    | Null ->
      let before = Awaitable.get t.awaitable in
      (match Portable_mpmc_queue.pop t.queue with
       | This v ->
         (match state with
          | Signaled -> Awaitable.signal t.awaitable
          | Never_awaited -> ());
         v
       | Null ->
         (match Awaitable.await await t.awaitable ~until_phys_unequal_to:before with
          | Terminated -> raise Await.Terminated
          | Signaled -> loop Signaled))
  in
  (loop Never_awaited).aliased_many
;;

let pop_or_cancel await c t : _ Or_canceled.t =
  let[@inline] rec loop state : _ Or_canceled.t =
    match Portable_mpmc_queue.pop t.queue with
    | This v ->
      (match state with
       | Signaled -> Awaitable.signal t.awaitable
       | Never_awaited -> ());
      Completed v
    | Null ->
      let before = Awaitable.get t.awaitable in
      (match Portable_mpmc_queue.pop t.queue with
       | This v ->
         (match state with
          | Signaled -> Awaitable.signal t.awaitable
          | Never_awaited -> ());
         Completed v
       | Null ->
         (match
            Awaitable.await_or_cancel await c t.awaitable ~until_phys_unequal_to:before
          with
          | Terminated -> raise Await.Terminated
          | Canceled -> Canceled
          | Signaled -> loop Signaled))
  in
  match loop Never_awaited with
  | Completed { aliased_many } -> Completed aliased_many
  | Canceled -> Canceled
;;

let pop_nonblocking t =
  match Portable_mpmc_queue.pop t.queue with
  | This { aliased_many } -> This aliased_many
  | Null -> Null
;;

let peek t =
  match Portable_mpmc_queue.peek t.queue with
  | This { aliased_many } -> This aliased_many
  | Null -> Null
;;

let[@inline] length t = Portable_mpmc_queue.length t.queue

module For_testing = struct
  let length t = Awaitable.For_testing.length t.awaitable
end