From 5b55f0f78ddc8d7900d5bdd82d426c6f9dceae8d Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Fri, 6 Mar 2020 00:45:37 +0200
Subject: [PATCH] track peer (re)transmissions in message cache

---
 mcache.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/mcache.go b/mcache.go
index e085297..86154cb 100644
--- a/mcache.go
+++ b/mcache.go
@@ -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]