ceremonyclient/node/p2p/peer_info_manager.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
	}
}