diff --git a/interface.go b/interface.go index c669ba9..e3d988c 100644 --- a/interface.go +++ b/interface.go @@ -56,17 +56,13 @@ type Peerstore interface { KeyBook PeerMetadata Metrics + ProtoBook // PeerInfo returns a peer.PeerInfo struct for given peer.ID. // This is a small slice of the information Peerstore has on // that peer, useful to other services. PeerInfo(peer.ID) PeerInfo - GetProtocols(peer.ID) ([]string, error) - AddProtocols(peer.ID, ...string) error - SetProtocols(peer.ID, ...string) error - SupportsProtocols(peer.ID, ...string) ([]string, error) - // Peers returns all of the peer IDs stored across all inner stores. Peers() peer.IDSlice } @@ -142,3 +138,11 @@ type KeyBook interface { // PeersWithKeys returns all the peer IDs stored in the KeyBook PeersWithKeys() peer.IDSlice } + +// ProtoBook tracks the protocols supported by peers +type ProtoBook interface { + GetProtocols(peer.ID) ([]string, error) + AddProtocols(peer.ID, ...string) error + SetProtocols(peer.ID, ...string) error + SupportsProtocols(peer.ID, ...string) ([]string, error) +} diff --git a/peerstore.go b/peerstore.go index fde9cd4..ea29a17 100644 --- a/peerstore.go +++ b/peerstore.go @@ -3,74 +3,30 @@ package peerstore import ( "fmt" "io" - "sync" peer "github.com/libp2p/go-libp2p-peer" ) var _ Peerstore = (*peerstore)(nil) -const maxInternedProtocols = 64 -const maxInternedProtocolSize = 128 - -type segment struct { - lk sync.RWMutex - interned map[string]string - protocols map[peer.ID]map[string]struct{} -} - -type segments [256]*segment - -func (s *segments) get(id peer.ID) *segment { - b := []byte(id) - return s[b[len(b)-1]] -} - -func (s *segment) internProtocol(proto string) string { - if len(proto) > maxInternedProtocolSize { - return proto - } - - if interned, ok := s.interned[proto]; ok { - return interned - } - - if len(s.interned) >= maxInternedProtocols { - s.interned = make(map[string]string, maxInternedProtocols) - } - - s.interned[proto] = proto - return proto -} - type peerstore struct { Metrics KeyBook AddrBook + ProtoBook PeerMetadata - - // segments for protocol information - segments segments } // NewPeerstore creates a data structure that stores peer data, backed by the // supplied implementations of KeyBook, AddrBook and PeerMetadata. -func NewPeerstore(kb KeyBook, ab AddrBook, md PeerMetadata) Peerstore { +func NewPeerstore(kb KeyBook, ab AddrBook, pb ProtoBook, md PeerMetadata) Peerstore { return &peerstore{ KeyBook: kb, AddrBook: ab, + ProtoBook: pb, PeerMetadata: md, Metrics: NewMetrics(), - segments: func() (ret segments) { - for i := range ret { - ret[i] = &segment{ - interned: make(map[string]string), - protocols: make(map[peer.ID]map[string]struct{}), - } - } - return ret - }(), } } @@ -86,6 +42,7 @@ func (ps *peerstore) Close() (err error) { weakClose("keybook", ps.KeyBook) weakClose("addressbook", ps.AddrBook) + weakClose("protobook", ps.ProtoBook) weakClose("peermetadata", ps.PeerMetadata) if len(errs) > 0 { @@ -117,67 +74,6 @@ func (ps *peerstore) PeerInfo(p peer.ID) PeerInfo { } } -func (ps *peerstore) SetProtocols(p peer.ID, protos ...string) error { - s := ps.segments.get(p) - s.lk.Lock() - defer s.lk.Unlock() - - newprotos := make(map[string]struct{}, len(protos)) - for _, proto := range protos { - newprotos[s.internProtocol(proto)] = struct{}{} - } - - s.protocols[p] = newprotos - - return nil -} - -func (ps *peerstore) AddProtocols(p peer.ID, protos ...string) error { - s := ps.segments.get(p) - s.lk.Lock() - defer s.lk.Unlock() - - protomap, ok := s.protocols[p] - if !ok { - protomap = make(map[string]struct{}) - s.protocols[p] = protomap - } - - for _, proto := range protos { - protomap[s.internProtocol(proto)] = struct{}{} - } - - return nil -} - -func (ps *peerstore) GetProtocols(p peer.ID) ([]string, error) { - s := ps.segments.get(p) - s.lk.RLock() - defer s.lk.RUnlock() - - out := make([]string, 0, len(s.protocols)) - for k := range s.protocols[p] { - out = append(out, k) - } - - return out, nil -} - -func (ps *peerstore) SupportsProtocols(p peer.ID, protos ...string) ([]string, error) { - s := ps.segments.get(p) - s.lk.RLock() - defer s.lk.RUnlock() - - out := make([]string, 0, len(protos)) - for _, proto := range protos { - if _, ok := s.protocols[p][proto]; ok { - out = append(out, proto) - } - } - - return out, nil -} - func PeerInfos(ps Peerstore, peers peer.IDSlice) []PeerInfo { pi := make([]PeerInfo, len(peers)) for i, p := range peers { diff --git a/pstoremem/addr_book.go b/pstoremem/addr_book.go index cd7a920..6baee9e 100644 --- a/pstoremem/addr_book.go +++ b/pstoremem/addr_book.go @@ -27,18 +27,18 @@ func (e *expiringAddr) ExpiredBy(t time.Time) bool { return t.After(e.Expires) } -type segments [256]*segment +type addrSegments [256]*addrSegment -type segment struct { +type addrSegment struct { size uint32 lk sync.RWMutex addrs map[peer.ID]map[string]*expiringAddr } -func (s *segments) get(id peer.ID) *segment { +func (s *addrSegments) get(id peer.ID) *addrSegment { b := []byte(id) - return s[b[len(b)-1]%byte(255)] + return s[b[len(b)-1]] } // memoryAddrBook manages addresses. @@ -47,7 +47,7 @@ type memoryAddrBook struct { // Use pointers to save memory. Maps always leave some fraction of their // space unused. storing the *values* directly in the map will // drastically increase the space waste. In our case, by 6x. - segments segments + segments addrSegments ctx context.Context cancel func() @@ -61,9 +61,9 @@ func NewAddrBook() pstore.AddrBook { ctx, cancel := context.WithCancel(context.Background()) ab := &memoryAddrBook{ - segments: func() (ret segments) { + segments: func() (ret addrSegments) { for i, _ := range ret { - ret[i] = &segment{addrs: make(map[peer.ID]map[string]*expiringAddr)} + ret[i] = &addrSegment{addrs: make(map[peer.ID]map[string]*expiringAddr)} } return ret }(), diff --git a/pstoremem/peerstore.go b/pstoremem/peerstore.go index 7d87313..c7cbd67 100644 --- a/pstoremem/peerstore.go +++ b/pstoremem/peerstore.go @@ -7,5 +7,6 @@ func NewPeerstore() pstore.Peerstore { return pstore.NewPeerstore( NewKeyBook(), NewAddrBook(), + NewProtoBook(), NewPeerMetadata()) } diff --git a/pstoremem/protobook.go b/pstoremem/protobook.go new file mode 100644 index 0000000..9074696 --- /dev/null +++ b/pstoremem/protobook.go @@ -0,0 +1,123 @@ +package pstoremem + +import ( + "sync" + + peer "github.com/libp2p/go-libp2p-peer" + + pstore "github.com/libp2p/go-libp2p-peerstore" +) + +const maxInternedProtocols = 64 +const maxInternedProtocolSize = 128 + +type protoSegment struct { + lk sync.RWMutex + interned map[string]string + protocols map[peer.ID]map[string]struct{} +} + +type protoSegments [256]*protoSegment + +func (s *protoSegments) get(id peer.ID) *protoSegment { + b := []byte(id) + return s[b[len(b)-1]] +} + +func (s *protoSegment) internProtocol(proto string) string { + if len(proto) > maxInternedProtocolSize { + return proto + } + + if interned, ok := s.interned[proto]; ok { + return interned + } + + if len(s.interned) >= maxInternedProtocols { + s.interned = make(map[string]string, maxInternedProtocols) + } + + s.interned[proto] = proto + return proto +} + +type memoryProtoBook struct { + segments protoSegments +} + +var _ pstore.ProtoBook = (*memoryProtoBook)(nil) + +func NewProtoBook() pstore.ProtoBook { + return &memoryProtoBook{ + segments: func() (ret protoSegments) { + for i := range ret { + ret[i] = &protoSegment{ + interned: make(map[string]string), + protocols: make(map[peer.ID]map[string]struct{}), + } + } + return ret + }(), + } +} + +func (pb *memoryProtoBook) SetProtocols(p peer.ID, protos ...string) error { + s := pb.segments.get(p) + s.lk.Lock() + defer s.lk.Unlock() + + newprotos := make(map[string]struct{}, len(protos)) + for _, proto := range protos { + newprotos[s.internProtocol(proto)] = struct{}{} + } + + s.protocols[p] = newprotos + + return nil +} + +func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...string) error { + s := pb.segments.get(p) + s.lk.Lock() + defer s.lk.Unlock() + + protomap, ok := s.protocols[p] + if !ok { + protomap = make(map[string]struct{}) + s.protocols[p] = protomap + } + + for _, proto := range protos { + protomap[s.internProtocol(proto)] = struct{}{} + } + + return nil +} + +func (pb *memoryProtoBook) GetProtocols(p peer.ID) ([]string, error) { + s := pb.segments.get(p) + s.lk.RLock() + defer s.lk.RUnlock() + + out := make([]string, 0, len(s.protocols)) + for k := range s.protocols[p] { + out = append(out, k) + } + + return out, nil +} + +func (pb *memoryProtoBook) SupportsProtocols(p peer.ID, protos ...string) ([]string, error) { + s := pb.segments.get(p) + s.lk.RLock() + defer s.lk.RUnlock() + + out := make([]string, 0, len(protos)) + for _, proto := range protos { + if _, ok := s.protocols[p][proto]; ok { + out = append(out, proto) + } + } + + return out, nil +}