Merge pull request #2817 from ipfs/feat/dht-reuse-stream

reuse streams for dht messaging

dht.go

@@ -58,6 +58,9 @@ type IpfsDHT struct {
 
 	ctx  context.Context
 	proc goprocess.Process
+
+	strmap map[peer.ID]*messageSender
+	smlk   sync.Mutex
 }
 
 // NewDHT creates a new DHT object with the given peer as the 'local' host
@@ -77,6 +80,7 @@ func NewDHT(ctx context.Context, h host.Host, dstore ds.Datastore) *IpfsDHT {
 		return nil
 	})
 
+	dht.strmap = make(map[peer.ID]*messageSender)
 	dht.ctx = ctx
 
 	h.SetStreamHandler(ProtocolDHT, dht.handleNewStream)

dht_net.go

 package dht
 
 import (
-	"errors"
+	"sync"
 	"time"
 
 	pb "github.com/ipfs/go-ipfs/routing/dht/pb"
@@ -27,6 +27,7 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
 	w := ggio.NewDelimitedWriter(cw)
 	mPeer := s.Conn().RemotePeer()
 
+	for {
 		// receive msg
 		pmes := new(pb.Message)
 		if err := r.ReadMsg(pmes); err != nil {
@@ -54,7 +55,7 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
 		// if nil response, return it before serializing
 		if rpmes == nil {
 			log.Debug("Got back nil response from request.")
-		return
+			continue
 		}
 
 		// send out response msg
@@ -62,6 +63,7 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
 			log.Debugf("send response error: %s", err)
 			return
 		}
+	}
 
 	return
 }
@@ -70,32 +72,14 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
 // measure the RTT for latency measurements.
 func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
 
-	log.Debugf("%s DHT starting stream", dht.self)
-	s, err := dht.host.NewStream(ctx, ProtocolDHT, p)
-	if err != nil {
-		return nil, err
-	}
-	defer s.Close()
-
-	cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func
-	cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
-	r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax)
-	w := ggio.NewDelimitedWriter(cw)
+	ms := dht.messageSenderForPeer(p)
 
 	start := time.Now()
 
-	if err := w.WriteMsg(pmes); err != nil {
-		return nil, err
-	}
-	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
-
-	rpmes := new(pb.Message)
-	if err := r.ReadMsg(rpmes); err != nil {
+	rpmes, err := ms.SendRequest(ctx, pmes)
+	if err != nil {
 		return nil, err
 	}
-	if rpmes == nil {
-		return nil, errors.New("no response to request")
-	}
 
 	// update the peer (on valid msgs only)
 	dht.updateFromMessage(ctx, p, rpmes)
@@ -108,24 +92,138 @@ func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message
 // sendMessage sends out a message
 func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
 
-	log.Debugf("%s DHT starting stream", dht.self)
-	s, err := dht.host.NewStream(ctx, ProtocolDHT, p)
+	ms := dht.messageSenderForPeer(p)
+
+	if err := ms.SendMessage(ctx, pmes); err != nil {
+		return err
+	}
+	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
+	return nil
+}
+
+func (dht *IpfsDHT) updateFromMessage(ctx context.Context, p peer.ID, mes *pb.Message) error {
+	dht.Update(ctx, p)
+	return nil
+}
+
+func (dht *IpfsDHT) messageSenderForPeer(p peer.ID) *messageSender {
+	dht.smlk.Lock()
+	defer dht.smlk.Unlock()
+
+	ms, ok := dht.strmap[p]
+	if !ok {
+		ms = dht.newMessageSender(p)
+		dht.strmap[p] = ms
+	}
+
+	return ms
+}
+
+type messageSender struct {
+	s   inet.Stream
+	r   ggio.ReadCloser
+	w   ggio.WriteCloser
+	lk  sync.Mutex
+	p   peer.ID
+	dht *IpfsDHT
+
+	singleMes int
+}
+
+func (dht *IpfsDHT) newMessageSender(p peer.ID) *messageSender {
+	return &messageSender{p: p, dht: dht}
+}
+
+func (ms *messageSender) prep() error {
+	if ms.s != nil {
+		return nil
+	}
+
+	nstr, err := ms.dht.host.NewStream(ms.dht.ctx, ProtocolDHT, ms.p)
 	if err != nil {
 		return err
 	}
-	defer s.Close()
 
-	cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
-	w := ggio.NewDelimitedWriter(cw)
+	ms.r = ggio.NewDelimitedReader(nstr, inet.MessageSizeMax)
+	ms.w = ggio.NewDelimitedWriter(nstr)
+	ms.s = nstr
 
-	if err := w.WriteMsg(pmes); err != nil {
+	return nil
+}
+
+// streamReuseTries is the number of times we will try to reuse a stream to a
+// given peer before giving up and reverting to the old one-message-per-stream
+// behaviour.
+const streamReuseTries = 3
+
+func (ms *messageSender) SendMessage(ctx context.Context, pmes *pb.Message) error {
+	ms.lk.Lock()
+	defer ms.lk.Unlock()
+	if err := ms.prep(); err != nil {
 		return err
 	}
-	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
+
+	if err := ms.writeMessage(pmes); err != nil {
+		return err
+	}
+
+	if ms.singleMes > streamReuseTries {
+		ms.s.Close()
+		ms.s = nil
+	}
+
 	return nil
 }
 
-func (dht *IpfsDHT) updateFromMessage(ctx context.Context, p peer.ID, mes *pb.Message) error {
-	dht.Update(ctx, p)
+func (ms *messageSender) writeMessage(pmes *pb.Message) error {
+	err := ms.w.WriteMsg(pmes)
+	if err != nil {
+		// If the other side isnt expecting us to be reusing streams, we're gonna
+		// end up erroring here. To make sure things work seamlessly, lets retry once
+		// before continuing
+
+		log.Infof("error writing message: ", err)
+		ms.s.Close()
+		ms.s = nil
+		if err := ms.prep(); err != nil {
+			return err
+		}
+
+		if err := ms.w.WriteMsg(pmes); err != nil {
+			return err
+		}
+
+		// keep track of this happening. If it happens a few times, its
+		// likely we can assume the otherside will never support stream reuse
+		ms.singleMes++
+	}
 	return nil
 }
+
+func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb.Message, error) {
+	ms.lk.Lock()
+	defer ms.lk.Unlock()
+	if err := ms.prep(); err != nil {
+		return nil, err
+	}
+
+	if err := ms.writeMessage(pmes); err != nil {
+		return nil, err
+	}
+
+	log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
+
+	mes := new(pb.Message)
+	if err := ms.r.ReadMsg(mes); err != nil {
+		ms.s.Close()
+		ms.s = nil
+		return nil, err
+	}
+
+	if ms.singleMes > streamReuseTries {
+		ms.s.Close()
+		ms.s = nil
+	}
+
+	return mes, nil
+}

ext_test.go

@@ -2,7 +2,6 @@ package dht
 
 import (
 	"io"
-	"io/ioutil"
 	"math/rand"
 	"testing"
 	"time"
@@ -40,8 +39,7 @@ func TestGetFailures(t *testing.T) {
 
 	// Reply with failures to every message
 	hosts[1].SetStreamHandler(ProtocolDHT, func(s inet.Stream) {
-		defer s.Close()
-		io.Copy(ioutil.Discard, s)
+		s.Close()
 	})
 
 	// This one should time out
@@ -51,7 +49,7 @@ func TestGetFailures(t *testing.T) {
 			err = merr[0]
 		}
 
-		if err.Error() != "process closing" {
+		if err != io.EOF {
 			t.Fatal("Got different error than we expected", err)
 		}
 	} else {