ceremonyclient/go-libp2p-blossomsub/randomsub.go

169 lines
3.2 KiB
Go
Raw Normal View History

package blossomsub
import (
"context"
"math"
"github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/libp2p/go-libp2p/core/protocol"
)
const (
RandomSubID = protocol.ID("/randomsub/1.0.0")
)
var (
RandomSubD = 6
)
// NewRandomSub returns a new PubSub object using RandomSubRouter as the router.
func NewRandomSub(ctx context.Context, h host.Host, size int, opts ...Option) (*PubSub, error) {
rt := &RandomSubRouter{
size: size,
peers: make(map[peer.ID]protocol.ID),
}
return NewPubSub(ctx, h, rt, opts...)
}
// RandomSubRouter is a router that implements a random propagation strategy.
// For each message, it selects the square root of the network size peers, with a min of RandomSubD,
// and forwards the message to them.
type RandomSubRouter struct {
p *PubSub
peers map[peer.ID]protocol.ID
size int
tracer *pubsubTracer
}
func (rs *RandomSubRouter) Protocols() []protocol.ID {
return []protocol.ID{RandomSubID, FloodSubID}
}
func (rs *RandomSubRouter) Attach(p *PubSub) {
rs.p = p
rs.tracer = p.tracer
}
func (rs *RandomSubRouter) AddPeer(p peer.ID, proto protocol.ID) {
rs.tracer.AddPeer(p, proto)
rs.peers[p] = proto
}
func (rs *RandomSubRouter) RemovePeer(p peer.ID) {
rs.tracer.RemovePeer(p)
delete(rs.peers, p)
}
func (rs *RandomSubRouter) EnoughPeers(bitmask []byte, suggested int) bool {
// check all peers in the bitmask
tmap, ok := rs.p.bitmasks[string(bitmask)]
if !ok {
return false
}
fsPeers := 0
rsPeers := 0
// count floodsub and randomsub peers
for p := range tmap {
switch rs.peers[p] {
case FloodSubID:
fsPeers++
case RandomSubID:
rsPeers++
}
}
if suggested == 0 {
suggested = RandomSubD
}
if fsPeers+rsPeers >= suggested {
return true
}
if rsPeers >= RandomSubD {
return true
}
return false
}
func (rs *RandomSubRouter) AcceptFrom(peer.ID) AcceptStatus {
return AcceptAll
}
func (rs *RandomSubRouter) HandleRPC(rpc *RPC) {}
func (rs *RandomSubRouter) Publish(msg *Message) {
from := msg.ReceivedFrom
tosend := make(map[peer.ID]struct{})
rspeers := make(map[peer.ID]struct{})
src := peer.ID(msg.GetFrom())
bitmask := msg.GetBitmask()
tmap, ok := rs.p.bitmasks[string(bitmask)]
if !ok {
return
}
for p := range tmap {
if p == from || p == src {
continue
}
if rs.peers[p] == FloodSubID {
tosend[p] = struct{}{}
} else {
rspeers[p] = struct{}{}
}
}
if len(rspeers) > RandomSubD {
target := RandomSubD
sqrt := int(math.Ceil(math.Sqrt(float64(rs.size))))
if sqrt > target {
target = sqrt
}
if target > len(rspeers) {
target = len(rspeers)
}
xpeers := peerMapToList(rspeers)
shufflePeers(xpeers)
xpeers = xpeers[:target]
for _, p := range xpeers {
tosend[p] = struct{}{}
}
} else {
for p := range rspeers {
tosend[p] = struct{}{}
}
}
out := rpcWithMessages(msg.Message)
for p := range tosend {
mch, ok := rs.p.peers[p]
if !ok {
continue
}
select {
case mch <- out:
rs.tracer.SendRPC(out, p)
default:
log.Infof("dropping message to peer %s: queue full", p)
rs.tracer.DropRPC(out, p)
}
}
}
func (rs *RandomSubRouter) Join(bitmask []byte) {
rs.tracer.Join(bitmask)
}
func (rs *RandomSubRouter) Leave(bitmask []byte) {
rs.tracer.Join(bitmask)
}