peermanager.go 6.6 KB
Newer Older
1 2 3 4
package peermanager

import (
	"context"
dirkmc's avatar
dirkmc committed
5
	"sync"
6
	"time"
7

8
	logging "github.com/ipfs/go-log"
dirkmc's avatar
dirkmc committed
9
	"github.com/ipfs/go-metrics-interface"
10

dirkmc's avatar
dirkmc committed
11
	cid "github.com/ipfs/go-cid"
Raúl Kripalani's avatar
Raúl Kripalani committed
12
	peer "github.com/libp2p/go-libp2p-core/peer"
13 14
)

15 16
var log = logging.Logger("bs:peermgr")

ZenGround0's avatar
ZenGround0 committed
17
// PeerQueue provides a queue of messages to be sent for a single peer.
18
type PeerQueue interface {
dirkmc's avatar
dirkmc committed
19 20 21
	AddBroadcastWantHaves([]cid.Cid)
	AddWants([]cid.Cid, []cid.Cid)
	AddCancels([]cid.Cid)
22
	ResponseReceived(at time.Time, ks []cid.Cid)
23
	Startup()
24 25 26
	Shutdown()
}

dirkmc's avatar
dirkmc committed
27 28 29 30 31
type Session interface {
	ID() uint64
	SignalAvailability(peer.ID, bool)
}

32
// PeerQueueFactory provides a function that will create a PeerQueue.
33
type PeerQueueFactory func(ctx context.Context, p peer.ID) PeerQueue
34

35
// PeerManager manages a pool of peers and sends messages to peers in the pool.
36
type PeerManager struct {
dirkmc's avatar
dirkmc committed
37 38
	// sync access to peerQueues and peerWantManager
	pqLk sync.RWMutex
39
	// peerQueues -- interact through internal utility functions get/set/remove/iterate
40
	peerQueues map[peer.ID]PeerQueue
dirkmc's avatar
dirkmc committed
41
	pwm        *peerWantManager
42

43 44
	createPeerQueue PeerQueueFactory
	ctx             context.Context
dirkmc's avatar
dirkmc committed
45 46 47 48 49 50

	psLk         sync.RWMutex
	sessions     map[uint64]Session
	peerSessions map[peer.ID]map[uint64]struct{}

	self peer.ID
51 52
}

53
// New creates a new PeerManager, given a context and a peerQueueFactory.
dirkmc's avatar
dirkmc committed
54 55
func New(ctx context.Context, createPeerQueue PeerQueueFactory, self peer.ID) *PeerManager {
	wantGauge := metrics.NewCtx(ctx, "wantlist_total", "Number of items in wantlist.").Gauge()
56
	return &PeerManager{
57
		peerQueues:      make(map[peer.ID]PeerQueue),
dirkmc's avatar
dirkmc committed
58
		pwm:             newPeerWantManager(wantGauge),
59 60
		createPeerQueue: createPeerQueue,
		ctx:             ctx,
dirkmc's avatar
dirkmc committed
61 62 63 64
		self:            self,

		sessions:     make(map[uint64]Session),
		peerSessions: make(map[peer.ID]map[uint64]struct{}),
65 66 67
	}
}

dirkmc's avatar
dirkmc committed
68 69 70 71 72
func (pm *PeerManager) AvailablePeers() []peer.ID {
	// TODO: Rate-limit peers
	return pm.ConnectedPeers()
}

73
// ConnectedPeers returns a list of peers this PeerManager is managing.
74
func (pm *PeerManager) ConnectedPeers() []peer.ID {
dirkmc's avatar
dirkmc committed
75 76 77
	pm.pqLk.RLock()
	defer pm.pqLk.RUnlock()

78
	peers := make([]peer.ID, 0, len(pm.peerQueues))
hannahhoward's avatar
hannahhoward committed
79
	for p := range pm.peerQueues {
80
		peers = append(peers, p)
hannahhoward's avatar
hannahhoward committed
81
	}
82
	return peers
83 84
}

85 86
// Connected is called to add a new peer to the pool, and send it an initial set
// of wants.
87
func (pm *PeerManager) Connected(p peer.ID) {
dirkmc's avatar
dirkmc committed
88 89 90
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

hannahhoward's avatar
hannahhoward committed
91
	pq := pm.getOrCreate(p)
92 93

	// Inform the peer want manager that there's a new peer
94
	wants := pm.pwm.addPeer(p)
95
	// Broadcast any live want-haves to the newly connected peers
96
	pq.AddBroadcastWantHaves(wants)
97 98
	// Inform the sessions that the peer has connected
	pm.signalAvailability(p, true)
99 100
}

101 102
// Disconnected is called to remove a peer from the pool.
func (pm *PeerManager) Disconnected(p peer.ID) {
dirkmc's avatar
dirkmc committed
103 104 105
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

106
	pq, ok := pm.peerQueues[p]
107

hannahhoward's avatar
hannahhoward committed
108 109 110 111
	if !ok {
		return
	}

dirkmc's avatar
dirkmc committed
112 113 114 115
	// Inform the sessions that the peer has disconnected
	pm.signalAvailability(p, false)

	// Clean up the peer
116
	delete(pm.peerQueues, p)
117
	pq.Shutdown()
118
	pm.pwm.removePeer(p)
119 120
}

