lotus/chain/block_receipt_tracker.go
Steven Allen dcb49dc8ee
refactor: update cache to the new generic version ()
- Adds type safety.
- Reduces allocations.
- Fixes the drand cache (was storing by value, but retrieving by pointer)
2023-03-13 15:29:09 -07:00

72 lines
1.3 KiB
Go

package chain
import (
"sort"
"sync"
"time"
lru "github.com/hashicorp/golang-lru/v2"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/filecoin-project/lotus/build"
"github.com/filecoin-project/lotus/chain/types"
)
type blockReceiptTracker struct {
lk sync.Mutex
// using an LRU cache because i don't want to handle all the edge cases for
// manual cleanup and maintenance of a fixed size set
cache *lru.Cache[types.TipSetKey, *peerSet]
}
type peerSet struct {
peers map[peer.ID]time.Time
}
func newBlockReceiptTracker() *blockReceiptTracker {
c, _ := lru.New[types.TipSetKey, *peerSet](512)
return &blockReceiptTracker{
cache: c,
}
}
func (brt *blockReceiptTracker) Add(p peer.ID, ts *types.TipSet) {
brt.lk.Lock()
defer brt.lk.Unlock()
val, ok := brt.cache.Get(ts.Key())
if !ok {
pset := &peerSet{
peers: map[peer.ID]time.Time{
p: build.Clock.Now(),
},
}
brt.cache.Add(ts.Key(), pset)
return
}
val.peers[p] = build.Clock.Now()
}
func (brt *blockReceiptTracker) GetPeers(ts *types.TipSet) []peer.ID {
brt.lk.Lock()
defer brt.lk.Unlock()
ps, ok := brt.cache.Get(ts.Key())
if !ok {
return nil
}
out := make([]peer.ID, 0, len(ps.peers))
for p := range ps.peers {
out = append(out, p)
}
sort.Slice(out, func(i, j int) bool {
return ps.peers[out[i]].Before(ps.peers[out[j]])
})
return out
}