track peer (re)transmissions in message cache

This commit is contained in:
vyzo 2020-03-06 00:45:37 +02:00
parent f9d29c47b6
commit 5b55f0f78d
1 changed files with 21 additions and 0 deletions

View File

@ -4,6 +4,8 @@ import (
"fmt"
pb "github.com/libp2p/go-libp2p-pubsub/pb"
"github.com/libp2p/go-libp2p-core/peer"
)
// NewMessageCache creates a sliding window cache that remembers messages for as
@ -26,6 +28,7 @@ func NewMessageCache(gossip, history int) *MessageCache {
}
return &MessageCache{
msgs: make(map[string]*pb.Message),
peertx: make(map[string]map[peer.ID]int),
history: make([][]CacheEntry, history),
gossip: gossip,
msgID: DefaultMsgIdFn,
@ -34,6 +37,7 @@ func NewMessageCache(gossip, history int) *MessageCache {
type MessageCache struct {
msgs map[string]*pb.Message
peertx map[string]map[peer.ID]int
history [][]CacheEntry
gossip int
msgID MsgIdFunction
@ -59,6 +63,22 @@ func (mc *MessageCache) Get(mid string) (*pb.Message, bool) {
return m, ok
}
func (mc *MessageCache) GetForPeer(mid string, p peer.ID) (*pb.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(topic string) []string {
var mids []string
for _, entries := range mc.history[:mc.gossip] {
@ -78,6 +98,7 @@ 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]