Merge pull request #656 from libp2p/fix/context-cancel-bug

fix: lookup context cancellation race condition

query.go

@@ -117,10 +117,12 @@ func (dht *IpfsDHT) runLookupWithFollowup(ctx context.Context, target string, qu
 	}
 
 	// wait for all queries to complete before returning, aborting ongoing queries if we've been externally stopped
+	followupsCompleted := 0
 processFollowUp:
 	for i := 0; i < len(queryPeers); i++ {
 		select {
 		case <-doneCh:
+			followupsCompleted++
 			if stopFn() {
 				cancelFollowUp()
 				if i < len(queryPeers)-1 {
@@ -130,10 +132,17 @@ processFollowUp:
 			}
 		case <-ctx.Done():
 			lookupRes.completed = false
+			cancelFollowUp()
 			break processFollowUp
 		}
 	}
 
+	if !lookupRes.completed {
+		for i := followupsCompleted; i < len(queryPeers); i++ {
+			<-doneCh
+		}
+	}
+
 	return lookupRes, nil
 }