(*****************************************************************************) (* *) (* Open Source License *) (* Copyright (c) 2023 Nomadic Labs, *) (* *) (* Permission is hereby granted, free of charge, to any person obtaining a *) (* copy of this software and associated documentation files (the "Software"),*) (* to deal in the Software without restriction, including without limitation *) (* the rights to use, copy, modify, merge, publish, distribute, sublicense, *) (* and/or sell copies of the Software, and to permit persons to whom the *) (* Software is furnished to do so, subject to the following conditions: *) (* *) (* The above copyright notice and this permission notice shall be included *) (* in all copies or substantial portions of the Software. *) (* *) (* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*) (* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, *) (* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL *) (* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*) (* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING *) (* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER *) (* DEALINGS IN THE SOFTWARE. *) (* *) (*****************************************************************************) (* FIXME https://gitlab.com/tezos/tezos/-/issues/4966 We should decide whether we want to implement a tracing mechanism as in the go implementation. *) module Gossipsub_intf = Gossipsub_intf open Gossipsub_intf module Make (C : AUTOMATON_CONFIG) : AUTOMATON with type Time.t = C.Time.t and module Span = C.Span and module Peer = C.Subconfig.Peer and module Topic = C.Subconfig.Topic and module Message_id = C.Subconfig.Message_id and module Message = C.Subconfig.Message = struct module Peer = C.Subconfig.Peer module Topic = C.Subconfig.Topic module Message_id = C.Subconfig.Message_id module Message = C.Subconfig.Message module Span = C.Span module Time = C.Time (** This module allows to compute a score for each peers. *) module Score : SCORE = struct (* FIXME https://gitlab.com/tezos/tezos/-/issues/4967 This is incomplete *) type t = {behaviour_penalty : int} let zero = {behaviour_penalty = 0} let float {behaviour_penalty} = -behaviour_penalty |> float_of_int let penalty {behaviour_penalty} penalty = {behaviour_penalty = behaviour_penalty + penalty} let compare s1 s2 = let f1 = float s1 in let f2 = float s2 in Float.compare f1 f2 include Compare.Make (struct type nonrec t = t let compare = compare end) end type message = Message.t type time = Time.t type span = Time.span type nonrec limits = (Peer.t, Message_id.t, span) limits type nonrec parameters = (Peer.t, Message_id.t) parameters type add_peer = {direct : bool; outbound : bool; peer : Peer.t} type remove_peer = {peer : Peer.t} type ihave = {peer : Peer.t; topic : Topic.t; message_ids : Message_id.t list} type iwant = {peer : Peer.t; message_ids : Message_id.t list} type graft = {peer : Peer.t; topic : Topic.t} type prune = { peer : Peer.t; topic : Topic.t; px : Peer.t Seq.t; backoff : span; } type publish = { sender : Peer.t option; topic : Topic.t; message_id : Message_id.t; message : message; } type join = {topic : Topic.t} type leave = {topic : Topic.t} type subscribe = {topic : Topic.t; peer : Peer.t} type unsubscribe = {topic : Topic.t; peer : Peer.t} (* FIXME not sure subtyping for output is useful. If it is, it is probably for few ouputs and could be removed. *) type _ output = | Negative_peer_score : Score.t -> [`IHave] output | Too_many_recv_ihave_messages : {count : int; max : int} -> [`IHave] output | Too_many_sent_iwant_messages : {count : int; max : int} -> [`IHave] output | Message_topic_not_tracked : [`IHave] output | Message_requested_message_ids : Message_id.t list -> [`IHave] output | On_iwant_messages_to_route : { routed_message_ids : [`Ignored | `Not_found | `Too_many_requests | `Message of message] Message_id.Map.t; } -> [`IWant] output | Peer_filtered : [`Graft] output | Unknown_topic : [`Graft] output | Peer_already_in_mesh : [`Graft] output | Grafting_direct_peer : [`Graft] output | Unexpected_grafting_peer : [`Graft] output | Grafting_peer_with_negative_score : [`Graft] output | Grafting_successfully : [`Graft] output | Peer_backed_off : [`Graft] output | Mesh_full : [`Graft] output | No_peer_in_mesh : [`Prune] output | Ignore_PX_score_too_low : Score.t -> [`Prune] output | No_PX : [`Prune] output | PX : Peer.Set.t -> [`Prune] output | Publish_message : {to_publish : Peer.Set.t} -> [`Publish] output | Already_published : [`Publish] output | Already_joined : [`Join] output | Joining_topic : {to_graft : Peer.Set.t} -> [`Join] output | Not_joined : [`Leave] output | Leaving_topic : {to_prune : Peer.Set.t} -> [`Leave] output | Heartbeat : { to_graft : Topic.Set.t Peer.Map.t; to_prune : Topic.Set.t Peer.Map.t; noPX_peers : Peer.Set.t; } -> [`Heartbeat] output | Peer_added : [`Add_peer] output | Peer_already_known : [`Add_peer] output | Removing_peer : [`Remove_peer] output | Subscribed : [`Subscribe] output | Subscribe_to_unknown_peer : [`Subscribe] output | Unsubscribed : [`Unsubscribe] output | Unsubscribe_from_unknown_peer : [`Unsubscribe] output type connection = { topics : Topic.Set.t; (** The set of topics the peer subscribed to. *) direct : bool; (** A direct (aka explicit) connection is a connection to which we forward all the messages. *) outbound : bool; (** An outbound connection is a connection we initiated. *) } type peer_status = | Connected | Disconnected of { expires : Time.t; (** The time after which the score can be cleared. *) } type peer_score = { score : Score.t; (** The score associated to a peer. *) peer_status : peer_status; } type fanout_peers = {peers : Peer.Set.t; last_published_time : time} module Message_cache = Message_cache.Make (C.Subconfig) (* FIXME https://gitlab.com/tezos/tezos/-/issues/4983 This data-structure should be documented. *) type state = { limits : limits; parameters : parameters; connections : connection Peer.Map.t; (** [connections] is the set of active connections. A connection is added through the `Add_peer` message and removed through the `Remove_peer` message. *) scores : peer_score Peer.Map.t; (** The scores are used to drive peer selection mechanisms. Scores are kept for at least [retain_duration] after a connection is removed, hence they can't be stored in {!connections}. Any peer having an active connection is associated to a score. *) ihave_per_heartbeat : int Peer.Map.t; iwant_per_heartbeat : int Peer.Map.t; mesh : Peer.Set.t Topic.Map.t; (** The mesh for a topic is a random subset of the connected peers subscribed to that topic and that have non-negative score, are not backed-off, and are not direct peers. The local peer routes full-messages to these peers. A topic is in the domain of the mesh iff the local peer has joined the topic (we also say that it tracks the topic). *) fanout : fanout_peers Topic.Map.t; (** The fanout for a topic is a random subset of the connected peers subscribed to that topic and that are not direct and have a score above some given threshold. The local peer routes full-messages to these peers. In contrast to the mesh, the fanout map contains topics the local peer has not joined. *) backoff : time Peer.Map.t Topic.Map.t; (** The backoff times associated to a topic and a peer. When a backoff is set, we refuse any `graft` from this peer. As a consequence, if a backoff is set for a peer and a topic, this peer is not expected to be in the mesh of the given topic. *) (* FIXME: https://gitlab.com/tezos/tezos/-/issues/5302 Introduce TTL to [seen_messages]. *) seen_messages : Message_id.Set.t; message_cache : Message_cache.t; rng : Random.State.t; heartbeat_ticks : int64; } (* Invariants: - Forall t, Topic.Map.mem t mesh <=> not (Map.mem t fanout) - Forall p c t, Peer.Map.find connections p = Some ({ connection = Expires _}) -> Topic.Map.find t mesh = None && Topic.Map.find t fanout = None - Forall p t, Peer.Map.find connections p = None -> Topic.Map.find t mesh = None && Topic.Map.find t fanout = None - Forall p. Peer.Map.mem connections p -> Peer.Map.mem scores p - Forall p, Peer.Map.mem connections p <-> Peer.Map.find scores p = Some {peer_status = Connected; _} *) (* FIXME https://gitlab.com/tezos/tezos/-/issues/4984 Test the those invariants *) module Monad = struct type 'a t = (state, 'a) State_monad.t type ('pass, 'fail) check = (state, 'pass, 'fail) State_monad.check include State_monad.M end let check_limits l = (* TODO: https://gitlab.com/tezos/tezos/-/issues/5129 Replace the asserts by something more informative. *) assert (l.degree_low > 0) ; assert (l.degree_out >= 0) ; assert (l.degree_score >= 0) ; assert (l.degree_low <= l.degree_optimal) ; assert (l.degree_high >= l.degree_optimal) ; assert (l.backoff_cleanup_ticks > 0) ; assert (l.score_cleanup_ticks > 0) ; (* TODO: https://gitlab.com/tezos/tezos/-/issues/5052 This requirement is not imposed in the spec/Go implementation. Relax this requirement or delete the todo. *) assert (l.degree_score + l.degree_out <= l.degree_optimal) ; assert (l.degree_out <= l.degree_low) ; assert (l.degree_out <= l.degree_optimal / 2) ; assert (l.history_gossip_length > 0) ; assert (l.history_gossip_length <= l.history_length) let make : Random.State.t -> limits -> parameters -> state = fun rng limits parameters -> check_limits limits ; { limits; parameters; connections = Peer.Map.empty; scores = Peer.Map.empty; ihave_per_heartbeat = Peer.Map.empty; iwant_per_heartbeat = Peer.Map.empty; mesh = Topic.Map.empty; fanout = Topic.Map.empty; backoff = Topic.Map.empty; seen_messages = Message_id.Set.empty; message_cache = Message_cache.create ~history_slots:limits.history_length ~gossip_slots:limits.history_gossip_length; rng; heartbeat_ticks = 0L; } module Helpers = struct let fail_if cond output_err = let open Monad.Syntax in if cond then fail output_err else unit let fail_if_not cond output_err = fail_if (not cond) output_err (* These projections enable let-punning. *) let max_recv_ihave_per_heartbeat state = state.limits.max_recv_ihave_per_heartbeat let max_sent_iwant_per_heartbeat state = state.limits.max_sent_iwant_per_heartbeat let degree_optimal state = state.limits.degree_optimal let publish_threshold state = state.limits.publish_threshold let do_px state = state.limits.do_px let peers_to_px state = state.limits.peers_to_px let accept_px_threshold state = state.limits.accept_px_threshold let prune_backoff state = state.limits.prune_backoff let unsubscribe_backoff state = state.limits.unsubscribe_backoff let graft_flood_backoff state = state.limits.graft_flood_backoff let retain_duration state = state.limits.retain_duration let fanout_ttl state = state.limits.fanout_ttl let heartbeat_interval state = state.limits.heartbeat_interval let backoff_cleanup_ticks state = state.limits.backoff_cleanup_ticks let score_cleanup_ticks state = state.limits.score_cleanup_ticks let degree_low state = state.limits.degree_low let degree_high state = state.limits.degree_high let degree_score state = state.limits.degree_score let degree_out state = state.limits.degree_out let max_gossip_retransmission state = state.limits.max_gossip_retransmission let mesh state = state.mesh let fanout state = state.fanout let backoff state = state.backoff let connections state = state.connections let scores state = state.scores let seen_messages state = state.seen_messages let peer_filter state = state.parameters.peer_filter let message_cache state = state.message_cache let rng state = state.rng let update ?(delta = 1) key map = Peer.Map.update key (function None -> Some delta | Some n -> Some (n + delta)) map let update_and_get ?(delta = 1) key map = let res = ref delta in Peer.Map.update key (function | None -> Some delta | Some n -> let value = n + delta in res := value ; Some value) map |> fun x -> (x, !res) let update_and_get_ihave_per_heartbeat ?delta key state = let ihave_per_heartbeat, res = update_and_get ?delta key state.ihave_per_heartbeat in let state = {state with ihave_per_heartbeat} in (state, res) let update_iwant_per_heartbeat ?delta key state = let iwant_per_heartbeat = update ?delta key state.iwant_per_heartbeat in let state = {state with iwant_per_heartbeat} in (state, ()) let find ?(default = 0) key map = match Peer.Map.find key map with None -> default | Some n -> n let find_iwant_per_heartbeat ?default key state = find ?default key state.iwant_per_heartbeat let reset_ihave_per_heartbeat state = ({state with ihave_per_heartbeat = Peer.Map.empty}, ()) let reset_iwant_per_heartbeat state = ({state with iwant_per_heartbeat = Peer.Map.empty}, ()) let heartbeat_ticks state = state.heartbeat_ticks let set_heartbeat_ticks heartbeat_ticks state = ({state with heartbeat_ticks}, ()) let set_connections connections state = ({state with connections}, ()) let set_scores scores state = ({state with scores}, ()) let topic_is_tracked topic state = let {mesh; _} = state in match Topic.Map.find topic mesh with None -> false | Some _ -> true let set_message_cache message_cache state = ({state with message_cache}, ()) let get_scores_score scores ~default peer = match Peer.Map.find peer scores with | None -> default | Some {score; peer_status = _} -> score let get_score ~default peer state = get_scores_score state.scores ~default peer let peer_has_outbound_connection peers ~default peer = Peer.Map.find_opt peer peers |> Option.map (fun connection -> connection.outbound) |> Option.value ~default (* TODO: https://gitlab.com/tezos/tezos/-/issues/5391 Optimize by having a topic to peers map *) let select_connections_peers connections scores rng topic ~filter ~max = Peer.Map.bindings connections |> List.filter_map (fun (peer, connection) -> let score = get_scores_score ~default:Score.zero scores peer in let topics = connection.topics in if filter peer connection score && Topic.Set.mem topic topics then Some peer else None) |> List.shuffle ~rng |> List.take_n max let select_peers topic ~filter ~max = let open Monad.Syntax in let*! connections in let*! scores in let*! rng in select_connections_peers connections scores rng topic ~filter ~max |> Peer.Set.of_list |> return let set_mesh_topic topic peers state = let state = {state with mesh = Topic.Map.add topic peers state.mesh} in (state, ()) let set_mesh mesh state = let state = {state with mesh} in (state, ()) let find_mesh topic state = Topic.Map.find topic state.mesh let find_fanout topic state = Topic.Map.find topic state.fanout let set_fanout_topic topic last_published_time peers state = if Peer.Set.is_empty peers then (state, ()) else let state = { state with fanout = Topic.Map.add topic {peers; last_published_time} state.fanout; } in (state, ()) let set_fanout fanout state = let state = {state with fanout} in (state, ()) let delete_mesh topic state = let state = {state with mesh = Topic.Map.remove topic state.mesh} in (state, ()) let delete_fanout topic state = let state = {state with fanout = Topic.Map.remove topic state.fanout} in (state, ()) let put_message_in_cache message_id message topic state = let state = { state with message_cache = Message_cache.add_message message_id message topic state.message_cache; } in (state, ()) let put_in_seen_messages message_id state = let state = { state with seen_messages = Message_id.Set.add message_id state.seen_messages; } in (state, ()) let get_backoff topic peer backoff = Option.bind (Topic.Map.find topic backoff) (fun per_peer_backoff -> Peer.Map.find peer per_peer_backoff) let exists_backoff topic peer backoff = get_backoff topic peer backoff |> Option.is_some let set_backoff backoff state = let state = {state with backoff} in (state, ()) let add_backoff_for_peer delay topic peer backoffs = let now = Time.now () in let backoff_expire = Time.add now delay in Topic.Map.update topic (function | None -> Some (Peer.Map.singleton peer backoff_expire) | Some peer_backoff -> Peer.Map.update peer (function | None -> Some backoff_expire | Some old_backoff -> if Time.(old_backoff < backoff_expire) then Some backoff_expire else Some old_backoff) peer_backoff |> Option.some) backoffs let set_backoff_for_peer delay topic peer = let open Monad.Syntax in let*! backoff in add_backoff_for_peer delay topic peer backoff |> set_backoff let add_scores_score peer score scores = Peer.Map.update peer (Option.map (fun peer_info -> {peer_info with score})) scores let add_score peer score = let open Monad.Syntax in let*! scores in let peers = add_scores_score peer score scores in set_scores peers let _check_peer_score peer = let open Monad.Syntax in let*! peer_score = get_score ~default:Score.zero peer in fail_if Score.(peer_score < zero) @@ Negative_peer_score peer_score end include Helpers module IHave = struct let check_too_many_recv_ihave_message count = let open Monad.Syntax in let*! max_recv_ihave_per_heartbeat in fail_if (count > max_recv_ihave_per_heartbeat) @@ Too_many_recv_ihave_messages {count; max = max_recv_ihave_per_heartbeat} (* FIXME https://gitlab.com/tezos/tezos/-/issues/5016 This check is not correct if the distant peer uses a different value for [max_recv_ihave_per_heartbeat] then our value for [max_sent_iwant_per_heartbeat]. *) let check_too_many_sent_iwant_message count = let open Monad.Syntax in let*! max_sent_iwant_per_heartbeat in fail_if (count >= max_sent_iwant_per_heartbeat) @@ Too_many_sent_iwant_messages {count; max = max_sent_iwant_per_heartbeat} let check_topic_tracked topic = let open Monad.Syntax in let*! is_topic_tracked = topic_is_tracked topic in fail_if_not is_topic_tracked Message_topic_not_tracked let check_not_empty iwant_message_ids = fail_if (List.is_empty iwant_message_ids) @@ Message_requested_message_ids [] let filter peer message_ids : Message_id.t list Monad.t = let open Monad.Syntax in let*! peer_filter in let*! seen_messages in let should_handle_message_id message_id : bool = (not (Message_id.Set.mem message_id seen_messages)) && peer_filter peer (`IHave message_id) in List.filter should_handle_message_id message_ids |> return let shuffle_and_trunc message_ids ~limit : (int * Message_id.t list) Monad.t = let open Monad.Syntax in let*! rng in let iwant_message_ids_len = List.length message_ids in (* Do not send more messages than [max_sent_iwant_per_heartbeat] *) let iwant_ids_to_send_n = min iwant_message_ids_len limit in let shuffle_iwant_ids = List.shuffle ~rng message_ids in let requested_message_ids = List.take_n iwant_ids_to_send_n shuffle_iwant_ids in return (iwant_ids_to_send_n, requested_message_ids) let handle peer topic message_ids : [`IHave] output Monad.t = let open Monad.Syntax in (* FIXME https://gitlab.com/tezos/tezos/-/issues/5009 Score check is missing. *) let* count_ihave_received = update_and_get_ihave_per_heartbeat peer in let*! count_iwant_sent = find_iwant_per_heartbeat peer in let*? () = check_too_many_recv_ihave_message count_ihave_received in let*? () = check_too_many_sent_iwant_message count_iwant_sent in let*? () = check_topic_tracked topic in let* iwant_message_ids = filter peer message_ids in let*? () = check_not_empty iwant_message_ids in let*! max_sent_iwant_per_heartbeat in let limit = max_sent_iwant_per_heartbeat - count_iwant_sent in (* Invariant: limit > 0 *) let* length, requested_message_ids = shuffle_and_trunc iwant_message_ids ~limit in let* () = update_iwant_per_heartbeat ~delta:length peer in (* FIXME https://gitlab.com/tezos/tezos/-/issues/4966 The go implementation traces some of the messages requested. *) Message_requested_message_ids requested_message_ids |> return end let handle_ihave : ihave -> [`IHave] output Monad.t = fun {peer; topic; message_ids} -> IHave.handle peer topic message_ids module IWant = struct let handle peer message_ids : [`IWant] output Monad.t = let open Monad.Syntax in (* FIXME https://gitlab.com/tezos/tezos/-/issues/5008 Score check is missing. *) let routed_message_ids = Message_id.Map.empty in let*! message_cache in let*! peer_filter in let*! max_gossip_retransmission in let message_cache, routed_message_ids = List.fold_left (fun (message_cache, messages) message_id -> let message_cache, info = match Message_cache.get_message_for_peer peer message_id message_cache with | None -> (message_cache, `Not_found) | Some (message_cache, message, access_counter) -> ( message_cache, if access_counter > max_gossip_retransmission then `Too_many_requests else if peer_filter peer (`IWant message_id) then `Message message else `Ignored ) in (message_cache, Message_id.Map.add message_id info messages)) (message_cache, routed_message_ids) message_ids in let* () = set_message_cache message_cache in On_iwant_messages_to_route {routed_message_ids} |> return end let handle_iwant : iwant -> [`IWant] output Monad.t = fun {peer; message_ids} -> IWant.handle peer message_ids module Graft = struct let check_filter peer = let open Monad.Syntax in let*! peer_filter in fail_if_not (peer_filter peer `Graft) Peer_filtered let check_topic_known topic = let open Monad.Syntax in let*! mesh_opt = find_mesh topic in match mesh_opt with | None -> Unknown_topic |> fail | Some mesh -> pass mesh let check_not_in_mesh mesh peer = fail_if (Peer.Set.mem peer mesh) Peer_already_in_mesh let check_active peer = let open Monad.Syntax in let*! connections in match Peer.Map.find peer connections with | None -> Unexpected_grafting_peer |> fail | Some connection -> pass connection let check_not_direct connection = let open Monad.Syntax in if connection.direct then Grafting_direct_peer |> fail else pass () let check_score peer topic score = let open Monad.Syntax in let*! prune_backoff in if Score.(score >= zero) then unit else let* () = set_backoff_for_peer prune_backoff topic peer in Grafting_peer_with_negative_score |> fail let check_mesh_size mesh connection = let open Monad.Syntax in let*! degree_high in (* Check the number of mesh peers; if it is at (or over) [degree_high], we only accept grafts from peers with outbound connections; this is a defensive check to restrict potential mesh takeover attacks combined with love bombing *) fail_if ((not connection.outbound) && Peer.Set.cardinal mesh >= degree_high) Mesh_full let check_backoff peer topic score = let open Monad.Syntax in let*! backoff in match get_backoff topic peer backoff with | None -> unit | Some backoff_expire -> let current = Time.now () in if Time.(current >= backoff_expire) then unit else let score = Score.penalty score 1 in let*! graft_flood_backoff in let score = if Time.(current < add backoff_expire graft_flood_backoff) then Score.penalty score 1 else score in let*! prune_backoff in let* () = set_backoff_for_peer prune_backoff topic peer in let* () = add_score peer score in fail Peer_backed_off let handle peer topic = (* FIXME: https://gitlab.com/tezos/tezos/-/issues/5264 Be sure that the peers that requested failed/rejected graft are properly pruned. *) (* FIXME: https://gitlab.com/tezos/tezos/-/issues/5305 A peer can be grafted only if he subscribed to this topic. *) let open Monad.Syntax in let*? () = check_filter peer in let*? mesh = check_topic_known topic in let*? () = check_not_in_mesh mesh peer in let*? connection = check_active peer in let*? () = check_not_direct connection in let*! score = get_score ~default:Score.zero peer in let*? () = check_backoff peer topic score in let*? () = check_score peer topic score in let*? () = check_mesh_size mesh connection in let* () = set_mesh_topic topic (Peer.Set.add peer mesh) in Grafting_successfully |> return end let handle_graft : graft -> [`Graft] output Monad.t = fun {peer; topic} -> Graft.handle peer topic module Prune = struct let check_px_score peer = let open Monad.Syntax in let*! accept_px_threshold in let*! score = get_score ~default:Score.zero peer in fail_if Compare.Float.(score |> Score.float < accept_px_threshold) @@ Ignore_PX_score_too_low score let check_topic_known topic = let open Monad.Syntax in let*! mesh_opt = find_mesh topic in match mesh_opt with | None -> No_peer_in_mesh |> fail | Some mesh -> pass mesh let handle peer topic ~px ~backoff = let open Monad.Syntax in let*? mesh = check_topic_known topic in let mesh = Peer.Set.remove peer mesh in let* () = set_mesh_topic topic mesh in let* () = set_backoff_for_peer backoff topic peer in let px = Peer.Set.of_seq px in if Peer.Set.is_empty px then No_PX |> return else let*? () = check_px_score peer in return (PX px) end let handle_prune : prune -> [`Prune] output Monad.t = fun {peer; topic; px; backoff} -> Prune.handle peer topic ~px ~backoff (* TODO: https://gitlab.com/tezos/tezos/-/issues/5148 Consider merging subcribe/unsubscribe with join/leave. *) module Subscribe = struct let handle topic peer = let open Monad.Syntax in let*! connections in match Peer.Map.find peer connections with | None -> return Subscribe_to_unknown_peer | Some connection -> let connection = {connection with topics = Topic.Set.add topic connection.topics} in let connections = Peer.Map.add peer connection connections in let* () = set_connections connections in (* TODO: https://gitlab.com/tezos/tezos/-/issues/5143 rust-libp2p adds the peer to the mesh if needed here. *) return Subscribed end let handle_subscribe : subscribe -> [`Subscribe] output Monad.t = fun {topic; peer} -> Subscribe.handle topic peer module Unsubscribe = struct let handle topic peer = let open Monad.Syntax in let*! connections in match Peer.Map.find peer connections with | None -> return @@ Unsubscribe_from_unknown_peer | Some connection -> let connection = {connection with topics = Topic.Set.remove topic connection.topics} in let peers = Peer.Map.add peer connection connections in let* () = set_connections peers in (* FIXME: https://gitlab.com/tezos/tezos/-/issues/5143 Remove unsubscribed peers from the mesh. *) return Unsubscribed end let handle_unsubscribe : unsubscribe -> [`Unsubscribe] output Monad.t = fun {topic; peer} -> Unsubscribe.handle topic peer module Publish = struct let check_seen message_id = let open Monad.Syntax in let*! seen_messages in fail_if (Message_id.Set.mem message_id seen_messages) Already_published let get_peers_for_unsubscribed_topic topic = let open Monad.Syntax in let*! publish_threshold in let*! degree_optimal in let now = Time.now () in let*! fanout_opt = find_fanout topic in match fanout_opt with | None -> let filter_by_score score = Compare.Float.(score |> Score.float >= publish_threshold) in let filter _peer {direct; _} score = (not direct) && filter_by_score score in let* not_direct_peers = select_peers topic ~filter ~max:degree_optimal in let* () = set_fanout_topic topic now not_direct_peers in return not_direct_peers | Some fanout -> let* () = set_fanout_topic topic now fanout.peers in return fanout.peers let handle ~sender topic message_id message : [`Publish] output Monad.t = let open Monad.Syntax in let*? () = check_seen message_id in let* () = put_message_in_cache message_id message topic in let*! mesh_opt = find_mesh topic in let* peers = match mesh_opt with | Some peers -> return peers | None -> get_peers_for_unsubscribed_topic topic in let peers = Option.fold ~none:peers ~some:(fun peer -> Peer.Set.remove peer peers) sender in let* () = put_in_seen_messages message_id in (* TODO: https://gitlab.com/tezos/tezos/-/issues/5272 Filter out peers from which we already received the message, or an IHave message? *) let* direct_peers = let filter _peer {direct; _} _score = direct in select_peers topic ~filter ~max:max_int in let to_publish = Peer.Set.union peers direct_peers in (* TODO: https://gitlab.com/tezos/tezos/-/issues/5010 Have a dedicated structure for direct peers? *) Publish_message {to_publish} |> return end let publish : publish -> [`Publish] output Monad.t = fun {sender; topic; message_id; message} -> Publish.handle ~sender topic message_id message module Join = struct let check_is_not_subscribed topic : (unit, [`Join] output) Monad.check = let open Monad.Syntax in let*! mesh in fail_if (Topic.Map.mem topic mesh) Already_joined let init_mesh topic : [`Join] output Monad.t = let open Monad.Syntax in let*! degree_optimal in let*! connections in let*! backoff in let*! scores in let is_valid peer = match Peer.Map.find peer connections with | None -> (* FIXME https://gitlab.com/tezos/tezos/-/issues/5005 Not supposed to happen. But maybe it is better to return a value for defensive programming. *) false | Some _peer_info -> let score = get_scores_score ~default:Score.zero scores peer in let backed_off = exists_backoff topic peer backoff in not (backed_off || Score.(score < zero)) in let*! fanout in let valid_fanout_peers = match Topic.Map.find topic fanout with | None -> Peer.Set.empty | Some fanout_peers -> Peer.Set.filter is_valid fanout_peers.peers in let* peers = (* We prioritize fanout peers to be in the mesh for this topic. If we need more peers, we look at all our peers subscribed to this topic. *) let valid_fanout_peers_len = Peer.Set.cardinal valid_fanout_peers in if valid_fanout_peers_len >= degree_optimal then return valid_fanout_peers else let max = max 0 (degree_optimal - valid_fanout_peers_len) in let* more_peers = let filter peer {direct; _} score = let backed_off = exists_backoff topic peer backoff in not (direct || backed_off || Score.(score < zero) || Peer.Set.mem peer valid_fanout_peers) in select_peers topic ~filter ~max in return (Peer.Set.union more_peers valid_fanout_peers) in let* () = set_mesh_topic topic peers in let* () = delete_fanout topic in Joining_topic {to_graft = peers} |> return let handle topic : [`Join] output Monad.t = let open Monad.Syntax in let*? () = check_is_not_subscribed topic in init_mesh topic end let join : join -> [`Join] output Monad.t = fun {topic} -> Join.handle topic module Leave = struct type mesh = Peer.Set.t let check_already_subscribed topic : (mesh, [`Leave] output) Monad.check = let open Monad.Syntax in let*! mesh in match Topic.Map.find topic mesh with | None -> Not_joined |> fail | Some mesh -> pass mesh let handle_mesh topic mesh : [`Leave] output Monad.t = let open Monad.Syntax in let*! unsubscribe_backoff in let*! backoff in let* () = Peer.Set.fold (fun peer backoff -> add_backoff_for_peer unsubscribe_backoff topic peer backoff) mesh backoff |> set_backoff in Leaving_topic {to_prune = mesh} |> return let handle topic : [`Leave] output Monad.t = let open Monad.Syntax in let*? mesh = check_already_subscribed topic in let* () = delete_mesh topic in handle_mesh topic mesh end let leave : leave -> [`Leave] output Monad.t = fun {topic} -> Leave.handle topic module Heartbeat = struct let clear_backoff = let open Monad.Syntax in let*! heartbeat_ticks in let*! heartbeat_interval in let*! backoff_cleanup_ticks in (* NOTE: Probably the cleanup can also be done lazily: at use, if a backoff time is expired, then remove it *) (* We only clear once every [backoff_cleanup_ticks] ticks to avoid iterating over the map(s) too much *) if Int64.(rem heartbeat_ticks (of_int backoff_cleanup_ticks)) <> 0L then return () else let current = Time.now () in let current_with_slack = (* Subtract some slack time to the current time to account for the message latency; for details, see https://github.com/libp2p/go-libp2p-pubsub/issues/368 *) Time.sub current (Time.mul_span heartbeat_interval 2) in let*! backoff in Topic.Map.filter_map (fun _topic peer_backoff -> let peer_backoff = Peer.Map.filter (fun _peer expire -> Time.(expire > current_with_slack)) peer_backoff in if Peer.Map.is_empty peer_backoff then None else Some peer_backoff) backoff |> set_backoff let clear_expired_scores = let open Monad.Syntax in let*! heartbeat_ticks in let*! score_cleanup_ticks in let*! scores in (* We only clear once every [score_cleanup_ticks] ticks to avoid iterating over the map(s) too much *) if Int64.(rem heartbeat_ticks (of_int score_cleanup_ticks)) <> 0L then return () else let current = Time.now () in Peer.Map.filter (fun _peer {score = _; peer_status} -> match peer_status with | Disconnected {expires = at} -> Time.(at > current) | Connected -> true) scores |> set_scores let cleanup = let open Monad.Syntax in (* Clean up expired backoffs *) let* () = clear_backoff in (* Clean up expired scores *) let* () = clear_expired_scores in (* Clean up IHave and IWant counters *) let* () = reset_ihave_per_heartbeat in let* () = reset_iwant_per_heartbeat in (* TODO: https://gitlab.com/tezos/tezos/-/issues/4967 Apply IWANT request penalties *) return () (* Update mesh for grafted and pruned peers. Note that in the Go implementation this update is done on-the-fly. The semantics is however the same, given that a peer cannot be grafted and then pruned or vice-versa. The reasoning for why that is the case is as follows. There are three blocks of updates in the code above: 1) graft peers if [d_mesh < degree_low] 2) prune peers if [d_mesh > degree_high] 3) graft peers if [d_mesh > degree_low] and not enough outbound peers ([d_mesh] denotes [Peer.Set.cardinal peers]) Condition 1) is mutually exclusive with the other two. Now, a peer p cannot be pruned in 2) and then grafted in 3) because in 2): A) Either we have enough outbound peers, then we don't execute 3). B) Or we don't have enough outbound peers, so p is not outbound and in 3) we only graft outbound peers. *) (* TODO: https://gitlab.com/tezos/tezos/-/issues/4967 Revisit this comment if opportunistic grafting is implemented. *) let update_mesh mesh ~to_graft ~to_prune = let update f = Peer.Map.fold (fun peer topicset mesh -> Topic.Set.fold (fun topic mesh -> Topic.Map.update topic (f peer) mesh) topicset mesh) in let add_peer peer = function | None -> Peer.Set.singleton peer |> Option.some | Some peers -> (* Note: [peer] should not be in [peers] already *) Peer.Set.add peer peers |> Option.some in let remove_peer peer = function | None -> (* Note: this should not occur *) None | Some peers -> Peer.Set.remove peer peers |> Option.some in mesh |> update add_peer to_graft |> update remove_peer to_prune |> set_mesh (* Mesh maintenance. For each topic, do in order: - Prune all peers with negative score, do not enable peer exchange for these peers. - If the number of remaining peers in the topic mesh is less than [degree_low], then select as many random peers (not already in the mesh topic) to graft as possible so that to have [degree_optimal] in the topic mesh. The selected peers should have a non-negative score, should not be backed off, and should not be direct peers (and should be subscribed to the topic). - If the number of remaining peers in the topic mesh is higher than [degree_high], then select as many peers (not already in the mesh topic) to prune as possible so that to have [degree_optimal] in the topic mesh. See [select_peers_to_prune] to see how the selection is performed. - If the number of remaining peers in the topic mesh is higher than [degree_low] and the number of outbound peers therein is smaller than [degree_out], then select an additional number of peers to graft (with the same conditions for grafting) to ensure that there are at least [degree_out] outbound peers in the mesh. Finally, for pruned peers, back them off for [prune_backoff] time. *) let maintain_mesh = let open Monad.Syntax in let*! connections in let*! scores in let*! backoff in let*! rng in let*! prune_backoff in let*! degree_optimal in let*! degree_low in let*! degree_high in let*! degree_score in let*! degree_out in let has_outbound_connection = peer_has_outbound_connection connections ~default:false in let get_score = get_scores_score scores ~default:Score.zero in (* [add_to_peers_topic_set peers_topicset topic peers] adds [topic] to each [peer] from the [peers] list into the [peers_topicset], which is a peer to topicset map. It returns the new map. Note that the following invariant is maintained by the caller: the topicset for [peer] does not already contain [topic]. *) let add_to_peers_topic_set map topic peers = List.fold_left (fun acc peer -> Peer.Map.update peer (function | None -> Some (Topic.Set.singleton topic) | Some topics -> Some (Topic.Set.add topic topics)) acc) map peers in (* Update [to_prune] and [old_peers] as follows: - In [to_prune], add [topic] for each of the peers in [new_peers]. - Remove [new_peers] from [old_peers]. *) let prune topic to_prune ~old_peers new_peers = let to_prune = add_to_peers_topic_set to_prune topic new_peers in let peers = List.fold_left (fun acc peer -> Peer.Set.remove peer acc) old_peers new_peers in (to_prune, peers) in (* Keep the first [degree_score] peers by score and the remaining up to [degree_optimal] randomly, under the constraint that we keep [degree_out] peers in the mesh (if we have that many). *) let select_peers_to_prune peers = (* TODO: https://gitlab.com/tezos/tezos/-/issues/5052 Consider first selecting [degree_out] peers and then ordering on the score. *) (* Sort descendingly by score, but shuffle first for the case we don't use the score. Head of list has highest score. *) let peers = peers |> Peer.Set.elements |> List.shuffle ~rng |> List.sort (fun peer1 peer2 -> let s1 = get_score peer1 in let s2 = get_score peer2 in Score.compare s2 s1) in let peers_high_score, peers_low_score = List.split_n degree_score peers in let peers_low_score = List.shuffle ~rng peers_low_score in let peers_to_keep, peers_to_prune = let to_keep, peers_to_prune = (* Recall that [degree_score <= degree_optimal] *) List.split_n (degree_optimal - degree_score) peers_low_score in (peers_high_score @ to_keep, peers_to_prune) in (* Count the outbound peers we are keeping. *) let outbound_peers_to_keep, inbound_peers_to_keep = List.partition has_outbound_connection peers_to_keep in let num_outbound_to_keep = List.length outbound_peers_to_keep in (* If [num_outbound] is less than [degree_out], swap some outbound peers from the peers to prune with inbound peers from the peers to keep. *) if num_outbound_to_keep < degree_out then let outbound_peers_to_prune, inbound_peers_to_prune = List.partition has_outbound_connection peers_to_prune in let num_outbound_to_prune = List.length outbound_peers_to_prune in let num_inbound_to_keep = List.length inbound_peers_to_keep in let num_to_swap = min (max 0 (degree_out - num_outbound_to_keep)) (min num_outbound_to_prune num_inbound_to_keep) in if num_to_swap > 0 then (* We additionally prune the [num_to_swap] inbound peers among the ones with a low score (or were shuffled); that's why we revert [inbound_peers_to_keep]. *) let inbound_peers_to_prune = List.take_n num_to_swap (List.rev inbound_peers_to_keep) @ inbound_peers_to_prune in (* Actually keep [num_to_swap] outbound peers. *) let outbound_peers_to_prune = List.drop_n num_to_swap outbound_peers_to_prune in inbound_peers_to_prune @ outbound_peers_to_prune else peers_to_prune else peers_to_prune in (* [maintain_topic_mesh topic peers (to_prune, to_graft, noPX_peers)] maintains the mesh for [topic] where [peers] are the original [peers] in this topic. [to_prune], [to_graft], [noPX_peers] are the peers to be pruned, grafted, and those for which no peer exchange should be done performed, accumulated from the maintenance for other mesh topics. *) let maintain_topic_mesh topic peers (to_prune, to_graft, noPX_peers) = let to_prune, peers, noPX_peers = (* Drop all peers with negative score, without PX *) Peer.Set.fold (fun peer (to_prune, peers, noPX_peers) -> if Score.(get_score peer < zero) then let to_prune, peers = prune topic to_prune ~old_peers:peers [peer] in let noPX_peers = Peer.Set.add peer noPX_peers in (to_prune, peers, noPX_peers) else (to_prune, peers, noPX_peers)) peers (to_prune, peers, noPX_peers) in (* Do we have too few peers? *) let num_peers = Peer.Set.cardinal peers in let to_graft = if num_peers < degree_low then let max = degree_optimal - num_peers in (* Filter out our current and direct peers, peers we are backing off, and peers with negative score. *) let filter peer connection score = let in_mesh = Peer.Set.mem peer peers in let backed_off = exists_backoff topic peer backoff in (not in_mesh) && (not backed_off) && (not connection.direct) && Score.(score >= zero) in select_connections_peers connections scores rng topic ~filter ~max |> add_to_peers_topic_set to_graft topic else to_graft in (* Do we have too many peers? *) let to_prune, peers = if num_peers > degree_high then (* We'll prune [num_peers - degree_optimal] peers. *) select_peers_to_prune peers |> prune topic to_prune ~old_peers:peers else (to_prune, peers) in (* Do we have enough outbound peers? *) let num_peers = Peer.Set.cardinal peers in let to_graft = if num_peers >= degree_low then let num_outbound = Peer.Set.fold (fun peer count -> if has_outbound_connection peer then count + 1 else count) peers 0 in if num_outbound < degree_out then let max = degree_out - num_outbound in (* Filter out our current and direct peers, peers we are backing off, and peers with negative score *) let filter peer connection score = let in_mesh = Peer.Set.mem peer peers in let backed_off = exists_backoff topic peer backoff in (not in_mesh) && (not backed_off) && (not connection.direct) && Score.(score >= zero) && has_outbound_connection peer in select_connections_peers connections scores rng topic ~filter ~max |> add_to_peers_topic_set to_graft topic else to_graft else to_graft in (* TODO: https://gitlab.com/tezos/tezos/-/issues/4967 Try to improve the mesh with opportunistic grafting *) (to_graft, to_prune, noPX_peers) in let*! mesh in let to_graft, to_prune, noPX_peers = Topic.Map.fold maintain_topic_mesh mesh (Peer.Map.empty, Peer.Map.empty, Peer.Set.empty) in (* Update backoff for pruned peers. *) let* () = Peer.Map.fold (fun peer topicset backoff -> Topic.Set.fold (fun topic backoff -> add_backoff_for_peer prune_backoff topic peer backoff) topicset backoff) to_prune backoff |> set_backoff in (* Update mesh for grafted and pruned peers *) let* () = update_mesh mesh ~to_graft ~to_prune in return (to_graft, to_prune, noPX_peers) let update_fanout fanout ~to_add ~to_remove = let update f topic_peers_list fanout = List.fold_left (fun fanout (topic, peers) -> Topic.Map.update topic (f peers) fanout) fanout topic_peers_list in let add_peers peers_to_add = function | None -> (* impossible: in [maintain_fanout] we only consider topics that are in the domain of the fanout map *) assert false | Some v -> let peers = List.fold_left (fun peers peer -> Peer.Set.add peer peers) v.peers peers_to_add in Some {v with peers} in let remove_peers peers_to_remove = function | None -> (* impossible, as in the previous case *) assert false | Some v -> let peers = List.fold_left (fun peers peer -> Peer.Set.remove peer peers) v.peers peers_to_remove in Some {v with peers} in fanout |> update add_peers to_add |> update remove_peers to_remove |> set_fanout (* Maintain the fanout map as follows: - Remove topics to which the local peer has not published in the [fanout_ttl] time. - Remove peers that are not subscribed anymore or that have a score below [publish_threshold]. - Remove peers that are expiring. - If for a topic the set of fanout peers is below [degree_optimal], then try to fill the map so that to have [degree_optimal] in the topic fanout. The selected peers should have a score above [publish_threshold] and should not be direct peers (and should be subscribed to the topic). *) let maintain_fanout = let open Monad.Syntax in let*! connections in let*! scores in let*! rng in let*! degree_optimal in let*! publish_threshold in let*! fanout_ttl in let expire_fanout = let current = Time.now () in (* TODO: https://gitlab.com/tezos/tezos/-/issues/5184 Optimize by having a min and a max last published time to avoid traversing the map when not needed? *) Topic.Map.filter (fun _topic {last_published_time; peers = _} -> Time.(add last_published_time fanout_ttl >= current)) in let maintain_topic_fanout topic {peers; _} (to_add, to_remove) = (* Check whether our peers are still in the topic and have a score above the publish threshold *) let peers_to_keep, peers_to_remove = Peer.Set.fold (fun peer acc -> match Peer.Map.find peer connections with | None -> (* impossible, given the global invariants on the state *) assert false | Some connection -> let score = get_scores_score ~default:Score.zero scores peer in if Topic.Set.mem topic connection.topics && Compare.Float.(Score.float score >= publish_threshold) then acc else let peers_to_keep, peers_to_remove = acc in (Peer.Set.remove peer peers_to_keep, peer :: peers_to_remove)) peers (peers, []) in let to_remove = (topic, peers_to_remove) :: to_remove in (* Do we need more peers? *) let num_peers = Peer.Set.cardinal peers_to_keep in if num_peers < degree_optimal then let ineed = degree_optimal - num_peers in (* Filter our current and direct peers and peers with score above the publish threshold *) let filter peer connection score = let in_fanout = Peer.Set.mem peer peers_to_keep in (not in_fanout) && (not connection.direct) && Compare.Float.(score |> Score.float >= publish_threshold) in let new_peers = select_connections_peers connections scores rng topic ~filter ~max:ineed in let to_add = (topic, new_peers) :: to_add in (to_add, to_remove) else (to_add, to_remove) in let*! fanout in (* Expire fanout for topics we haven't published to in a while. *) let fanout = expire_fanout fanout in let to_add, to_remove = Topic.Map.fold maintain_topic_fanout fanout ([], []) in (* Update the fanout map. *) update_fanout fanout ~to_add ~to_remove let handle = let open Monad.Syntax in let*! heartbeat_ticks in let* () = set_heartbeat_ticks (Int64.succ heartbeat_ticks) in (* cleaning up *) let* () = cleanup in (* TODO: https://gitlab.com/tezos/tezos/-/issues/4967 Cache scores throughout the heartbeat *) (* Maintain the mesh for topics we have joined. Concretely, in case the number of peers per topic (in the mesh) is lower than [degree_low] or higher than [degree_high], then select peers to graft or prune respectively, so that the new number of peers per topic becomes [degree_optimal]. *) let* to_graft, to_prune, noPX_peers = maintain_mesh in (* Maintain our fanout for topics we are publishing to, but we have not joined. *) let* () = maintain_fanout in (* Advance the message history sliding window. *) let*! message_cache in let* () = Message_cache.shift message_cache |> set_message_cache in Heartbeat {to_graft; to_prune; noPX_peers} |> return end let heartbeat : [`Heartbeat] output Monad.t = Heartbeat.handle module Add_peer = struct let handle ~direct ~outbound peer : [`Add_peer] output Monad.t = let open Monad.Syntax in let*! connections in let*! scores in match Peer.Map.find peer connections with | None -> let connection = {direct; topics = Topic.Set.empty; outbound} in let connections = Peer.Map.add peer connection connections in let scores = Peer.Map.update peer (function | None -> Some {score = Score.zero; peer_status = Connected} | Some score -> Some {score with peer_status = Connected}) scores in let* () = set_connections connections in let* () = set_scores scores in return Peer_added | Some _ -> return Peer_already_known end let add_peer : add_peer -> [`Add_peer] output Monad.t = fun {direct; outbound; peer} -> Add_peer.handle ~direct ~outbound peer module Remove_peer = struct let handle peer : [`Remove_peer] output Monad.t = let open Monad.Syntax in let*! mesh in let mesh = Topic.Map.map (fun peers -> Peer.Set.remove peer peers) mesh in let* () = set_mesh mesh in let*! fanout in let fanout = Topic.Map.map (fun fanout_peers -> {fanout_peers with peers = Peer.Set.remove peer fanout_peers.peers}) fanout in let* () = set_fanout fanout in let*! retain_duration in let*! scores in let*! connections in let* () = Peer.Map.remove peer connections |> set_connections in let* () = Peer.Map.update peer (function | None -> (* This should not happen (global invariant). *) None | Some score -> if Score.(score.score > zero) then (* We only retain non-positive scores to dissuade attacks on the score function. *) None else let now = Time.now () in let at = Time.add now retain_duration in let score = {score with peer_status = Disconnected {expires = at}} in Some score) scores |> set_scores in Removing_peer |> return end let remove_peer : remove_peer -> [`Remove_peer] output Monad.t = fun {peer} -> Remove_peer.handle peer let select_px_peers state ~peer_to_prune topic ~noPX_peers = let do_px = do_px state in let peers_to_px = peers_to_px state in let filter peer _conn score = (not (Peer.equal peer_to_prune peer)) && Score.(score >= zero) in if do_px && not (Peer.Set.mem peer_to_prune noPX_peers) then select_connections_peers state.connections state.scores state.rng topic ~filter ~max:peers_to_px else [] let select_gossip_messages state = let rng = state.rng in let select_gossip_for_peer message_ids = (* We shuffle the message ids so that we emit a different set for each peer. *) (* TODO: https://gitlab.com/tezos/tezos/-/issues/5396 Can this be optimized? *) message_ids |> List.shuffle ~rng |> List.take_n state.limits.max_sent_iwant_per_heartbeat in let select_gossip_for_topic topic excluded_peers = let message_ids = Message_cache.get_message_ids_to_gossip topic state.message_cache in if message_ids = [] then [] else (* We collect the peers with a score above [gossip_threshold] that are not in the excluded set and are not direct peers. *) let filter peer {direct; _} score = (not direct) && Compare.Float.( score |> Score.float >= state.limits.gossip_threshold) && not (Peer.Set.mem peer excluded_peers) in (* We first select all peers satisfying the criterion and then we see if we have too many. *) let peers = select_connections_peers state.connections state.scores rng topic ~filter ~max:Int.max_int in let num_peers = List.length peers in let target_num = max state.limits.degree_lazy (int_of_float (state.limits.gossip_factor *. float_of_int num_peers)) in let selected_peers = List.take_n target_num peers in (* Prepare the IHave gossip to the selected peers. *) List.fold_left (fun messages peer -> let message_ids = select_gossip_for_peer message_ids in {peer; topic; message_ids} :: messages) [] selected_peers in (* Prepare the IHave gossip for each topic in the mesh or fanout maps. Peers in the mesh/fanout are excluded from gossip (because we send full messages to them). *) let add_gossip_for_topic topic peers gossip_msgs = let new_msgs = select_gossip_for_topic topic peers in List.rev_append new_msgs gossip_msgs in Topic.Map.fold add_gossip_for_topic state.mesh [] |> Topic.Map.fold (fun topic {peers; _} -> add_gossip_for_topic topic peers) state.fanout (* Helpers. *) let pp_add_peer fmtr ({direct; outbound; peer} : add_peer) = let open Format in fprintf fmtr "{ direct=%b; outbound=%b; peer=%a }" direct outbound Peer.pp peer let pp_remove_peer fmtr ({peer} : remove_peer) = let open Format in fprintf fmtr "{ peer=%a }" Peer.pp peer let pp_ihave fmtr ({peer; topic; message_ids} : ihave) = let open Format in fprintf fmtr "{ peer=%a; topic=%a; message_ids=[%a] }" Peer.pp peer Topic.pp topic (pp_print_list ~pp_sep:(fun fmtr () -> fprintf fmtr ";") Message_id.pp) message_ids let pp_iwant fmtr ({peer : Peer.t; message_ids : Message_id.t list} : iwant) = let open Format in fprintf fmtr "{ peer=%a; message_ids=[%a] }" Peer.pp peer (pp_print_list ~pp_sep:(fun fmtr () -> fprintf fmtr ";") Message_id.pp) message_ids let pp_graft fmtr ({peer; topic} : graft) = let open Format in fprintf fmtr "{ peer=%a; topic=%a }" Peer.pp peer Topic.pp topic let pp_prune fmtr ({peer; topic; px; backoff} : prune) = let open Format in fprintf fmtr "{ peer=%a; topic=%a; px=%a; backoff=%a }" Peer.pp peer Topic.pp topic (pp_print_list ~pp_sep:(fun fmtr () -> fprintf fmtr ";") Peer.pp) (List.of_seq px) Span.pp backoff let pp_publish fmtr ({sender; topic; message_id; message} : publish) = let open Format in fprintf fmtr "{ sender=%a; topic=%a; message_id=%a; message=%a }" (pp_print_option ~none:(fun fmtr () -> pp_print_string fmtr "None") Peer.pp) sender Topic.pp topic Message_id.pp message_id Message.pp message let pp_join fmtr ({topic} : join) = let open Format in fprintf fmtr "{ topic=%a }" Topic.pp topic let pp_leave fmtr ({topic} : leave) = let open Format in fprintf fmtr "{ topic=%a }" Topic.pp topic let pp_subscribe fmtr ({topic; peer} : subscribe) = let open Format in fprintf fmtr "{ topic=%a; peer=%a }" Topic.pp topic Peer.pp peer let pp_unsubscribe fmtr ({topic; peer} : unsubscribe) = let open Format in fprintf fmtr "{ topic=%a; peer=%a }" Topic.pp topic Peer.pp peer module Introspection = struct (* This module reexport datatypes so that it can be used for introspection. While at the moment, this module reexport purely the datatype, we can decide that for abstraction purpose, we do not export all those fields in the future. This may have a small overhead cost, but for introspection it should be irrelevant. *) type nonrec connection = connection = { topics : Topic.Set.t; direct : bool; outbound : bool; } type nonrec peer_status = peer_status = | Connected | Disconnected of {expires : Time.t} type nonrec peer_score = peer_score = { score : Score.t; peer_status : peer_status; } type nonrec fanout_peers = fanout_peers = { peers : Peer.Set.t; last_published_time : time; } module Message_cache = Message_cache type view = state = { limits : limits; parameters : parameters; connections : connection Peer.Map.t; scores : peer_score Peer.Map.t; ihave_per_heartbeat : int Peer.Map.t; iwant_per_heartbeat : int Peer.Map.t; mesh : Peer.Set.t Topic.Map.t; fanout : fanout_peers Topic.Map.t; backoff : time Peer.Map.t Topic.Map.t; seen_messages : Message_id.Set.t; message_cache : Message_cache.t; rng : Random.State.t; heartbeat_ticks : int64; } let view state = state type connected_peers_filter = | Direct | Subscribed_to of Topic.t | Score_above of {threshold : float} (* Add other cases here if needed. *) let connected_peers_filter _peer connection score = function | Direct -> connection.direct | Subscribed_to topic -> Topic.Set.mem topic connection.topics | Score_above {threshold} -> Compare.Float.(Score.float score >= threshold) let get_connected_peers = let rec filter_rec peer connection score = function | [] -> true | filter :: filters -> connected_peers_filter peer connection score filter && filter_rec peer connection score filters in fun ?(filters = []) view -> Peer.Map.fold (fun peer connection acc -> let score = get_score ~default:Score.zero peer view in if filter_rec peer connection score filters then peer :: acc else acc) view.connections [] let get_peers_in_topic_mesh topic state = match Topic.Map.find topic state.mesh with | None -> [] | Some peers -> Peer.Set.elements peers let get_subscribed_topics peer state = match Peer.Map.find peer state.connections with | None -> [] | Some connection -> Topic.Set.elements connection.topics let get_our_topics state = Topic.Map.fold (fun topic _peers acc -> topic :: acc) state.mesh [] let get_fanout_peers topic state = match Topic.Map.find topic state.fanout with | None -> [] | Some fanout_peers -> Peer.Set.elements fanout_peers.peers let limits state = state.limits let has_joined topic {mesh; _} = Topic.Map.mem topic mesh end end