reuse streams in the dht networking code
License: MIT Signed-off-by: Jeromy <why@ipfs.io>
Jeromy committed
Jun 6, 2016 at 17:35 UTC
3e9717aa20ae0dd994d253573e6432634541f9ff
3 files changed
+134
-70
routing/dht/dht.go
+4
@@ -58,6 +58,9 @@ type IpfsDHT struct {
58
59
ctx context.Context
60
proc goprocess.Process
61
+
62
+ strmap map[peer.ID]*messageSender
63
+ smlk sync.Mutex
64
}
65
66
// 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 {
80
return nil
81
})
82
83
+ dht.strmap = make(map[peer.ID]*messageSender)
84
dht.ctx = ctx
85
86
h.SetStreamHandler(ProtocolDHT, dht.handleNewStream)
routing/dht/dht_net.go
+128
-66
@@ -1,7 +1,7 @@
1
package dht
2
3
import (
4
- "errors"
4
+ "sync"
5
"time"
6
7
pb "github.com/ipfs/go-ipfs/routing/dht/pb"
@@ -27,40 +27,42 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
27
w := ggio.NewDelimitedWriter(cw)
28
mPeer := s.Conn().RemotePeer()
29
30
- // receive msg
31
- pmes := new(pb.Message)
32
- if err := r.ReadMsg(pmes); err != nil {
33
- log.Debugf("Error unmarshaling data: %s", err)
34
- return
35
- }
36
-
37
- // update the peer (on valid msgs only)
38
- dht.updateFromMessage(ctx, mPeer, pmes)
39
-
40
- // get handler for this msg type.
41
- handler := dht.handlerForMsgType(pmes.GetType())
42
- if handler == nil {
43
- log.Debug("got back nil handler from handlerForMsgType")
44
- return
45
- }
46
-
47
- // dispatch handler.
48
- rpmes, err := handler(ctx, mPeer, pmes)
49
- if err != nil {
50
- log.Debugf("handle message error: %s", err)
51
- return
52
- }
53
-
54
- // if nil response, return it before serializing
55
- if rpmes == nil {
56
- log.Debug("Got back nil response from request.")
57
- return
58
- }
59
-
60
- // send out response msg
61
- if err := w.WriteMsg(rpmes); err != nil {
62
- log.Debugf("send response error: %s", err)
63
- return
30
+ for {
31
+ // receive msg
32
+ pmes := new(pb.Message)
33
+ if err := r.ReadMsg(pmes); err != nil {
34
+ log.Debugf("Error unmarshaling data: %s", err)
35
+ return
36
+ }
37
+
38
+ // update the peer (on valid msgs only)
39
+ dht.updateFromMessage(ctx, mPeer, pmes)
40
+
41
+ // get handler for this msg type.
42
+ handler := dht.handlerForMsgType(pmes.GetType())
43
+ if handler == nil {
44
+ log.Debug("got back nil handler from handlerForMsgType")
45
+ return
46
+ }
47
+
48
+ // dispatch handler.
49
+ rpmes, err := handler(ctx, mPeer, pmes)
50
+ if err != nil {
51
+ log.Debugf("handle message error: %s", err)
52
+ return
53
+ }
54
+
55
+ // if nil response, return it before serializing
56
+ if rpmes == nil {
57
+ log.Debug("Got back nil response from request.")
58
+ continue
59
+ }
60
+
61
+ // send out response msg
62
+ if err := w.WriteMsg(rpmes); err != nil {
63
+ log.Debugf("send response error: %s", err)
64
+ return
65
+ }
66
}
67
68
return
@@ -70,32 +72,14 @@ func (dht *IpfsDHT) handleNewMessage(s inet.Stream) {
72
// measure the RTT for latency measurements.
73
func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
74
73
- log.Debugf("%s DHT starting stream", dht.self)
74
- s, err := dht.host.NewStream(ctx, ProtocolDHT, p)
75
- if err != nil {
76
- return nil, err
77
- }
78
- defer s.Close()
79
-
80
- cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func
81
- cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
82
- r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax)
83
- w := ggio.NewDelimitedWriter(cw)
75
+ ms := dht.messageSenderForPeer(p)
76
77
start := time.Now()
78
87
- if err := w.WriteMsg(pmes); err != nil {
88
- return nil, err
89
- }
90
- log.Event(ctx, "dhtSentMessage", dht.self, p, pmes)
91
-
92
- rpmes := new(pb.Message)
93
- if err := r.ReadMsg(rpmes); err != nil {
79
+ rpmes, err := ms.SendRequest(ctx, pmes)
80
+ if err != nil {
81
return nil, err
82
}
96
- if rpmes == nil {
97
- return nil, errors.New("no response to request")
98
- }
83
84
// update the peer (on valid msgs only)
85
dht.updateFromMessage(ctx, p, rpmes)
@@ -108,17 +92,9 @@ func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message
92
// sendMessage sends out a message
93
func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message) error {
94
111
- log.Debugf("%s DHT starting stream", dht.self)
112
- s, err := dht.host.NewStream(ctx, ProtocolDHT, p)
113
- if err != nil {
114
- return err
115
- }
116
- defer s.Close()
95
+ ms := dht.messageSenderForPeer(p)
96
118
- cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func
119
- w := ggio.NewDelimitedWriter(cw)
120
-
121
- if err := w.WriteMsg(pmes); err != nil {
97
+ if err := ms.SendMessage(ctx, pmes); err != nil {
98
return err
99
}
100
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
105
dht.Update(ctx, p)
106
return nil
107
}
108
+
109
+func (dht *IpfsDHT) messageSenderForPeer(p peer.ID) *messageSender {
110
+ dht.smlk.Lock()
111
+ defer dht.smlk.Unlock()
112
+
113
+ ms, ok := dht.strmap[p]
114
+ if !ok {
115
+ ms = dht.newMessageSender(p)
116
+ dht.strmap[p] = ms
117
+ }
118
+
119
+ return ms
120
+}
121
+
122
+type messageSender struct {
123
+ s inet.Stream
124
+ r ggio.ReadCloser
125
+ w ggio.WriteCloser
126
+ lk sync.Mutex
127
+ p peer.ID
128
+ dht *IpfsDHT
129
+}
130
+
131
+func (dht *IpfsDHT) newMessageSender(p peer.ID) *messageSender {
132
+ return &messageSender{p: p, dht: dht}
133
+}
134
+
135
+func (ms *messageSender) prep() error {
136
+ if ms.s != nil {
137
+ return nil
138
+ }
139
+
140
+ nstr, err := ms.dht.host.NewStream(ms.dht.ctx, ProtocolDHT, ms.p)
141
+ if err != nil {
142
+ return err
143
+ }
144
+
145
+ ms.r = ggio.NewDelimitedReader(nstr, inet.MessageSizeMax)
146
+ ms.w = ggio.NewDelimitedWriter(nstr)
147
+ ms.s = nstr
148
+
149
+ return nil
150
+}
151
+
152
+func (ms *messageSender) SendMessage(ctx context.Context, pmes *pb.Message) error {
153
+ ms.lk.Lock()
154
+ defer ms.lk.Unlock()
155
+ if err := ms.prep(); err != nil {
156
+ return err
157
+ }
158
+
159
+ err := ms.w.WriteMsg(pmes)
160
+ if err != nil {
161
+ ms.s.Close()
162
+ ms.s = nil
163
+ return err
164
+ }
165
+ return nil
166
+}
167
+
168
+func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb.Message, error) {
169
+ ms.lk.Lock()
170
+ defer ms.lk.Unlock()
171
+ if err := ms.prep(); err != nil {
172
+ return nil, err
173
+ }
174
+
175
+ err := ms.w.WriteMsg(pmes)
176
+ if err != nil {
177
+ ms.s.Close()
178
+ ms.s = nil
179
+ return nil, err
180
+ }
181
+
182
+ log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
183
+
184
+ mes := new(pb.Message)
185
+ err = ms.r.ReadMsg(mes)
186
+ if err != nil {
187
+ ms.s.Close()
188
+ ms.s = nil
189
+ return nil, err
190
+ }
191
+
192
+ return mes, nil
193
+}
routing/dht/ext_test.go
+2
-4
@@ -2,7 +2,6 @@ package dht
2
3
import (
4
"io"
5
- "io/ioutil"
5
"math/rand"
6
"testing"
7
"time"
@@ -40,8 +39,7 @@ func TestGetFailures(t *testing.T) {
39
40
// Reply with failures to every message
41
hosts[1].SetStreamHandler(ProtocolDHT, func(s inet.Stream) {
43
- defer s.Close()
44
- io.Copy(ioutil.Discard, s)
42
+ s.Close()
43
})
44
45
// This one should time out
@@ -51,7 +49,7 @@ func TestGetFailures(t *testing.T) {
49
err = merr[0]
50
}
51
54
- if err.Error() != "process closing" {
52
+ if err != io.EOF {
53
t.Fatal("Got different error than we expected", err)
54
}
55
} else {