https://gitlab.com/tezos/tezos
Tip revision: ff91884efaa43c882de77cc08daf40ebb5f84668 authored by Gauthier SEBILLE on 13 April 2023, 13:38:03 UTC
DAC: fix tests
DAC: fix tests
Tip revision: ff91884
tezos_gossipsub.ml
(*****************************************************************************)
(* *)
(* Open Source License *)
(* Copyright (c) 2023 Nomadic Labs, <contact@nomadic-labs.com> *)
(* *)
(* 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