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

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

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

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

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

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

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

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

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

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

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

	self peer.ID
49 50
}

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

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

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

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

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

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

hannahhoward's avatar
hannahhoward committed
89
	pq := pm.getOrCreate(p)
90 91

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

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

104
	pq, ok := pm.peerQueues[p]
105

hannahhoward's avatar
hannahhoward committed
106 107 108 109
	if !ok {
		return
	}

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

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

119 120 121 122
// 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
123 124 125 126
func (pm *PeerManager) BroadcastWantHaves(ctx context.Context, wantHaves []cid.Cid) {
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

127
	for p, ks := range pm.pwm.prepareBroadcastWantHaves(wantHaves) {
128 129
		if pq, ok := pm.peerQueues[p]; ok {
			pq.AddBroadcastWantHaves(ks)
hannahhoward's avatar
hannahhoward committed
130
		}
dirkmc's avatar
dirkmc committed
131 132 133
	}
}

134 135
// 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
136 137 138 139
func (pm *PeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) {
	pm.pqLk.Lock()
	defer pm.pqLk.Unlock()

140
	if pq, ok := pm.peerQueues[p]; ok {
141
		wblks, whvs := pm.pwm.prepareSendWants(p, wantBlocks, wantHaves)
142
		pq.AddWants(wblks, whvs)
dirkmc's avatar
dirkmc committed
143 144 145
	}
}

146 147
// SendCancels sends cancels for the given keys to all peers who had previously
// received a want for those keys.
dirkmc's avatar
dirkmc committed
148 149 150 151 152
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
153
	for p, ks := range pm.pwm.prepareSendCancels(cancelKs) {
154 155
		if pq, ok := pm.peerQueues[p]; ok {
			pq.AddCancels(ks)
156 157 158
		}
	}
}
159

160
// CurrentWants returns the list of pending wants (both want-haves and want-blocks).
dirkmc's avatar
dirkmc committed
161 162 163 164
func (pm *PeerManager) CurrentWants() []cid.Cid {
	pm.pqLk.RLock()
	defer pm.pqLk.RUnlock()

165
	return pm.pwm.getWants()
166 167 168 169 170 171 172
}

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

173
	return pm.pwm.getWantBlocks()
dirkmc's avatar
dirkmc committed
174 175
}

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

181
	return pm.pwm.getWantHaves()
dirkmc's avatar
dirkmc committed
182 183
}

184 185
func (pm *PeerManager) getOrCreate(p peer.ID) PeerQueue {
	pq, ok := pm.peerQueues[p]
186
	if !ok {
187
		pq = pm.createPeerQueue(pm.ctx, p)
188
		pq.Startup()
189
		pm.peerQueues[p] = pq
190
	}
191
	return pq
192
}
dirkmc's avatar
dirkmc committed
193

194 195
// RegisterSession tells the PeerManager that the given session is interested
// in events about the given peer.
dirkmc's avatar
dirkmc committed
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212
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
}

213 214
// UnregisterSession tells the PeerManager that the given session is no longer
// interested in PeerManager events.
dirkmc's avatar
dirkmc committed
215 216 217 218 219 220 221 222 223 224 225 226 227 228
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)
}

229 230
// signalAvailability is called when a peer's connectivity changes.
// It informs interested sessions.
dirkmc's avatar
dirkmc committed
231
func (pm *PeerManager) signalAvailability(p peer.ID, isConnected bool) {
232 233 234 235 236 237 238
	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
239 240 241
		}
	}
}