this repo has no description
gdiazlo.tngl.sh
1open Types
2
3type t = {
4 config : config;
5 self : node_info;
6 members : Membership.t;
7 incarnation : incarnation Kcas.Loc.t;
8 sequence : int Kcas.Loc.t;
9 broadcast_queue : Dissemination.t;
10 pending_acks : Pending_acks.t;
11 probe_index : int Kcas.Loc.t;
12 send_pool : Buffer_pool.t;
13 recv_pool : Buffer_pool.t;
14 udp_sock : [ `Generic ] Eio.Net.datagram_socket_ty Eio.Resource.t;
15 event_stream : node_event Eio.Stream.t;
16 user_handlers : (node_info -> string -> string -> unit) list Kcas.Loc.t;
17 cipher_key : Crypto.key;
18 stats : stats Kcas.Loc.t;
19 shutdown : bool Kcas.Loc.t;
20 clock : float Eio.Time.clock_ty Eio.Resource.t;
21 mono_clock : Eio.Time.Mono.ty Eio.Resource.t;
22 secure_random : Eio.Flow.source_ty Eio.Resource.t;
23}
24
25let next_seq t =
26 Kcas.Xt.commit
27 {
28 tx =
29 (fun ~xt ->
30 let seq = Kcas.Xt.get ~xt t.sequence in
31 Kcas.Xt.set ~xt t.sequence (seq + 1);
32 seq);
33 }
34
35let get_incarnation t =
36 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.get ~xt t.incarnation) }
37
38let incr_my_incarnation t =
39 Kcas.Xt.commit
40 {
41 tx =
42 (fun ~xt ->
43 let inc = Kcas.Xt.get ~xt t.incarnation in
44 let new_inc = incr_incarnation inc in
45 Kcas.Xt.set ~xt t.incarnation new_inc;
46 new_inc);
47 }
48
49let is_shutdown t =
50 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.get ~xt t.shutdown) }
51
52let now_mtime t =
53 Eio.Time.Mono.now t.mono_clock
54 |> Mtime.to_uint64_ns |> Mtime.Span.of_uint64_ns
55
56let update_stats t f =
57 Kcas.Xt.commit
58 {
59 tx =
60 (fun ~xt ->
61 let s = Kcas.Xt.get ~xt t.stats in
62 Kcas.Xt.set ~xt t.stats (f s));
63 }
64
65let emit_event t ev = Eio.Stream.add t.event_stream ev
66
67let send_packet t ~dst (packet : packet) =
68 Buffer_pool.with_buffer t.send_pool (fun buf ->
69 match Codec.encode_packet packet ~buf with
70 | Error `Buffer_too_small -> ()
71 | Ok encoded_len ->
72 let encoded = Cstruct.sub buf 0 encoded_len in
73 let encrypted =
74 Crypto.encrypt ~key:t.cipher_key ~random:t.secure_random encoded
75 in
76 Transport.send_udp t.udp_sock dst encrypted;
77 update_stats t (fun s -> { s with msgs_sent = s.msgs_sent + 1 }))
78
79let make_packet t ~primary ~piggyback =
80 { cluster = t.config.cluster_name; primary; piggyback }
81
82let drain_piggyback t ~max_bytes =
83 Dissemination.drain t.broadcast_queue ~max_bytes
84 ~encode_size:Codec.encoded_size
85
86let enqueue_broadcast t msg =
87 let transmits =
88 Protocol_pure.retransmit_limit t.config
89 ~node_count:(Membership.count t.members)
90 in
91 Dissemination.enqueue t.broadcast_queue msg ~transmits ~created:(now_mtime t);
92 Dissemination.invalidate t.broadcast_queue
93 ~invalidates:Protocol_pure.invalidates msg
94
95let handle_ping t ~src (ping : protocol_msg) =
96 match ping with
97 | Ping { seq; sender = _ } ->
98 let piggyback =
99 drain_piggyback t ~max_bytes:(t.config.udp_buffer_size - 100)
100 in
101 let ack = Ack { seq; responder = t.self; payload = None } in
102 let packet = make_packet t ~primary:ack ~piggyback in
103 send_packet t ~dst:src packet
104 | _ -> ()
105
106let handle_ping_req t ~src:_ (ping_req : protocol_msg) =
107 match ping_req with
108 | Ping_req { seq; target; sender = _ } -> (
109 match Membership.find t.members target with
110 | None -> ()
111 | Some member ->
112 let target_addr = (Membership.Member.node member).addr in
113 let ping = Ping { seq; sender = t.self } in
114 let packet = make_packet t ~primary:ping ~piggyback:[] in
115 send_packet t ~dst:target_addr packet)
116 | _ -> ()
117
118let handle_ack t (ack : protocol_msg) =
119 match ack with
120 | Ack { seq; responder = _; payload } ->
121 ignore (Pending_acks.complete t.pending_acks ~seq ~payload)
122 | _ -> ()
123
124let apply_member_transition t member_id transition_fn =
125 let now = now_mtime t in
126 match Membership.find t.members member_id with
127 | None -> ()
128 | Some member ->
129 let snap = Membership.Member.snapshot_now member in
130 let transition = transition_fn snap ~now in
131 if transition.Protocol_pure.new_state.state <> snap.state then begin
132 Membership.update_member t.members member_id
133 {
134 update =
135 (fun m ~xt ->
136 match transition.new_state.state with
137 | Alive ->
138 Membership.Member.set_alive ~xt m
139 ~incarnation:transition.new_state.incarnation ~now
140 | Suspect ->
141 Membership.Member.set_suspect ~xt m
142 ~incarnation:transition.new_state.incarnation ~now
143 | Dead ->
144 Membership.Member.set_dead ~xt m
145 ~incarnation:transition.new_state.incarnation ~now);
146 }
147 |> ignore
148 end;
149 List.iter (fun msg -> enqueue_broadcast t msg) transition.broadcasts;
150 List.iter (emit_event t) transition.events
151
152let handle_alive_msg t (msg : protocol_msg) =
153 match msg with
154 | Alive { node; incarnation = _ } ->
155 apply_member_transition t node.id (fun snap ~now ->
156 Protocol_pure.handle_alive ~self:t.self.id snap msg ~now)
157 | _ -> ()
158
159let handle_suspect_msg t (msg : protocol_msg) =
160 match msg with
161 | Suspect { node; incarnation = _; suspector = _ } ->
162 apply_member_transition t node (fun snap ~now ->
163 Protocol_pure.handle_suspect ~self:t.self.id snap msg ~now)
164 | _ -> ()
165
166let handle_dead_msg t (msg : protocol_msg) =
167 match msg with
168 | Dead { node; incarnation = _; declarator = _ } ->
169 apply_member_transition t node (fun snap ~now ->
170 Protocol_pure.handle_dead snap msg ~now)
171 | _ -> ()
172
173let handle_user_msg t (msg : protocol_msg) =
174 match msg with
175 | User_msg { topic; payload; origin } -> (
176 let handlers =
177 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.get ~xt t.user_handlers) }
178 in
179 match Membership.find t.members origin with
180 | None -> ()
181 | Some member ->
182 let node = Membership.Member.node member in
183 List.iter (fun h -> h node topic payload) handlers)
184 | _ -> ()
185
186let handle_message t ~src (msg : protocol_msg) =
187 match msg with
188 | Ping _ -> handle_ping t ~src msg
189 | Ping_req _ -> handle_ping_req t ~src msg
190 | Ack _ -> handle_ack t msg
191 | Alive _ -> handle_alive_msg t msg
192 | Suspect _ -> handle_suspect_msg t msg
193 | Dead _ -> handle_dead_msg t msg
194 | User_msg _ -> handle_user_msg t msg
195
196let handle_packet t ~src (packet : packet) =
197 if String.equal packet.cluster t.config.cluster_name then begin
198 handle_message t ~src packet.primary;
199 List.iter (handle_message t ~src) packet.piggyback;
200 update_stats t (fun s -> { s with msgs_received = s.msgs_received + 1 })
201 end
202
203let process_udp_packet t ~buf ~src =
204 match Crypto.decrypt ~key:t.cipher_key buf with
205 | Error _ ->
206 update_stats t (fun s -> { s with msgs_dropped = s.msgs_dropped + 1 })
207 | Ok decrypted -> (
208 match Codec.decode_packet decrypted with
209 | Error _ ->
210 update_stats t (fun s -> { s with msgs_dropped = s.msgs_dropped + 1 })
211 | Ok packet -> handle_packet t ~src packet)
212
213let run_udp_receiver t =
214 while not (is_shutdown t) do
215 Buffer_pool.with_buffer t.recv_pool (fun buf ->
216 let n, src = Transport.recv_udp t.udp_sock buf in
217 let received = Cstruct.sub buf 0 n in
218 process_udp_packet t ~buf:received ~src)
219 done
220
221let probe_member t (member : Membership.Member.t) =
222 let target = Membership.Member.node member in
223 let seq = next_seq t in
224 let piggyback =
225 drain_piggyback t ~max_bytes:(t.config.udp_buffer_size - 100)
226 in
227 let ping = Ping { seq; sender = t.self } in
228 let packet = make_packet t ~primary:ping ~piggyback in
229
230 let waiter = Pending_acks.register t.pending_acks ~seq in
231 send_packet t ~dst:target.addr packet;
232
233 match
234 Pending_acks.wait waiter ~timeout:t.config.probe_timeout ~clock:t.clock
235 with
236 | Some _ ->
237 let now = now_mtime t in
238 Membership.update_member t.members target.id
239 { update = (fun m ~xt -> Membership.Member.record_ack ~xt m ~now) }
240 |> ignore;
241 true
242 | None ->
243 Pending_acks.cancel t.pending_acks ~seq;
244 false
245
246let indirect_probe t (member : Membership.Member.t) =
247 let target = Membership.Member.node member in
248 let seq = next_seq t in
249 let ping_req = Ping_req { seq; target = target.id; sender = t.self } in
250
251 let all_members = Membership.to_node_list t.members in
252 let indirect_targets =
253 Protocol_pure.select_indirect_targets ~self:t.self.id ~exclude:target.id
254 ~count:t.config.indirect_checks ~members:all_members
255 in
256
257 let waiter = Pending_acks.register t.pending_acks ~seq in
258 List.iter
259 (fun node ->
260 let packet = make_packet t ~primary:ping_req ~piggyback:[] in
261 send_packet t ~dst:node.addr packet)
262 indirect_targets;
263
264 match
265 Pending_acks.wait waiter ~timeout:t.config.probe_timeout ~clock:t.clock
266 with
267 | Some _ ->
268 let now = now_mtime t in
269 Membership.update_member t.members target.id
270 { update = (fun m ~xt -> Membership.Member.record_ack ~xt m ~now) }
271 |> ignore;
272 true
273 | None ->
274 Pending_acks.cancel t.pending_acks ~seq;
275 false
276
277let suspect_member t (member : Membership.Member.t) =
278 let node = Membership.Member.node member in
279 let inc = get_incarnation t in
280 let msg =
281 Suspect { node = node.id; incarnation = inc; suspector = t.self.id }
282 in
283 apply_member_transition t node.id (fun snap ~now ->
284 Protocol_pure.handle_suspect ~self:t.self.id snap msg ~now)
285
286let probe_cycle t =
287 let members = Membership.to_list t.members in
288 let member_nodes = List.map Membership.Member.node members in
289 let probe_idx =
290 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.get ~xt t.probe_index) }
291 in
292
293 match
294 Protocol_pure.next_probe_target ~self:t.self.id ~probe_index:probe_idx
295 ~members:member_nodes
296 with
297 | None -> ()
298 | Some (target_node, new_idx) -> (
299 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.set ~xt t.probe_index new_idx) };
300 match Membership.find t.members target_node.id with
301 | None -> ()
302 | Some member ->
303 let direct_ok = probe_member t member in
304 if not direct_ok then
305 let indirect_ok = indirect_probe t member in
306 if not indirect_ok then suspect_member t member)
307
308let run_protocol t =
309 while not (is_shutdown t) do
310 probe_cycle t;
311 Eio.Time.sleep t.clock t.config.protocol_interval
312 done
313
314let create ~config ~self ~udp_sock ~clock ~mono_clock ~secure_random =
315 match Crypto.init_key config.secret_key with
316 | Error _ -> Error `Invalid_key
317 | Ok cipher_key ->
318 Ok
319 {
320 config;
321 self;
322 members = Membership.create ();
323 incarnation = Kcas.Loc.make zero_incarnation;
324 sequence = Kcas.Loc.make 0;
325 broadcast_queue = Dissemination.create ();
326 pending_acks = Pending_acks.create ();
327 probe_index = Kcas.Loc.make 0;
328 send_pool =
329 Buffer_pool.create ~size:config.udp_buffer_size
330 ~count:config.send_buffer_count;
331 recv_pool =
332 Buffer_pool.create ~size:config.udp_buffer_size
333 ~count:config.recv_buffer_count;
334 udp_sock;
335 event_stream = Eio.Stream.create 100;
336 user_handlers = Kcas.Loc.make [];
337 cipher_key;
338 stats = Kcas.Loc.make empty_stats;
339 shutdown = Kcas.Loc.make false;
340 clock;
341 mono_clock;
342 secure_random;
343 }
344
345let shutdown t =
346 Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.set ~xt t.shutdown true) }
347
348let add_member t node_info =
349 let now = now_mtime t in
350 let member = Membership.Member.create ~now node_info in
351 Membership.add t.members member;
352 emit_event t (Join node_info)
353
354let remove_member t node_id =
355 match Membership.find t.members node_id with
356 | None -> false
357 | Some member ->
358 let node = Membership.Member.node member in
359 let removed = Membership.remove t.members node_id in
360 if removed then emit_event t (Leave node);
361 removed
362
363let local_node t = t.self
364let members t = Membership.to_list t.members
365let member_count t = Membership.count t.members
366let events t = t.event_stream
367
368let stats t =
369 let base = Kcas.Xt.commit { tx = (fun ~xt -> Kcas.Xt.get ~xt t.stats) } in
370 let alive, suspect, dead =
371 Membership.snapshot_all t.members
372 |> List.fold_left
373 (fun (a, s, d) snap ->
374 match snap.state with
375 | Alive -> (a + 1, s, d)
376 | Suspect -> (a, s + 1, d)
377 | Dead -> (a, s, d + 1))
378 (0, 0, 0)
379 in
380 {
381 base with
382 nodes_alive = alive;
383 nodes_suspect = suspect;
384 nodes_dead = dead;
385 queue_depth = Dissemination.depth t.broadcast_queue;
386 buffers_available =
387 Buffer_pool.available t.send_pool + Buffer_pool.available t.recv_pool;
388 buffers_total =
389 Buffer_pool.total t.send_pool + Buffer_pool.total t.recv_pool;
390 }
391
392let broadcast t ~topic ~payload =
393 let msg = User_msg { topic; payload; origin = t.self.id } in
394 enqueue_broadcast t msg
395
396let on_message t handler =
397 Kcas.Xt.commit
398 {
399 tx =
400 (fun ~xt ->
401 let handlers = Kcas.Xt.get ~xt t.user_handlers in
402 Kcas.Xt.set ~xt t.user_handlers (handler :: handlers));
403 }