reuse streams in the dht networking code

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>

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

@@ -1,7 +1,7 @@
 package dht
 
 import (
-	"errors"
+	"sync"
 	"time"
 
 	pb "github.com/ipfs/go-ipfs/routing/dht/pb"
@@ -27,40 +27,42 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
 	w := ggio.NewDelimitedWriter(cw)
 	mPeer := s.Conn().RemotePeer()
 
-	// receive msg
-	pmes := new(pb.Message)
-	if err := r.ReadMsg(pmes); err != nil {
-		log.Debugf("Error unmarshaling data: %s", err)
-		return
-	}
+	for {
+		// receive msg
+		pmes := new(pb.Message)
+		if err := r.ReadMsg(pmes); err != nil {
+			log.Debugf("Error unmarshaling data: %s", err)
+			return
+		}
 
-	// update the peer (on valid msgs only)
-	dht.updateFromMessage(ctx, mPeer, pmes)
+		// update the peer (on valid msgs only)
+		dht.updateFromMessage(ctx, mPeer, pmes)
 
-	// get handler for this msg type.
-	handler := dht.handlerForMsgType(pmes.GetType())
-	if handler == nil {
-		log.Debug("got back nil handler from handlerForMsgType")
-		return
-	}
+		// get handler for this msg type.
+		handler := dht.handlerForMsgType(pmes.GetType())
+		if handler == nil {
+			log.Debug("got back nil handler from handlerForMsgType")
+			return
+		}
 
-	// dispatch handler.
-	rpmes, err := handler(ctx, mPeer, pmes)
-	if err != nil {
-		log.Debugf("handle message error: %s", err)
-		return
-	}
+		// dispatch handler.
+		rpmes, err := handler(ctx, mPeer, pmes)
+		if err != nil {
+			log.Debugf("handle message error: %s", err)
+			return
+		}
 
-	// if nil response, return it before serializing
-	if rpmes == nil {
-		log.Debug("Got back nil response from request.")
-		return
-	}
+		// if nil response, return it before serializing
+		if rpmes == nil {
+			log.Debug("Got back nil response from request.")
+			continue
+		}
 
-	// send out response msg
-	if err := w.WriteMsg(rpmes); err != nil {
-		log.Debugf("send response error: %s", err)
-		return
+		// send out response msg
+		if err := w.WriteMsg(rpmes); err != nil {
+			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 {
+	rpmes, err := ms.SendRequest(ctx, pmes)
+	if err != nil {
 		return nil, err
 	}
-	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
-
-	rpmes := new(pb.Message)
-	if err := r.ReadMsg(rpmes); 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,17 +92,9 @@ 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)
-	if err != nil {
-		return err
-	}
-	defer s.Close()
+	ms := dht.messageSenderForPeer(p)
 
-	cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
-	w := ggio.NewDelimitedWriter(cw)
-
-	if err := w.WriteMsg(pmes); err != nil {
+	if err := ms.SendMessage(ctx, pmes); err != nil {
 		return err
 	}
 	log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
@@ -129,3 +105,89 @@ func (dht *IpfsDHT) updateFromMessage(ctx context.Context, p peer.ID, mes *pb.Me
 	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
+}
+
+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
+	}
+
+	ms.r = ggio.NewDelimitedReader(nstr, inet.MessageSizeMax)
+	ms.w = ggio.NewDelimitedWriter(nstr)
+	ms.s = nstr
+
+	return nil
+}
+
+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
+	}
+
+	err := ms.w.WriteMsg(pmes)
+	if err != nil {
+		ms.s.Close()
+		ms.s = nil
+		return err
+	}
+	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
+	}
+
+	err := ms.w.WriteMsg(pmes)
+	if err != nil {
+		ms.s.Close()
+		ms.s = nil
+		return nil, err
+	}
+
+	log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
+
+	mes := new(pb.Message)
+	err = ms.r.ReadMsg(mes)
+	if err != nil {
+		ms.s.Close()
+		ms.s = nil
+		return nil, err
+	}
+
+	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 {