stream back diagnostic responses as they are received
Jeromy committed
Jan 18, 2015 at 04:53 UTC
65b657e0ea82a1d3fee66a48529d46b56d7ab1e6
1 file changed
+106
-104
diagnostics/diag.go
+106
-104
@@ -4,11 +4,9 @@
4
package diagnostics
5
6
import (
7
- "bytes"
7
"encoding/json"
8
"errors"
9
"fmt"
11
- "io"
10
"sync"
11
"time"
12
@@ -33,6 +31,8 @@ var log = util.Logger("diagnostics")
31
// ProtocolDiag is the diagnostics protocol.ID
32
var ProtocolDiag protocol.ID = "/ipfs/diagnostics"
33
34
+var ErrAlreadyRunning = errors.New("diagnostic with that ID already running")
35
+
36
const ResponseTimeout = time.Second * 10
37
const HopTimeoutDecrement = time.Second * 2
38
@@ -159,86 +159,56 @@ func (d *Diagnostics) GetDiagnostic(timeout time.Duration) ([]*DiagInfo, error)
159
return nil, fmt.Errorf("diagnostic from peers err: %s", err)
160
}
161
162
- var out []*DiagInfo
162
di := d.getDiagInfo()
164
- out = append(out, di)
165
- for _, dpi := range dpeers {
166
- out = appendDiagnostics(out, dpi)
163
+ out := []*DiagInfo{di}
164
+ for dpi := range dpeers {
165
+ out = append(out, dpi)
166
}
167
return out, nil
168
}
169
171
-func appendDiagnostics(cur []*DiagInfo, data []byte) []*DiagInfo {
172
- buf := bytes.NewBuffer(data)
173
- dec := json.NewDecoder(buf)
174
- for {
175
- di := new(DiagInfo)
176
- err := dec.Decode(di)
177
- if err != nil {
178
- if err != io.EOF {
179
- log.Errorf("error decoding DiagInfo: %v", err)
180
- }
181
- break
182
- }
183
- cur = append(cur, di)
170
+func decodeDiagJson(data []byte) (*DiagInfo, error) {
171
+ di := new(DiagInfo)
172
+ err := json.Unmarshal(data, di)
173
+ if err != nil {
174
+ return nil, err
175
}
185
- return cur
186
-}
176
188
-func (d *Diagnostics) getDiagnosticFromPeers(ctx context.Context, peers map[peer.ID]int, pmes *pb.Message) ([][]byte, error) {
189
- timeout := pmes.GetTimeoutDuration()
190
- if timeout < 1 {
191
- return nil, fmt.Errorf("timeout too short: %s", timeout)
192
- }
193
- ctx, _ = context.WithTimeout(ctx, timeout)
177
+ return di, nil
178
+}
179
195
- respdata := make(chan []byte)
196
- sendcount := 0
180
+func (d *Diagnostics) getDiagnosticFromPeers(ctx context.Context, peers map[peer.ID]int, pmes *pb.Message) (<-chan *DiagInfo, error) {
181
+ respdata := make(chan *DiagInfo)
182
+ wg := sync.WaitGroup{}
183
for p, _ := range peers {
184
+ wg.Add(1)
185
log.Debugf("Sending diagnostic request to peer: %s", p)
199
- sendcount++
186
go func(p peer.ID) {
187
+ defer wg.Done()
188
out, err := d.getDiagnosticFromPeer(ctx, p, pmes)
189
if err != nil {
203
- log.Errorf("getDiagnostic error: %v", err)
204
- respdata <- nil
190
+ log.Errorf("Error getting diagnostic from %s: %s", p, err)
191
return
192
}
207
- respdata <- out
193
+ for d := range out {
194
+ respdata <- d
195
+ }
196
}(p)
197
}
198
211
- outall := make([][]byte, 0, len(peers))
212
- for i := 0; i < sendcount; i++ {
213
- out := <-respdata
214
- outall = append(outall, out)
215
- }
216
-
217
- return outall, nil
218
-}
219
-
220
-// TODO: this method no longer needed.
221
-func (d *Diagnostics) getDiagnosticFromPeer(ctx context.Context, p peer.ID, mes *pb.Message) ([]byte, error) {
222
- rpmes, err := d.sendRequest(ctx, p, mes)
223
- if err != nil {
224
- return nil, err
225
- }
226
- return rpmes.GetData(), nil
227
-}
199
+ go func() {
200
+ wg.Wait()
201
+ close(respdata)
202
+ }()
203
229
-func newMessage(diagID string) *pb.Message {
230
- pmes := new(pb.Message)
231
- pmes.DiagID = proto.String(diagID)
232
- return pmes
204
+ return respdata, nil
205
}
206
235
-func (d *Diagnostics) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
236
-
207
+func (d *Diagnostics) getDiagnosticFromPeer(ctx context.Context, p peer.ID, pmes *pb.Message) (<-chan *DiagInfo, error) {
208
s, err := d.host.NewStream(ProtocolDiag, p)
209
if err != nil {
210
return nil, err
211
}
241
- defer s.Close()
212
213
cr := ctxutil.NewReader(ctx, s) // ok to use. we defer close stream in this func
214
cw := ctxutil.NewWriter(ctx, s) // ok to use. we defer close stream in this func
@@ -251,51 +221,57 @@ func (d *Diagnostics) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Messa
221
return nil, err
222
}
223
254
- rpmes := new(pb.Message)
255
- if err := r.ReadMsg(rpmes); err != nil {
256
- return nil, err
257
- }
258
- if rpmes == nil {
259
- return nil, errors.New("no response to request")
260
- }
224
+ out := make(chan *DiagInfo)
225
+ go func() {
226
262
- rtt := time.Since(start)
263
- log.Infof("diagnostic request took: %s", rtt.String())
264
- return rpmes, nil
265
-}
227
+ defer func() {
228
+ close(out)
229
+ s.Close()
230
+ rtt := time.Since(start)
231
+ log.Infof("diagnostic request took: %s", rtt.String())
232
+ }()
233
267
-func (d *Diagnostics) handleDiagnostic(p peer.ID, pmes *pb.Message) (*pb.Message, error) {
268
- log.Debugf("HandleDiagnostic from %s for id = %s", p, util.Key(pmes.GetDiagID()).B58String())
269
- resp := newMessage(pmes.GetDiagID())
234
+ for {
235
+ rpmes := new(pb.Message)
236
+ if err := r.ReadMsg(rpmes); err != nil {
237
+ log.Errorf("Error reading diagnostic from stream: %s", err)
238
+ return
239
+ }
240
+ if rpmes == nil {
241
+ log.Error("Got no response back from diag request.")
242
+ return
243
+ }
244
271
- // Make sure we havent already handled this request to prevent loops
272
- d.diagLock.Lock()
273
- _, found := d.diagMap[pmes.GetDiagID()]
274
- if found {
275
- d.diagLock.Unlock()
276
- return resp, nil
277
- }
278
- d.diagMap[pmes.GetDiagID()] = time.Now()
279
- d.diagLock.Unlock()
245
+ di, err := decodeDiagJson(rpmes.GetData())
246
+ if err != nil {
247
+ log.Error(err)
248
+ return
249
+ }
250
281
- di := d.getDiagInfo()
282
- resp.Data = di.Marshal()
283
- dpeers, err := d.getDiagnosticFromPeers(context.TODO(), d.getPeers(), pmes)
284
- if err != nil {
285
- log.Errorf("diagnostic from peers err: %s", err)
286
- } else {
287
- for _, b := range dpeers {
288
- resp.Data = append(resp.Data, b...) // concatenate them all.
251
+ select {
252
+ case out <- di:
253
+ case <-ctx.Done():
254
+ return
255
+ }
256
}
290
- }
257
292
- return resp, nil
258
+ }()
259
+
260
+ return out, nil
261
+}
262
+
263
+func newMessage(diagID string) *pb.Message {
264
+ pmes := new(pb.Message)
265
+ pmes.DiagID = proto.String(diagID)
266
+ return pmes
267
}
268
269
func (d *Diagnostics) HandleMessage(ctx context.Context, s inet.Stream) error {
270
297
- r := ggio.NewDelimitedReader(s, 32768) // maxsize
298
- w := ggio.NewDelimitedWriter(s)
271
+ cr := ctxutil.NewReader(ctx, s)
272
+ cw := ctxutil.NewWriter(ctx, s)
273
+ r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax) // maxsize
274
+ w := ggio.NewDelimitedWriter(cw)
275
276
// deserialize msg
277
pmes := new(pb.Message)
@@ -308,25 +284,51 @@ func (d *Diagnostics) HandleMessage(ctx context.Context, s inet.Stream) error {
284
log.Infof("[peer: %s] Got message from [%s]\n",
285
d.self.Pretty(), s.Conn().RemotePeer())
286
311
- // dispatch handler.
312
- p := s.Conn().RemotePeer()
313
- rpmes, err := d.handleDiagnostic(p, pmes)
314
- if err != nil {
315
- log.Errorf("handleDiagnostic error: %s", err)
287
+ // Make sure we havent already handled this request to prevent loops
288
+ if err := d.startDiag(pmes.GetDiagID()); err != nil {
289
return nil
290
}
291
319
- // if nil response, return it before serializing
320
- if rpmes == nil {
321
- return nil
292
+ resp := newMessage(pmes.GetDiagID())
293
+ resp.Data = d.getDiagInfo().Marshal()
294
+ if err := w.WriteMsg(resp); err != nil {
295
+ log.Errorf("Failed to write protobuf message over stream: %s", err)
296
+ return err
297
}
298
324
- // serialize + send response msg
325
- if err := w.WriteMsg(rpmes); err != nil {
326
- log.Errorf("Failed to encode protobuf message: %v", err)
327
- return nil
299
+ timeout := pmes.GetTimeoutDuration()
300
+ if timeout < HopTimeoutDecrement {
301
+ return fmt.Errorf("timeout too short: %s", timeout)
302
+ }
303
+ ctx, _ = context.WithTimeout(ctx, timeout)
304
+ pmes.SetTimeoutDuration(timeout - HopTimeoutDecrement)
305
+
306
+ dpeers, err := d.getDiagnosticFromPeers(ctx, d.getPeers(), pmes)
307
+ if err != nil {
308
+ log.Errorf("diagnostic from peers err: %s", err)
309
+ return err
310
}
311
+ for b := range dpeers {
312
+ resp := newMessage(pmes.GetDiagID())
313
+ resp.Data = b.Marshal()
314
+ if err := w.WriteMsg(resp); err != nil {
315
+ log.Errorf("Failed to write protobuf message over stream: %s", err)
316
+ return err
317
+ }
318
+ }
319
+
320
+ return nil
321
+}
322
323
+func (d *Diagnostics) startDiag(id string) error {
324
+ d.diagLock.Lock()
325
+ _, found := d.diagMap[id]
326
+ if found {
327
+ d.diagLock.Unlock()
328
+ return ErrAlreadyRunning
329
+ }
330
+ d.diagMap[id] = time.Now()
331
+ d.diagLock.Unlock()
332
return nil
333
}
334