Merge pull request #900 from jbenet/fix/bitswap-ctx-respect

respect contexts in a more timely manner

exchange/bitswap/bitswap.go

@@ -227,7 +227,14 @@ func (bs *Bitswap) HasBlock(ctx context.Context, blk *blocks.Block) error {
 func (bs *Bitswap) sendWantlistMsgToPeers(ctx context.Context, m bsmsg.BitSwapMessage, peers <-chan peer.ID) error {
 	set := pset.New()
 	wg := sync.WaitGroup{}
-	for peerToQuery := range peers {
+
+loop:
+	for {
+		select {
+		case peerToQuery, ok := <-peers:
+			if !ok {
+				break loop
+			}
 
 			if !set.TryAdd(peerToQuery) { //Do once per peer
 				continue
@@ -240,8 +247,23 @@ func (bs *Bitswap) sendWantlistMsgToPeers(ctx context.Context, m bsmsg.BitSwapMe
 					log.Debug(err) // TODO remove if too verbose
 				}
 			}(peerToQuery)
+		case <-ctx.Done():
+			return nil
+		}
 	}
+	done := make(chan struct{})
+	go func() {
 		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		// NB: we may be abandoning goroutines here before they complete
+		// this shouldnt be an issue because they will complete soon anyways
+		// we just don't want their being slow to impact bitswap transfer speeds
+	}
 	return nil
 }
 
@@ -385,7 +407,18 @@ func (bs *Bitswap) wantNewBlocks(ctx context.Context, bkeys []u.Key) {
 			}
 		}(p)
 	}
+	done := make(chan struct{})
+	go func() {
 		wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-ctx.Done():
+		// NB: we may be abandoning goroutines here before they complete
+		// this shouldnt be an issue because they will complete soon anyways
+		// we just don't want their being slow to impact bitswap transfer speeds
+	}
 }
 
 func (bs *Bitswap) ReceiveError(err error) {