ceremonyclient/go-libp2p-blossomsub/mcache.go
2023-07-07 01:07:10 -05:00

106 lines
2.4 KiB
Go

package blossomsub
import (
"bytes"
"fmt"
"github.com/libp2p/go-libp2p/core/peer"
)
// NewMessageCache creates a sliding window cache that remembers messages for as
// long as `history` slots.
//
// When queried for messages to advertise, the cache only returns messages in
// the last `gossip` slots.
//
// The `gossip` parameter must be smaller or equal to `history`, or this
// function will panic.
//
// The slack between `gossip` and `history` accounts for the reaction time
// between when a message is advertised via IHAVE gossip, and the peer pulls it
// via an IWANT command.
func NewMessageCache(gossip, history int) *MessageCache {
if gossip > history {
err := fmt.Errorf("invalid parameters for message cache; gossip slots (%d) cannot be larger than history slots (%d)",
gossip, history)
panic(err)
}
return &MessageCache{
msgs: make(map[string]*Message),
peertx: make(map[string]map[peer.ID]int),
history: make([][]CacheEntry, history),
gossip: gossip,
msgID: func(msg *Message) string {
return DefaultMsgIdFn(msg.Message)
},
}
}
type MessageCache struct {
msgs map[string]*Message
peertx map[string]map[peer.ID]int
history [][]CacheEntry
gossip int
msgID func(*Message) string
}
func (mc *MessageCache) SetMsgIdFn(msgID func(*Message) string) {
mc.msgID = msgID
}
type CacheEntry struct {
mid string
bitmask []byte
}
func (mc *MessageCache) Put(msg *Message) {
mid := mc.msgID(msg)
mc.msgs[mid] = msg
mc.history[0] = append(mc.history[0], CacheEntry{mid: mid, bitmask: msg.GetBitmask()})
}
func (mc *MessageCache) Get(mid string) (*Message, bool) {
m, ok := mc.msgs[mid]
return m, ok
}
func (mc *MessageCache) GetForPeer(mid string, p peer.ID) (*Message, int, bool) {
m, ok := mc.msgs[mid]
if !ok {
return nil, 0, false
}
tx, ok := mc.peertx[mid]
if !ok {
tx = make(map[peer.ID]int)
mc.peertx[mid] = tx
}
tx[p]++
return m, tx[p], true
}
func (mc *MessageCache) GetGossipIDs(bitmask []byte) []string {
var mids []string
for _, entries := range mc.history[:mc.gossip] {
for _, entry := range entries {
if bytes.Equal(entry.bitmask, bitmask) {
mids = append(mids, entry.mid)
}
}
}
return mids
}
func (mc *MessageCache) Shift() {
last := mc.history[len(mc.history)-1]
for _, entry := range last {
delete(mc.msgs, entry.mid)
delete(mc.peertx, entry.mid)
}
for i := len(mc.history) - 2; i >= 0; i-- {
mc.history[i+1] = mc.history[i]
}
mc.history[0] = nil
}