refactor protocol functionality into ProtoBook

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)
+}

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 {

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
 		}(),

pstoremem/peerstore.go

@@ -7,5 +7,6 @@ func NewPeerstore() pstore.Peerstore {
 	return pstore.NewPeerstore(
 		NewKeyBook(),
 		NewAddrBook(),
+		NewProtoBook(),
 		NewPeerMetadata())
 }

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
+}