use query for getClosestPeers

routing/dht/routing.go

@@ -159,52 +159,48 @@ func (dht *IpfsDHT) getClosestPeers(ctx context.Context, key u.Key, count int) (
 	peerset := pset.NewLimited(count)
 
 	for _, p := range tablepeers {
-		out <- p
+		select {
+		case out <- p:
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
 		peerset.Add(p)
 	}
 
-	wg := sync.WaitGroup{}
+	query := newQuery(key, dht.network, func(ctx context.Context, p peer.ID) (*dhtQueryResult, error) {
-	for _, p := range tablepeers {
+		closer, err := dht.closerPeersSingle(ctx, key, p)
-		wg.Add(1)
+		if err != nil {
-		go func(p peer.ID) {
+			log.Errorf("error getting closer peers: %s", err)
-			dht.getClosestPeersRecurse(ctx, key, p, peerset, out)
+			return nil, err
-			wg.Done()
+		}
-		}(p)
+
-	}
+		var filtered []peer.PeerInfo
+		for _, p := range closer {
+			if kb.Closer(p, dht.self, key) && peerset.TryAdd(p) {
+				select {
+				case out <- p:
+				case <-ctx.Done():
+					return nil, ctx.Err()
+				}
+				filtered = append(filtered, dht.peerstore.PeerInfo(p))
+			}
+		}
+
+		return &dhtQueryResult{closerPeers: filtered}, nil
+	})
 
 	go func() {
-		wg.Wait()
+		defer close(out)
-		close(out)
+		// run it!
+		_, err := query.Run(ctx, tablepeers)
+		if err != nil {
+			log.Errorf("closestPeers query run error: %s", err)
+		}
 	}()
 
 	return out, nil
 }
 
-func (dht *IpfsDHT) getClosestPeersRecurse(ctx context.Context, key u.Key, p peer.ID, peers *pset.PeerSet, peerOut chan<- peer.ID) {
-	closer, err := dht.closerPeersSingle(ctx, key, p)
-	if err != nil {
-		log.Errorf("error getting closer peers: %s", err)
-		return
-	}
-
-	wg := sync.WaitGroup{}
-	for _, p := range closer {
-		if kb.Closer(p, dht.self, key) && peers.TryAdd(p) {
-			select {
-			case peerOut <- p:
-			case <-ctx.Done():
-				return
-			}
-			wg.Add(1)
-			go func(p peer.ID) {
-				dht.getClosestPeersRecurse(ctx, key, p, peers, peerOut)
-				wg.Done()
-			}(p)
-		}
-	}
-	wg.Wait()
-}
-
 func (dht *IpfsDHT) closerPeersSingle(ctx context.Context, key u.Key, p peer.ID) ([]peer.ID, error) {
 	pmes, err := dht.findPeerSingle(ctx, p, peer.ID(key))
 	if err != nil {