sessionpeermanager.go 2.96 KB
Newer Older
1 2 3 4
package sessionpeermanager

import (
	"fmt"
5
	"sync"
6

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

Raúl Kripalani's avatar
Raúl Kripalani committed
9
	peer "github.com/libp2p/go-libp2p-core/peer"
10 11
)

dirkmc's avatar
dirkmc committed
12 13
var log = logging.Logger("bs:sprmgr")

hannahhoward's avatar
hannahhoward committed
14
const (
15 16 17
	// Connection Manager tag value for session peers. Indicates to connection
	// manager that it should keep the connection to the peer.
	sessionPeerTagValue = 5
hannahhoward's avatar
hannahhoward committed
18 19
)

20 21 22 23 24 25
// PeerTagger is an interface for tagging peers with metadata
type PeerTagger interface {
	TagPeer(peer.ID, string, int)
	UntagPeer(p peer.ID, tag string)
}

26 27
// SessionPeerManager keeps track of peers for a session, and takes care of
// ConnectionManager tagging.
28
type SessionPeerManager struct {
29 30
	tagger PeerTagger
	tag    string
31

32 33 34
	plk             sync.RWMutex
	peers           map[peer.ID]struct{}
	peersDiscovered bool
35 36
}

37
// New creates a new SessionPeerManager
38 39 40 41 42
func New(id uint64, tagger PeerTagger) *SessionPeerManager {
	return &SessionPeerManager{
		tag:    fmt.Sprint("bs-ses-", id),
		tagger: tagger,
		peers:  make(map[peer.ID]struct{}),
43
	}
hannahhoward's avatar
hannahhoward committed
44 45
}

46 47 48 49 50
// AddPeer adds the peer to the SessionPeerManager.
// Returns true if the peer is a new peer, false if it already existed.
func (spm *SessionPeerManager) AddPeer(p peer.ID) bool {
	spm.plk.Lock()
	defer spm.plk.Unlock()
dirkmc's avatar
dirkmc committed
51

52 53 54
	// Check if the peer is a new peer
	if _, ok := spm.peers[p]; ok {
		return false
dirkmc's avatar
dirkmc committed
55
	}
hannahhoward's avatar
hannahhoward committed
56

57 58
	spm.peers[p] = struct{}{}
	spm.peersDiscovered = true
59

60 61 62
	// Tag the peer with the ConnectionManager so it doesn't discard the
	// connection
	spm.tagger.TagPeer(p, spm.tag, sessionPeerTagValue)
63

64 65
	log.Infof("Added peer %s to session: %d peers\n", p, len(spm.peers))
	return true
66 67
}

68 69 70 71 72
// RemovePeer removes the peer from the SessionPeerManager.
// Returns true if the peer was removed, false if it did not exist.
func (spm *SessionPeerManager) RemovePeer(p peer.ID) bool {
	spm.plk.Lock()
	defer spm.plk.Unlock()
hannahhoward's avatar
hannahhoward committed
73

74 75
	if _, ok := spm.peers[p]; !ok {
		return false
hannahhoward's avatar
hannahhoward committed
76 77
	}

78 79 80
	delete(spm.peers, p)
	spm.tagger.UntagPeer(p, spm.tag)
	return true
hannahhoward's avatar
hannahhoward committed
81 82
}

83 84 85 86 87 88
// PeersDiscovered indicates whether peers have been discovered yet.
// Returns true once a peer has been discovered by the session (even if all
// peers are later removed from the session).
func (spm *SessionPeerManager) PeersDiscovered() bool {
	spm.plk.RLock()
	defer spm.plk.RUnlock()
89

90
	return spm.peersDiscovered
91 92
}

93 94 95
func (spm *SessionPeerManager) Peers() []peer.ID {
	spm.plk.RLock()
	defer spm.plk.RUnlock()
96

97 98 99
	peers := make([]peer.ID, 0, len(spm.peers))
	for p := range spm.peers {
		peers = append(peers, p)
100
	}
hannahhoward's avatar
hannahhoward committed
101

102
	return peers
hannahhoward's avatar
hannahhoward committed
103 104
}

105 106 107
func (spm *SessionPeerManager) HasPeers() bool {
	spm.plk.RLock()
	defer spm.plk.RUnlock()
108

109
	return len(spm.peers) > 0
110 111
}

Dirk McCormick's avatar
Dirk McCormick committed
112 113 114 115 116 117 118 119
func (spm *SessionPeerManager) HasPeer(p peer.ID) bool {
	spm.plk.RLock()
	defer spm.plk.RUnlock()

	_, ok := spm.peers[p]
	return ok
}

120 121 122 123
// Shutdown untags all the peers
func (spm *SessionPeerManager) Shutdown() {
	spm.plk.Lock()
	defer spm.plk.Unlock()
124

125 126 127
	// Untag the peers with the ConnectionManager so that it can release
	// connections to those peers
	for p := range spm.peers {
128
		spm.tagger.UntagPeer(p, spm.tag)
129 130
	}
}