diff --git a/dht.go b/dht.go index 6ddd432..1ef8245 100644 --- a/dht.go +++ b/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) diff --git a/dht_net.go b/dht_net.go index 4ebd6e3..abafb52 100644 --- a/dht_net.go +++ b/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 +} diff --git a/ext_test.go b/ext_test.go index 28868a9..b5fa640 100644 --- a/ext_test.go +++ b/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 {