forked from libp2p/go-libp2p
-
Notifications
You must be signed in to change notification settings - Fork 0
/
peerstore.go
114 lines (100 loc) · 2.59 KB
/
peerstore.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package pstoremem
import (
"fmt"
"io"
"github.com/thusharprakash/go-libp2p/core/peer"
"github.com/thusharprakash/go-libp2p/core/peerstore"
pstore "github.com/thusharprakash/go-libp2p/p2p/host/peerstore"
)
type pstoremem struct {
peerstore.Metrics
*memoryKeyBook
*memoryAddrBook
*memoryProtoBook
*memoryPeerMetadata
}
var _ peerstore.Peerstore = &pstoremem{}
type Option interface{}
// NewPeerstore creates an in-memory thread-safe collection of peers.
// It's the caller's responsibility to call RemovePeer to ensure
// that memory consumption of the peerstore doesn't grow unboundedly.
func NewPeerstore(opts ...Option) (ps *pstoremem, err error) {
ab := NewAddrBook()
defer func() {
if err != nil {
ab.Close()
}
}()
var protoBookOpts []ProtoBookOption
for _, opt := range opts {
switch o := opt.(type) {
case ProtoBookOption:
protoBookOpts = append(protoBookOpts, o)
case AddrBookOption:
o(ab)
default:
return nil, fmt.Errorf("unexpected peer store option: %v", o)
}
}
pb, err := NewProtoBook(protoBookOpts...)
if err != nil {
return nil, err
}
return &pstoremem{
Metrics: pstore.NewMetrics(),
memoryKeyBook: NewKeyBook(),
memoryAddrBook: ab,
memoryProtoBook: pb,
memoryPeerMetadata: NewPeerMetadata(),
}, nil
}
func (ps *pstoremem) Close() (err error) {
var errs []error
weakClose := func(name string, c interface{}) {
if cl, ok := c.(io.Closer); ok {
if err = cl.Close(); err != nil {
errs = append(errs, fmt.Errorf("%s error: %s", name, err))
}
}
}
weakClose("keybook", ps.memoryKeyBook)
weakClose("addressbook", ps.memoryAddrBook)
weakClose("protobook", ps.memoryProtoBook)
weakClose("peermetadata", ps.memoryPeerMetadata)
if len(errs) > 0 {
return fmt.Errorf("failed while closing peerstore; err(s): %q", errs)
}
return nil
}
func (ps *pstoremem) Peers() peer.IDSlice {
set := map[peer.ID]struct{}{}
for _, p := range ps.PeersWithKeys() {
set[p] = struct{}{}
}
for _, p := range ps.PeersWithAddrs() {
set[p] = struct{}{}
}
pps := make(peer.IDSlice, 0, len(set))
for p := range set {
pps = append(pps, p)
}
return pps
}
func (ps *pstoremem) PeerInfo(p peer.ID) peer.AddrInfo {
return peer.AddrInfo{
ID: p,
Addrs: ps.memoryAddrBook.Addrs(p),
}
}
// RemovePeer removes entries associated with a peer from:
// * the KeyBook
// * the ProtoBook
// * the PeerMetadata
// * the Metrics
// It DOES NOT remove the peer from the AddrBook.
func (ps *pstoremem) RemovePeer(p peer.ID) {
ps.memoryKeyBook.RemovePeer(p)
ps.memoryProtoBook.RemovePeer(p)
ps.memoryPeerMetadata.RemovePeer(p)
ps.Metrics.RemovePeer(p)
}