ceremonyclient/node/p2p/peer_info_manager.go
2024-03-21 02:14:45 -05:00

147 lines
3.2 KiB
Go

package p2p
import (
"bytes"
"sync"
"go.uber.org/zap"
)
type PeerInfoManager interface {
Start()
Stop()
AddPeerInfo(manifest *PeerManifest)
GetPeerInfo(peerId []byte) *PeerManifest
GetPeerMap() map[string]*PeerManifest
GetPeersBySpeed() [][]byte
}
type Capability struct {
ProtocolIdentifier uint32
AdditionalMetadata []byte
}
type PeerManifest struct {
PeerId []byte
Difficulty uint32
DifficultyMetric int64
Commit_16Metric int64
Commit_128Metric int64
Commit_1024Metric int64
Commit_65536Metric int64
Proof_16Metric int64
Proof_128Metric int64
Proof_1024Metric int64
Proof_65536Metric int64
Cores uint32
Memory []byte
Storage []byte
Capabilities []Capability
MasterHeadFrame uint64
Bandwidth uint64
}
type InMemoryPeerInfoManager struct {
logger *zap.Logger
peerInfoCh chan *PeerManifest
quitCh chan struct{}
peerInfoMx sync.RWMutex
peerMap map[string]*PeerManifest
fastestPeers []*PeerManifest
}
var _ PeerInfoManager = (*InMemoryPeerInfoManager)(nil)
func NewInMemoryPeerInfoManager(logger *zap.Logger) *InMemoryPeerInfoManager {
return &InMemoryPeerInfoManager{
logger: logger,
peerInfoCh: make(chan *PeerManifest),
fastestPeers: []*PeerManifest{},
peerMap: make(map[string]*PeerManifest),
}
}
func (m *InMemoryPeerInfoManager) Start() {
go func() {
for {
select {
case manifest := <-m.peerInfoCh:
m.peerInfoMx.Lock()
m.peerMap[string(manifest.PeerId)] = manifest
m.searchAndInsertPeer(manifest)
m.peerInfoMx.Unlock()
case <-m.quitCh:
return
}
}
}()
}
func (m *InMemoryPeerInfoManager) Stop() {
go func() {
m.quitCh <- struct{}{}
}()
}
func (m *InMemoryPeerInfoManager) AddPeerInfo(manifest *PeerManifest) {
go func() {
m.peerInfoCh <- manifest
}()
}
func (m *InMemoryPeerInfoManager) GetPeerInfo(peerId []byte) *PeerManifest {
m.peerInfoMx.RLock()
manifest, ok := m.peerMap[string(peerId)]
m.peerInfoMx.RUnlock()
if !ok {
return nil
}
return manifest
}
func (m *InMemoryPeerInfoManager) GetPeerMap() map[string]*PeerManifest {
data := make(map[string]*PeerManifest)
m.peerInfoMx.RLock()
for k, v := range m.peerMap {
data[k] = v
}
m.peerInfoMx.RUnlock()
return data
}
func (m *InMemoryPeerInfoManager) GetPeersBySpeed() [][]byte {
result := [][]byte{}
m.peerInfoMx.RLock()
for _, info := range m.fastestPeers {
result = append(result, info.PeerId)
}
m.peerInfoMx.RUnlock()
return result
}
// blatantly lifted from slices.BinarySearchFunc, optimized for direct insertion
// and uint64 comparison without overflow
func (m *InMemoryPeerInfoManager) searchAndInsertPeer(manifest *PeerManifest) {
n := len(m.fastestPeers)
i, j := 0, n
for i < j {
h := int(uint(i+j) >> 1)
if m.fastestPeers[h].Bandwidth > manifest.Bandwidth {
i = h + 1
} else {
j = h
}
}
if i < n && m.fastestPeers[i].Bandwidth == manifest.Bandwidth &&
bytes.Equal(m.fastestPeers[i].PeerId, manifest.PeerId) {
m.fastestPeers[i] = manifest
} else {
m.fastestPeers = append(m.fastestPeers, new(PeerManifest))
copy(m.fastestPeers[i+1:], m.fastestPeers[i:])
m.fastestPeers[i] = manifest
}
}