peerstore : add GetPeersForProtocol API

core/peerstore/peerstore.go

@@ -246,5 +246,7 @@ type ProtoBook interface {
 	FirstSupportedProtocol(peer.ID, ...protocol.ID) (protocol.ID, error)
 
 	// RemovePeer removes all protocols associated with a peer.
+
+	GetPeersForProtocol(protocol.ID) ([]peer.ID, error)
 	RemovePeer(peer.ID)
 }

p2p/host/peerstore/pstoreds/metadata.go

@@ -42,7 +42,11 @@
 }
 
 func (pm *dsPeerMetadata) Get(p peer.ID, key string) (interface{}, error) {
-	k := pmBase.ChildString(base32.RawStdEncoding.EncodeToString([]byte(p))).ChildString(key)
+	return pm.getFromDS(p.String(), key)
+}
+
+func (pm *dsPeerMetadata) getFromDS(uid string, key string) (interface{}, error) {
+	k := pmBase.ChildString(base32.RawStdEncoding.EncodeToString([]byte(uid))).ChildString(key)
 	value, err := pm.ds.Get(context.TODO(), k)
 	if err != nil {
 		if err == ds.ErrNotFound {
@@ -80,3 +84,16 @@
 		pm.ds.Delete(context.TODO(), ds.NewKey(entry.Key))
 	}
 }
+
+func (pm *dsPeerMetadata) getPeers(p protocol.ID, key string) (interface{}, error) {
+	return pm.getFromDS(string(p), key)
+}
+
+func (pm *dsPeerMetadata) setPeers(ctx context.Context, p protocol.ID, key string, val interface{}) error {
+	k := pmBase.ChildString(base32.RawStdEncoding.EncodeToString([]byte(p))).ChildString(key)
+	var buf pool.Buffer
+	if err := gob.NewEncoder(&buf).Encode(&val); err != nil {
+		return err
+	}
+	return pm.ds.Put(ctx, k, buf.Bytes())
+}

p2p/host/peerstore/pstoreds/protobook.go

@@ -31,10 +31,15 @@
 	}
 }
 
+type peersPerProtocol struct {
+	sync.RWMutex
+}
+
 type dsProtoBook struct {
-	segments  protoSegments
-	meta      pstore.PeerMetadata
-	maxProtos int
+	segments           protoSegments
+	meta               pstore.PeerMetadata
+	maxProtos          int
+	protocolToPeersMap peersPerProtocol
 }
 
 var _ pstore.ProtoBook = (*dsProtoBook)(nil)
@@ -71,9 +76,13 @@
 
 	s := pb.segments.get(p)
 	s.Lock()
-	defer s.Unlock()
+	err := pb.meta.Put(p, "protocols", protomap)
+	s.Unlock()
+	if err != nil {
+		return err
+	}
 
-	return pb.meta.Put(p, "protocols", protomap)
+	return nil
 }
 
 func (pb *dsProtoBook) AddProtocols(p peer.ID, protos ...protocol.ID) error {
@@ -193,4 +202,12 @@
 
 func (pb *dsProtoBook) RemovePeer(p peer.ID) {
 	pb.meta.RemovePeer(p)
+	//TODO: Remove peer from protolist as well.
+}
+
+func (pb *dsProtoBook) GetPeersForProtocol(proto protocol.ID) ([]peer.ID, error) {
+	/* 	pb.protocolToPeersMap.Lock()
+	   	defer pb.protocolToPeersMap.Unlock()
+	   	pb.meta.getPeers(proto, "peers") */
+	return nil, nil
 }

p2p/host/peerstore/pstoremem/protobook.go

@@ -2,6 +2,7 @@ package pstoremem
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 
 	"github.com/libp2p/go-libp2p/core/peer"
@@ -20,7 +21,13 @@ func (s *protoSegments) get(p peer.ID) *protoSegment {
 	return s[byte(p[len(p)-1])]
 }
 
+type peersPerProtocol struct {
+	sync.RWMutex
+	peers map[protocol.ID]map[peer.ID]peer.ID
+}
+
 var errTooManyProtocols = errors.New("too many protocols")
+var errNoPeersForProtocol = errors.New("no peers available for queried protocol")
 
 type memoryProtoBook struct {
 	segments protoSegments
@@ -29,6 +36,8 @@ type memoryProtoBook struct {
 
 	lk       sync.RWMutex
 	interned map[protocol.ID]protocol.ID
+
+	protocolToPeersMap peersPerProtocol
 }
 
 var _ pstore.ProtoBook = (*memoryProtoBook)(nil)
@@ -61,6 +70,7 @@ func NewProtoBook(opts ...ProtoBookOption) (*memoryProtoBook, error) {
 			return nil, err
 		}
 	}
+	pb.protocolToPeersMap.peers = make(map[protocol.ID]map[peer.ID]peer.ID, pb.maxProtos)
 	return pb, nil
 }
 