121 122 123 124 125 126 127 128 129 130 131 132 133
// ResponseReceived is called when a message is received from the network.
// ks is the set of blocks, HAVEs and DONT_HAVEs in the message
// Note that this is just used to calculate latency.
func (pm *PeerManager) ResponseReceived(p peer.ID, at time.Time, ks []cid.Cid) {
	pm.pqLk.Lock()
	pq, ok := pm.peerQueues[p]
	pm.pqLk.Unlock()

	if ok {
		pq.ResponseReceived(at, ks)
	}
}

134 135 136 137
// BroadcastWantHaves broadcasts want-haves to all peers (used by the session
// to discover seeds).
// For each peer it filters out want-haves that have previously been sent to
// the peer.
dirkmc's avatar
dirkmc committed
138 139 140 141
func (pm *PeerManager) BroadcastWantHaves(ctx context.Context, wantHaves []cid.Cid) {
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

142
	for p, ks := range pm.pwm.prepareBroadcastWantHaves(wantHaves) {
143 144
		if pq, ok := pm.peerQueues[p]; ok {
			pq.AddBroadcastWantHaves(ks)
hannahhoward's avatar
hannahhoward committed
145
		}
dirkmc's avatar
dirkmc committed
146 147 148
	}
}

149 150
// SendWants sends the given want-blocks and want-haves to the given peer.
// It filters out wants that have previously been sent to the peer.
dirkmc's avatar
dirkmc committed
151 152 153 154
func (pm *PeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) {
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

155
	if pq, ok := pm.peerQueues[p]; ok {
156
		wblks, whvs := pm.pwm.prepareSendWants(p, wantBlocks, wantHaves)
157
		pq.AddWants(wblks, whvs)
dirkmc's avatar
dirkmc committed
158 159 160
	}
}

161 162
// SendCancels sends cancels for the given keys to all peers who had previously
// received a want for those keys.
dirkmc's avatar
dirkmc committed
163 164 165 166 167
func (pm *PeerManager) SendCancels(ctx context.Context, cancelKs []cid.Cid) {
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

	// Send a CANCEL to each peer that has been sent a want-block or want-have
168
	for p, ks := range pm.pwm.prepareSendCancels(cancelKs) {
169 170
		if pq, ok := pm.peerQueues[p]; ok {
			pq.AddCancels(ks)
171 172 173
		}
	}
}
174

175
// CurrentWants returns the list of pending wants (both want-haves and want-blocks).
dirkmc's avatar
dirkmc committed
176 177 178 179
func (pm *PeerManager) CurrentWants() []cid.Cid {
	pm.pqLk.RLock()
	defer pm.pqLk.RUnlock()

180
	return pm.pwm.getWants()
181 182 183 184 185 186 187
}

// CurrentWantBlocks returns the list of pending want-blocks
func (pm *PeerManager) CurrentWantBlocks() []cid.Cid {
	pm.pqLk.RLock()
	defer pm.pqLk.RUnlock()

188
	return pm.pwm.getWantBlocks()
dirkmc's avatar
dirkmc committed
189 190
}

191
// CurrentWantHaves returns the list of pending want-haves
dirkmc's avatar
dirkmc committed
192 193 194 195
func (pm *PeerManager) CurrentWantHaves() []cid.Cid {
	pm.pqLk.RLock()
	defer pm.pqLk.RUnlock()

196
	return pm.pwm.getWantHaves()
dirkmc's avatar
dirkmc committed
197 198
}

199 200
func (pm *PeerManager) getOrCreate(p peer.ID) PeerQueue {
	pq, ok := pm.peerQueues[p]
201
	if !ok {
202
		pq = pm.createPeerQueue(pm.ctx, p)
203
		pq.Startup()
204
		pm.peerQueues[p] = pq
205
	}
206
	return pq
207
}
dirkmc's avatar
dirkmc committed
208

209 210
// RegisterSession tells the PeerManager that the given session is interested
// in events about the given peer.
dirkmc's avatar
dirkmc committed
211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
func (pm *PeerManager) RegisterSession(p peer.ID, s Session) bool {
	pm.psLk.Lock()
	defer pm.psLk.Unlock()

	if _, ok := pm.sessions[s.ID()]; !ok {
		pm.sessions[s.ID()] = s
	}

	if _, ok := pm.peerSessions[p]; !ok {
		pm.peerSessions[p] = make(map[uint64]struct{})
	}
	pm.peerSessions[p][s.ID()] = struct{}{}

	_, ok := pm.peerQueues[p]
	return ok
}

228 229
// UnregisterSession tells the PeerManager that the given session is no longer
// interested in PeerManager events.
dirkmc's avatar
dirkmc committed
230 231 232 233 234 235 236 237 238 239 240 241 242 243
func (pm *PeerManager) UnregisterSession(ses uint64) {
	pm.psLk.Lock()
	defer pm.psLk.Unlock()

	for p := range pm.peerSessions {
		delete(pm.peerSessions[p], ses)
		if len(pm.peerSessions[p]) == 0 {
			delete(pm.peerSessions, p)
		}
	}

	delete(pm.sessions, ses)
}

244 245
// signalAvailability is called when a peer's connectivity changes.
// It informs interested sessions.
dirkmc's avatar
dirkmc committed
246
func (pm *PeerManager) signalAvailability(p peer.ID, isConnected bool) {
247 248 249 250 251 252 253
	sesIds, ok := pm.peerSessions[p]
	if !ok {
		return
	}
	for sesId := range sesIds {
		if s, ok := pm.sessions[sesId]; ok {
			s.SignalAvailability(p, isConnected)
dirkmc's avatar
dirkmc committed
254 255 256
		}
	}
}