Merge pull request #147 from libp2p/feat/supports-no-alloc

Implement the FirstSupportedProtocol API.
libp2p · May 14, 2020 · ce24bfc · ce24bfc
2 parents b1f58c0 + 21a1c14
commit ce24bfc
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 1 deletion.
diff --git a/go.mod b/go.mod
@@ -8,13 +8,14 @@ require (
 	github.com/ipfs/go-ds-leveldb v0.4.2
 	github.com/ipfs/go-log v1.0.3
 	github.com/libp2p/go-buffer-pool v0.0.2
-	github.com/libp2p/go-libp2p-core v0.5.1
+	github.com/libp2p/go-libp2p-core v0.5.4
 	github.com/multiformats/go-base32 v0.0.3
 	github.com/multiformats/go-multiaddr v0.2.1
 	github.com/multiformats/go-multiaddr-fmt v0.1.0
 	github.com/multiformats/go-multiaddr-net v0.1.4
 	github.com/multiformats/go-multihash v0.0.13
 	github.com/pkg/errors v0.9.1
+	github.com/stretchr/testify v1.4.0
 	github.com/whyrusleeping/go-keyspace v0.0.0-20160322163242-5b898ac5add1
 	go.uber.org/goleak v1.0.0
 	golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5 // indirect

diff --git a/go.sum b/go.sum
@@ -95,6 +95,12 @@ github.com/libp2p/go-buffer-pool v0.0.2/go.mod h1:MvaB6xw5vOrDl8rYZGLFdKAuk/hRoR
 github.com/libp2p/go-flow-metrics v0.0.3/go.mod h1:HeoSNUrOJVK1jEpDqVEiUOIXqhbnS27omG0uWU5slZs=
 github.com/libp2p/go-libp2p-core v0.5.1 h1:6Cu7WljPQtGY2krBlMoD8L/zH3tMUsCbqNFH7cZwCoI=
 github.com/libp2p/go-libp2p-core v0.5.1/go.mod h1:uN7L2D4EvPCvzSH5SrhR72UWbnSGpt5/a35Sm4upn4Y=
+github.com/libp2p/go-libp2p-core v0.5.4-0.20200508062439-98b95a487749 h1:G0zRpRnpZ8sAZ5E5IRB6x7np9iuulPIWE+y3I/xGNfo=
+github.com/libp2p/go-libp2p-core v0.5.4-0.20200508062439-98b95a487749/go.mod h1:uN7L2D4EvPCvzSH5SrhR72UWbnSGpt5/a35Sm4upn4Y=
+github.com/libp2p/go-libp2p-core v0.5.4-0.20200511073005-dabf5bbfb028 h1:nQb4SDWcadn/kpWtMwsWweogIzYviVE5qOMT42bkyNE=
+github.com/libp2p/go-libp2p-core v0.5.4-0.20200511073005-dabf5bbfb028/go.mod h1:uN7L2D4EvPCvzSH5SrhR72UWbnSGpt5/a35Sm4upn4Y=
+github.com/libp2p/go-libp2p-core v0.5.4 h1:Z8Tt3R5or2pkl3Wgywfcc0GCNjf18aYWA30OjBpbmRs=
+github.com/libp2p/go-libp2p-core v0.5.4/go.mod h1:uN7L2D4EvPCvzSH5SrhR72UWbnSGpt5/a35Sm4upn4Y=
 github.com/libp2p/go-msgio v0.0.4/go.mod h1:63lBBgOTDKQL6EWazRMCwXsEeEeK9O2Cd+0+6OOuipQ=
 github.com/libp2p/go-openssl v0.0.4 h1:d27YZvLoTyMhIN4njrkr8zMDOM4lfpHIp6A+TK9fovg=
 github.com/libp2p/go-openssl v0.0.4/go.mod h1:unDrJpgy3oFr+rqXsarWifmJuNnJR4chtO1HmaZjggc=

diff --git a/pstoreds/protobook.go b/pstoreds/protobook.go
@@ -122,6 +122,28 @@ func (pb *dsProtoBook) SupportsProtocols(p peer.ID, protos ...string) ([]string,
 	return res, nil
 }
 
+func (pb *dsProtoBook) FirstSupportedProtocol(p peer.ID, protos ...string) (string, error) {
+	if err := p.Validate(); err != nil {
+		return "", err
+	}
+
+	s := pb.segments.get(p)
+	s.RLock()
+	defer s.RUnlock()
+
+	pmap, err := pb.getProtocolMap(p)
+	if err != nil {
+		return "", err
+	}
+	for _, proto := range protos {
+		if _, ok := pmap[proto]; ok {
+			return proto, nil
+		}
+	}
+
+	return "", nil
+}
+
 func (pb *dsProtoBook) RemoveProtocols(p peer.ID, protos ...string) error {
 	if err := p.Validate(); err != nil {
 		return err

diff --git a/pstoremem/protobook.go b/pstoremem/protobook.go
@@ -163,3 +163,20 @@ func (pb *memoryProtoBook) SupportsProtocols(p peer.ID, protos ...string) ([]str
 
 	return out, nil
 }
+
+func (pb *memoryProtoBook) FirstSupportedProtocol(p peer.ID, protos ...string) (string, error) {
+	if err := p.Validate(); err != nil {
+		return "", err
+	}
+
+	s := pb.segments.get(p)
+	s.RLock()
+	defer s.RUnlock()
+
+	for _, proto := range protos {
+		if _, ok := s.protocols[p][proto]; ok {
+			return proto, nil
+		}
+	}
+	return "", nil
+}
diff --git a/test/peerstore_suite.go b/test/peerstore_suite.go
@@ -14,6 +14,7 @@ import (
 	ma "github.com/multiformats/go-multiaddr"
 
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
+	"github.com/stretchr/testify/require"
 )
 
 var peerstoreSuite = map[string]func(pstore.Peerstore) func(*testing.T){
@@ -242,6 +243,18 @@ func testPeerstoreProtoStore(ps pstore.Peerstore) func(t *testing.T) {
 			t.Fatal("got wrong supported array: ", supported)
 		}
 
+		b, err := ps.FirstSupportedProtocol(p1, "q", "w", "a", "y", "b")
+		require.NoError(t, err)
+		require.Equal(t, "a", b)
+
+		b, err = ps.FirstSupportedProtocol(p1, "q", "x", "z")
+		require.NoError(t, err)
+		require.Empty(t, b)
+
+		b, err = ps.FirstSupportedProtocol(p1, "a")
+		require.NoError(t, err)
+		require.Equal(t, "a", b)
+
 		protos = []string{"other", "yet another", "one more"}
 		err = ps.SetProtocols(p1, protos...)
 		if err != nil {