@@ -92,22 +102,42 @@ func (pb *memoryProtoBook) SetProtocols(p peer.ID, protos ...protocol.ID) error
 	if len(protos) > pb.maxProtos {
 		return errTooManyProtocols
 	}
-
+	fmt.Println("Set Protocols for peer ", p, " protocols are :", protos)
 	newprotos := make(map[protocol.ID]struct{}, len(protos))
 	for _, proto := range protos {
 		newprotos[pb.internProtocol(proto)] = struct{}{}
 	}
-
 	s := pb.segments.get(p)
 	s.Lock()
 	s.protocols[p] = newprotos
 	s.Unlock()
-
+	pb.addPeersPerProtocol(p, protos...)
 	return nil
 }
 
-func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...protocol.ID) error {
+func (pb *memoryProtoBook) addPeersPerProtocol(p peer.ID, protos ...protocol.ID) {
+	fmt.Println("Add Protocols for peer ", p, " protocols are :", protos)
+
+	pb.protocolToPeersMap.Lock()
+	defer pb.protocolToPeersMap.Unlock()
+	for _, proto := range protos {
+		peers, ok := pb.protocolToPeersMap.peers[proto]
+		if !ok {
+			peers = make(map[peer.ID]peer.ID)
+			peers[p] = p
+			pb.protocolToPeersMap.peers[proto] = peers
+		} else {
+			_, ok := peers[p]
+			if !ok {
+				peers[p] = p
+			}
+		}
+	}
+}
+
+func (pb *memoryProtoBook) addProtocolsToSegment(p peer.ID, protos ...protocol.ID) error {
 	s := pb.segments.get(p)
+
 	s.Lock()
 	defer s.Unlock()
 
@@ -126,6 +156,15 @@ func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...protocol.ID) error
 	return nil
 }
 
