extract messaging components from IpfsDHT into its own struct. create…

… a new struct that manages sending DHT messages that can be used independently from the DHT.

dht.go

@@ -1,7 +1,6 @@
 package dht
 
 import (
- "bytes"
  "context"
  "errors"
  "fmt"
@@ -33,7 +32,6 @@ import (
  goprocessctx "github.com/jbenet/goprocess/context"
  "github.com/multiformats/go-base32"
  ma "github.com/multiformats/go-multiaddr"
- "github.com/multiformats/go-multihash"
  "go.opencensus.io/tag"
  "go.uber.org/zap"
 )
@@ -97,8 +95,7 @@ type IpfsDHT struct {
  ctx context.Context
  proc goprocess.Process
 
- strmap map[peer.ID]*messageSender
- smlk sync.Mutex
+ protoMessenger *ProtocolMessenger
 
  plk sync.Mutex
 
@@ -188,6 +185,7 @@ func New(ctx context.Context, h host.Host, options ...Option) (*IpfsDHT, error)
  dht.enableValues = cfg.enableValues
 
  dht.Validator = cfg.validator
+ dht.protoMessenger = NewProtocolMessenger(dht.host, dht.protocols, dht.Validator)
 
  dht.testAddressUpdateProcessing = cfg.testAddressUpdateProcessing
 
@@ -274,7 +272,6 @@ func makeDHT(ctx context.Context, h host.Host, cfg config) (*IpfsDHT, error) {
  selfKey: kb.ConvertPeerID(h.ID()),
  peerstore: h.Peerstore(),
  host: h,
- strmap: make(map[peer.ID]*messageSender),
  birth: time.Now(),
  protocols: protocols,
  protocolsStrs: protocol.ConvertToStrings(protocols),
@@ -507,67 +504,8 @@ func (dht *IpfsDHT) persistRTPeersInPeerStore() {
  }
 }
 
-// putValueToPeer stores the given key/value pair at the peer 'p'
-func (dht *IpfsDHT) putValueToPeer(ctx context.Context, p peer.ID, rec *recpb.Record) error {
- pmes := pb.NewMessage(pb.Message_PUT_VALUE, rec.Key, 0)
- pmes.Record = rec
- rpmes, err := dht.sendRequest(ctx, p, pmes)
- if err != nil {
- logger.Debugw("failed to put value to peer", "to", p, "key", loggableRecordKeyBytes(rec.Key), "error", err)
- return err
- }
-
- if !bytes.Equal(rpmes.GetRecord().Value, pmes.GetRecord().Value) {
- logger.Infow("value not put correctly", "put-message", pmes, "get-message", rpmes)
- return errors.New("value not put correctly")
- }
-
- return nil
-}
-
 var errInvalidRecord = errors.New("received invalid record")
 
-// getValueOrPeers queries a particular peer p for the value for
-// key. It returns either the value or a list of closer peers.
-// NOTE: It will update the dht's peerstore with any new addresses
-// it finds for the given peer.
-func (dht *IpfsDHT) getValueOrPeers(ctx context.Context, p peer.ID, key string) (*recpb.Record, []*peer.AddrInfo, error) {
- pmes, err := dht.getValueSingle(ctx, p, key)
- if err != nil {
- return nil, nil, err
- }
-
- // Perhaps we were given closer peers
- peers := pb.PBPeersToPeerInfos(pmes.GetCloserPeers())
-
- if rec := pmes.GetRecord(); rec != nil {
- // Success! We were given the value
- logger.Debug("got value")
-
- // make sure record is valid.
- err = dht.Validator.Validate(string(rec.GetKey()), rec.GetValue())
- if err != nil {
- logger.Debug("received invalid record (discarded)")
- // return a sentinal to signify an invalid record was received
- err = errInvalidRecord
- rec = new(recpb.Record)
- }
- return rec, peers, err
- }
-
- if len(peers) > 0 {
- return nil, peers, nil
- }
-
- return nil, nil, routing.ErrNotFound
-}
-
-// getValueSingle simply performs the get value RPC with the given parameters
-func (dht *IpfsDHT) getValueSingle(ctx context.Context, p peer.ID, key string) (*pb.Message, error) {
- pmes := pb.NewMessage(pb.Message_GET_VALUE, []byte(key), 0)
- return dht.sendRequest(ctx, p, pmes)
-}
-
 // getLocal attempts to retrieve the value from the datastore
 func (dht *IpfsDHT) getLocal(key string) (*recpb.Record, error) {
  logger.Debugw("finding value in datastore", "key", loggableRecordKeyString(key))
@@ -696,17 +634,6 @@ func (dht *IpfsDHT) FindLocal(id peer.ID) peer.AddrInfo {
  }
 }
 
-// findPeerSingle asks peer 'p' if they know where the peer with id 'id' is
-func (dht *IpfsDHT) findPeerSingle(ctx context.Context, p peer.ID, id peer.ID) (*pb.Message, error) {
- pmes := pb.NewMessage(pb.Message_FIND_NODE, []byte(id), 0)
- return dht.sendRequest(ctx, p, pmes)
-}
-
-func (dht *IpfsDHT) findProvidersSingle(ctx context.Context, p peer.ID, key multihash.Multihash) (*pb.Message, error) {
- pmes := pb.NewMessage(pb.Message_GET_PROVIDERS, key, 0)
- return dht.sendRequest(ctx, p, pmes)
-}
-
 // nearestPeersToQuery returns the routing tables closest peers.
 func (dht *IpfsDHT) nearestPeersToQuery(pmes *pb.Message, count int) []peer.ID {
  closer := dht.routingTable.NearestPeers(kb.ConvertKey(string(pmes.GetKey())), count)
@@ -847,15 +774,7 @@ func (dht *IpfsDHT) Host() host.Host {
 
 // Ping sends a ping message to the passed peer and waits for a response.
 func (dht *IpfsDHT) Ping(ctx context.Context, p peer.ID) error {
- req := pb.NewMessage(pb.Message_PING, nil, 0)
- resp, err := dht.sendRequest(ctx, p, req)
- if err != nil {
- return fmt.Errorf("sending request: %w", err)
- }
- if resp.Type != pb.Message_PING {
- return fmt.Errorf("got unexpected response type: %v", resp.Type)
- }
- return nil
+ return dht.protoMessenger.Ping(ctx, p)
 }
 
 // newContextWithLocalTags returns a new context.Context with the InstanceID and

dht_net.go

@@ -9,8 +9,10 @@ import (
  "time"
 
  "github.com/libp2p/go-libp2p-core/helpers"
+ "github.com/libp2p/go-libp2p-core/host"
  "github.com/libp2p/go-libp2p-core/network"
  "github.com/libp2p/go-libp2p-core/peer"
+ "github.com/libp2p/go-libp2p-core/protocol"
  "github.com/libp2p/go-msgio/protoio"
 
  "github.com/libp2p/go-libp2p-kad-dht/metrics"
@@ -207,12 +209,38 @@ func (dht *IpfsDHT) handleNewMessage(s network.Stream) bool {
  }
 }
 
+type messageManager struct {
+ host host.Host // the network services we need
+ strmap map[peer.ID]*messageSender
+ smlk sync.Mutex
+ protocols []protocol.ID
+}
+
+func (m *messageManager) streamDisconnect(ctx context.Context, p peer.ID) {
+ m.smlk.Lock()
+ defer m.smlk.Unlock()
+ ms, ok := m.strmap[p]
+ if !ok {
+ return
+ }
+ delete(m.strmap, p)
+
+ // Do this asynchronously as ms.lk can block for a while.
+ go func() {
+ if err := ms.lk.Lock(ctx); err != nil {
+ return
+ }
+ defer ms.lk.Unlock()
+ ms.invalidate()
+ }()
+}
+
 // sendRequest sends out a request, but also makes sure to
 // measure the RTT for latency measurements.
-func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
+func (m *messageManager) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
  ctx, _ = tag.New(ctx, metrics.UpsertMessageType(pmes))
 
- ms, err := dht.messageSenderForPeer(ctx, p)
+ ms, err := m.messageSenderForPeer(ctx, p)
  if err != nil {
  stats.Record(ctx,
  metrics.SentRequests.M(1),
@@ -239,15 +267,15 @@ func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message
  metrics.SentBytes.M(int64(pmes.Size())),
  metrics.OutboundRequestLatency.M(float64(time.Since(start))/float64(time.Millisecond)),
  )
- dht.peerstore.RecordLatency(p, time.Since(start))
+ m.host.Peerstore().RecordLatency(p, time.Since(start))
  return rpmes, nil
 }
 
 // sendMessage sends out a message
-func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
+func (m *messageManager) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
  ctx, _ = tag.New(ctx, metrics.UpsertMessageType(pmes))
 
- ms, err := dht.messageSenderForPeer(ctx, p)
+ ms, err := m.messageSenderForPeer(ctx, p)
  if err != nil {
  stats.Record(ctx,
  metrics.SentMessages.M(1),
@@ -273,30 +301,30 @@ func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message
  return nil
 }
 
-func (dht *IpfsDHT) messageSenderForPeer(ctx context.Context, p peer.ID) (*messageSender, error) {
- dht.smlk.Lock()
- ms, ok := dht.strmap[p]
+func (m *messageManager) messageSenderForPeer(ctx context.Context, p peer.ID) (*messageSender, error) {
+ m.smlk.Lock()
+ ms, ok := m.strmap[p]
  if ok {
- dht.smlk.Unlock()
+ m.smlk.Unlock()
  return ms, nil
  }
- ms = &messageSender{p: p, dht: dht, lk: newCtxMutex()}
- dht.strmap[p] = ms
- dht.smlk.Unlock()
+ ms = &messageSender{p: p, m: m, lk: newCtxMutex()}
+ m.strmap[p] = ms
+ m.smlk.Unlock()
 
  if err := ms.prepOrInvalidate(ctx); err != nil {
- dht.smlk.Lock()
- defer dht.smlk.Unlock()
+ m.smlk.Lock()
+ defer m.smlk.Unlock()
 
- if msCur, ok := dht.strmap[p]; ok {
+ if msCur, ok := m.strmap[p]; ok {
  // Changed. Use the new one, old one is invalid and
  // not in the map so we can just throw it away.
  if ms != msCur {
  return msCur, nil
  }
  // Not changed, remove the now invalid stream from the
  // map.
- delete(dht.strmap, p)
+ delete(m.strmap, p)
  }
  // Invalid but not in map. Must have been removed by a disconnect.
  return nil, err
@@ -306,11 +334,11 @@ func (dht *IpfsDHT) messageSenderForPeer(ctx context.Context, p peer.ID) (*messa
 }
 
 type messageSender struct {
- s  network.Stream
- r  msgio.ReadCloser
- lk  ctxMutex
- p  peer.ID
- dht *IpfsDHT
+ s network.Stream
+ r msgio.ReadCloser
+ lk ctxMutex
+ p peer.ID
+ m *messageManager
 
  invalid bool
  singleMes int
@@ -351,7 +379,7 @@ func (ms *messageSender) prep(ctx context.Context) error {
  // We only want to speak to peers using our primary protocols. We do not want to query any peer that only speaks
  // one of the secondary "server" protocols that we happen to support (e.g. older nodes that we can respond to for
  // backwards compatibility reasons).
- nstr, err := ms.dht.host.NewStream(ctx, ms.p, ms.dht.protocols...)
+ nstr, err := ms.m.host.NewStream(ctx, ms.p, ms.m.protocols...)
  if err != nil {
  return err
  }

dht_test.go

@@ -570,14 +570,14 @@ func TestInvalidMessageSenderTracking(t *testing.T) {
  defer dht.Close()
 
  foo := peer.ID("asdasd")
- _, err := dht.messageSenderForPeer(ctx, foo)
+ _, err := dht.protoMessenger.m.messageSenderForPeer(ctx, foo)
  if err == nil {
  t.Fatal("that shouldnt have succeeded")
  }
 
- dht.smlk.Lock()
- mscnt := len(dht.strmap)
- dht.smlk.Unlock()
+ dht.protoMessenger.m.smlk.Lock()
+ mscnt := len(dht.protoMessenger.m.strmap)
+ dht.protoMessenger.m.smlk.Unlock()
 
  if mscnt > 0 {
  t.Fatal("should have no message senders in map")

go.sum

@@ -57,6 +57,7 @@ github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4er
 github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.0/go.mod h1:Qd/q+1AKNOZr9uGQzbzCmRO6sUih6GTPZv6a1/R87v0=
+github.com/golang/protobuf v1.3.1 h1:YF8+flBXS5eO826T4nzqPrxfhQThhXl0YzfuUPu4SBg=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8=
 github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA=
@@ -66,6 +67,7 @@ github.com/golang/protobuf v1.4.0 h1:oOuy+ugB+P/kBdUnG5QaMXSIyJ1q38wWSojYCb3z5VQ
 github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0=
 github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
+github.com/google/go-cmp v0.3.0 h1:crn/baboCvb5fXaQ0IJ1SGTsTVrWpDsCWC8EGETZijY=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4=
@@ -91,6 +93,7 @@ github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ
 github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
+github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/huin/goupnp v1.0.0 h1:wg75sLpL6DZqwHQN6E1Cfk6mtfzS45z8OV+ic+DtHRo=
 github.com/huin/goupnp v1.0.0/go.mod h1:n9v9KO1tAxYH82qOn+UTIFQDmx5n1Zxd/ClZDMX7Bnc=
@@ -467,6 +470,7 @@ github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
+github.com/onsi/ginkgo v1.12.0 h1:Iw5WCbBcaAAd0fpRb1c9r5YCylv4XDoCSigm1zLevwU=
 github.com/onsi/ginkgo v1.12.0/go.mod h1:oUhWkIvk5aDxtKvDDuw8gItl8pKl42LzjC9KZE0HfGg=
 github.com/onsi/ginkgo v1.12.1 h1:mFwc4LvZ0xpSvDZ3E+k8Yte0hLOMxXUlP+yXtJqkYfQ=
 github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
@@ -626,6 +630,7 @@ golang.org/x/tools v0.0.0-20191108193012-7d206e10da11/go.mod h1:b+2E5dAYhXwXZwtn
 golang.org/x/tools v0.0.0-20191216052735-49a3e744a425 h1:VvQyQJN0tSuecqgcIxMWnnfG5kSmgy9KZR9sW3W5QeA=
 golang.org/x/tools v0.0.0-20191216052735-49a3e744a425/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898 h1:/atklqdjdhuosWIl6AIbOeHJjicWYPqR9bpxqxYG2pA=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -650,6 +655,7 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
+gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/src-d/go-cli.v0 v0.0.0-20181105080154-d492247bbc0d/go.mod h1:z+K8VcOYVYcSwSjGebuDL6176A1XskgbtNl64NSg+n8=
 gopkg.in/src-d/go-log.v1 v1.0.1/go.mod h1:GN34hKP0g305ysm2/hctJ0Y8nWP3zxXXJ8GFabTyABE=

lookup.go

@@ -8,7 +8,6 @@ import (
  "github.com/libp2p/go-libp2p-core/peer"
  "github.com/libp2p/go-libp2p-core/routing"
 
- pb "github.com/libp2p/go-libp2p-kad-dht/pb"
  kb "github.com/libp2p/go-libp2p-kbucket"
 )
 
@@ -30,12 +29,11 @@ func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key string) (<-chan pee
  ID: p,
  })
 
- pmes, err := dht.findPeerSingle(ctx, p, peer.ID(key))
+ peers, err := dht.protoMessenger.GetClosestPeers(ctx, p, peer.ID(key))
  if err != nil {
  logger.Debugf("error getting closer peers: %s", err)
  return nil, err
  }
- peers := pb.PBPeersToPeerInfos(pmes.GetCloserPeers())
 
  // For DHT query command
  routing.PublishQueryEvent(ctx, &routing.QueryEvent{