Merge pull request #2826 from ipfs/fix/dht-ctx-respect

respect contexts while reading messages in dht
ipfs · Jun 9, 2016 · b731c79 · b731c79
2 parents 6fd7ce3 + 73463cd
commit b731c79
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/routing/dht/dht_net.go b/routing/dht/dht_net.go
@@ -214,7 +214,7 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
 	log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
 
 	mes := new(pb.Message)
-	if err := ms.r.ReadMsg(mes); err != nil {
+	if err := ms.ctxReadMsg(ctx, mes); err != nil {
 		ms.s.Close()
 		ms.s = nil
 		return nil, err
@@ -227,3 +227,17 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
 
 	return mes, nil
 }
+
+func (ms *messageSender) ctxReadMsg(ctx context.Context, mes *pb.Message) error {
+	errc := make(chan error, 1)
+	go func() {
+		errc <- ms.r.ReadMsg(mes)
+	}()
+
+	select {
+	case err := <-errc:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}