+func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...protocol.ID) error {
+	err := pb.addProtocolsToSegment(p, protos...)
+	if err != nil {
+		return err
+	}
+	pb.addPeersPerProtocol(p, protos...)
+	return nil
+}
+
 func (pb *memoryProtoBook) GetProtocols(p peer.ID) ([]protocol.ID, error) {
 	s := pb.segments.get(p)
 	s.RLock()
@@ -139,7 +178,8 @@ func (pb *memoryProtoBook) GetProtocols(p peer.ID) ([]protocol.ID, error) {
 	return out, nil
 }
 
-func (pb *memoryProtoBook) RemoveProtocols(p peer.ID, protos ...protocol.ID) error {
+func (pb *memoryProtoBook) removeProtocolsFromSegment(p peer.ID, protos ...protocol.ID) error {
+
 	s := pb.segments.get(p)
 	s.Lock()
 	defer s.Unlock()
@@ -156,7 +196,30 @@ func (pb *memoryProtoBook) RemoveProtocols(p peer.ID, protos ...protocol.ID) err
 	return nil
 }
 
+func (pb *memoryProtoBook) removePeersFromProtocols(p peer.ID, protos ...protocol.ID) {
+	pb.protocolToPeersMap.Lock()
+	defer pb.protocolToPeersMap.Unlock()
+	for _, proto := range protos {
+		if peerMap, ok := pb.protocolToPeersMap.peers[proto]; ok {
+			delete(peerMap, p)
+		}
+	}
+}
+
+func (pb *memoryProtoBook) RemoveProtocols(p peer.ID, protos ...protocol.ID) error {
+	fmt.Println("RemoveProtocols for peer ", p, " protocols are :", protos)
+
+	err := pb.removeProtocolsFromSegment(p, protos...)
+	if err != nil {
+		return err
+	}
+	pb.removePeersFromProtocols(p, protos...)
+	return nil
+}
+
 func (pb *memoryProtoBook) SupportsProtocols(p peer.ID, protos ...protocol.ID) ([]protocol.ID, error) {
+	fmt.Println("SupportsProtocols for peer ", p, " queried protocols are :", protos)
+
 	s := pb.segments.get(p)
 	s.RLock()
 	defer s.RUnlock()
@@ -167,6 +230,7 @@ func (pb *memoryProtoBook) SupportsProtocols(p peer.ID, protos ...protocol.ID) (
 			out = append(out, proto)
 		}
 	}
+	fmt.Println("SupportsProtocols for peer ", p, " supported protocols are :", out)
 
 	return out, nil
 }
@@ -186,7 +250,40 @@ func (pb *memoryProtoBook) FirstSupportedProtocol(p peer.ID, protos ...protocol.
 
 func (pb *memoryProtoBook) RemovePeer(p peer.ID) {
 	s := pb.segments.get(p)
+	//TODO: Is a read lock required for the segment??
+	pb.protocolToPeersMap.Lock()
+	for _, protos := range s.protocols {
+		for proto := range protos {
+			if peers, ok := pb.protocolToPeersMap.peers[proto]; ok {
+				fmt.Println("Removing Peer for protocol ", proto, " peers before removal ", peers)
+				delete(peers, p)
+				fmt.Println("Removing Peer for protocol ", proto, " peers after removal ", peers)
+			}
+		}
+	}
+	pb.protocolToPeersMap.Unlock()
+
 	s.Lock()
 	delete(s.protocols, p)
 	s.Unlock()
+
+}
+
+func (pb *memoryProtoBook) GetPeersForProtocol(proto protocol.ID) ([]peer.ID, error) {
+	pb.protocolToPeersMap.RLock()
+	defer pb.protocolToPeersMap.RUnlock()
+
+	peers, ok := pb.protocolToPeersMap.peers[proto]
+	if !ok {
+		return nil, errNoPeersForProtocol
+	}
+	peerIDs := make([]peer.ID, len(peers))
+
+	i := 0
+	for k := range peers {
+		peerIDs[i] = k
+		i++
+	}
+
+	return peerIDs, nil
 }

p2p/host/peerstore/test/peerstore_suite.go

@@ -227,11 +227,21 @@ func testPeerstoreProtoStore(ps pstore.Peerstore) func(t *testing.T) {
 					t.Fatal("got wrong protocol")
 				}
 			}
+			peers, err := ps.GetPeersForProtocol(protocol.ID("a"))
+			t.Log("peers returned for protocol a is :", peers)
+			require.NoError(t, err)
+			require.Equal(t, 1, len(peers))
+			require.Equal(t, p1, peers[0])
+
+			peers, err = ps.GetPeersForProtocol(protocol.ID("non-existent"))
+			t.Log("peers returned for protocol non-existent is :", peers)
+
+			require.Error(t, err)
+			require.Equal(t, 0, len(peers))
 
 			supported, err := ps.SupportsProtocols(p1, "q", "w", "a", "y", "b")
 			require.NoError(t, err)
 			require.Len(t, supported, 2, "only expected 2 supported")
-
 			if supported[0] != "a" || supported[1] != "b" {
 				t.Fatal("got wrong supported array: ", supported)
 			}
@@ -281,10 +291,22 @@ func testPeerstoreProtoStore(ps pstore.Peerstore) func(t *testing.T) {
 			out, err := ps.GetProtocols(p)
 			require.NoError(t, err)
 			require.Len(t, out, 2)
+
+			peers, err := ps.GetPeersForProtocol(protocol.ID("a"))
+			t.Log("peers returned for protocol a is :", peers)
+			require.NoError(t, err)
+			require.Equal(t, 2, len(peers))
+
 			ps.RemovePeer(p)
 			out, err = ps.GetProtocols(p)
+
 			require.NoError(t, err)
 			require.Empty(t, out)
+
+			peers, err = ps.GetPeersForProtocol(protocol.ID("a"))
+			t.Log("peers returned for protocol a is :", peers)
+			require.NoError(t, err)
+			require.Equal(t, 1, len(peers))
 		})
 	}
 }