respect contexts while reading messages in dht
License: MIT Signed-off-by: Jeromy <why@ipfs.io>
Jeromy committed
Jun 8, 2016 at 16:12 UTC
73463cd208c92d180cd4c0977299d1e18eced54d
1 file changed
+15
-1
routing/dht/dht_net.go
+15
-1
@@ -214,7 +214,7 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
214
log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
215
216
mes := new(pb.Message)
217
- if err := ms.r.ReadMsg(mes); err != nil {
217
+ if err := ms.ctxReadMsg(ctx, mes); err != nil {
218
ms.s.Close()
219
ms.s = nil
220
return nil, err
@@ -227,3 +227,17 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
227
228
return mes, nil
229
}
230
+
231
+func (ms *messageSender) ctxReadMsg(ctx context.Context, mes *pb.Message) error {
232
+ errc := make(chan error, 1)
233
+ go func() {
234
+ errc <- ms.r.ReadMsg(mes)
235
+ }()
236
+
237
+ select {
238
+ case err := <-errc:
239
+ return err
240
+ case <-ctx.Done():
241
+ return ctx.Err()
242
+ }
243
+}