cmd/dht-get-peers: Add -once flag to do a single scrape

diff --git a/cmd/dht-get-peers/main.go b/cmd/dht-get-peers/main.go
index a61a26602ae92d0d1e6aec86c3f547498f4a4124..fa6b5dc8508539e9e80e69a7d2e9275f63ac0f33 100644 (file)
--- a/cmd/dht-get-peers/main.go
+++ b/cmd/dht-get-peers/main.go
@@ -8,7 +8,6 @@ import (
        "net"
        "os"
        "os/signal"
-       "time"
 
        "bitbucket.org/anacrolix/go.torrent/dht"
        "bitbucket.org/anacrolix/go.torrent/util"
@@ -24,6 +23,7 @@ var (
        tableFileName = flag.String("tableFile", "", "name of file for storing node info")
        serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
        infoHash      = flag.String("infoHash", "", "torrent infohash")
+       once          = flag.Bool("once", false, "only do one scrape iteration")
 
        s        *dht.Server
        quitting = make(chan struct{})
@@ -138,8 +138,13 @@ getPeers:
                if err != nil {
                        log.Fatal(err)
                }
-               go func() {
-                       for v := range ps.Values {
+       values:
+               for {
+                       select {
+                       case v, ok := <-ps.Values:
+                               if !ok {
+                                       break values
+                               }
                                log.Printf("received %d peers from %x", len(v.Peers), v.NodeInfo.ID)
                                for _, p := range v.Peers {
                                        if _, ok := seen[p]; ok {
@@ -151,12 +156,12 @@ getPeers:
                                                Port: int(p.Port),
                                        }).String())
                                }
+                       case <-quitting:
+                               break getPeers
                        }
-               }()
-               select {
-               case <-time.After(15 * time.Second):
-               case <-quitting:
-                       break getPeers
+               }
+               if *once {
+                       break
                }
        }
        if err := saveTable(); err != nil {