From 5aca9d35631ab143ccadcbd821309df9d8f01131 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 8 Dec 2015 00:45:42 +1100
Subject: [PATCH] Rewrite cmd/dht-ping

---
 cmd/dht-ping/main.go | 123 +++++++++++++++++++++++--------------------
 1 file changed, 66 insertions(+), 57 deletions(-)

diff --git a/cmd/dht-ping/main.go b/cmd/dht-ping/main.go
index 94d93368..b23aa0cf 100644
--- a/cmd/dht-ping/main.go
+++ b/cmd/dht-ping/main.go
@@ -2,83 +2,92 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	"log"
+	"math"
 	"net"
 	"os"
 	"time"
 
+	"github.com/anacrolix/tagflag"
+	"github.com/bradfitz/iter"
+
 	"github.com/anacrolix/torrent/dht"
 )
 
-type pingResponse struct {
-	addr  string
-	krpc  dht.Msg
-	msgOk bool
-	rtt   time.Duration
-}
-
 func main() {
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
-	timeout := flag.Duration("timeout", -1, "maximum timeout")
-	flag.Parse()
-	pingStrAddrs := flag.Args()
-	if len(pingStrAddrs) == 0 {
-		os.Stderr.WriteString("u must specify addrs of nodes to ping e.g. router.bittorrent.com:6881\n")
-		os.Exit(2)
+	var args = struct {
+		Timeout time.Duration
+		Nodes   []string `type:"pos" arity:"+" help:"nodes to ping e.g. router.bittorrent.com:6881"`
+	}{
+		Timeout: math.MaxInt64,
 	}
+	tagflag.Parse(&args)
 	s, err := dht.NewServer(nil)
 	if err != nil {
 		log.Fatal(err)
 	}
 	log.Printf("dht server on %s", s.Addr())
-	pingResponsesChan := make(chan pingResponse)
-	timeoutChan := make(chan struct{})
-	go func() {
-		for i, netloc := range pingStrAddrs {
-			if i != 0 {
-				time.Sleep(1 * time.Millisecond)
-			}
-			addr, err := net.ResolveUDPAddr("udp4", netloc)
-			if err != nil {
-				log.Fatal(err)
-			}
-			t, err := s.Ping(addr)
-			if err != nil {
-				log.Fatal(err)
-			}
-			start := time.Now()
-			t.SetResponseHandler(func(addr string) func(dht.Msg, bool) {
-				return func(resp dht.Msg, ok bool) {
-					pingResponsesChan <- pingResponse{
-						addr:  addr,
-						krpc:  resp,
-						rtt:   time.Now().Sub(start),
-						msgOk: ok,
-					}
-				}
-			}(netloc))
-		}
-		if *timeout >= 0 {
-			time.Sleep(*timeout)
-			close(timeoutChan)
-		}
-	}()
-	responses := 0
-pingResponsesLoop:
-	for _ = range pingStrAddrs {
+	timeout := time.After(args.Timeout)
+	pongChan := make(chan pong)
+	startPings(s, pongChan, args.Nodes)
+	numResp := receivePongs(pongChan, timeout, len(args.Nodes))
+	fmt.Printf("%d/%d responses (%f%%)\n", numResp, len(args.Nodes), 100*float64(numResp)/float64(len(args.Nodes)))
+}
+
+func receivePongs(pongChan chan pong, timeout <-chan time.Time, maxPongs int) (numResp int) {
+	for range iter.N(maxPongs) {
 		select {
-		case resp := <-pingResponsesChan:
-			if !resp.msgOk {
+		case pong := <-pongChan:
+			if !pong.msgOk {
 				break
 			}
-			responses++
-			fmt.Printf("%-65s %s\n", fmt.Sprintf("%x (%s):", resp.krpc.SenderID(), resp.addr), resp.rtt)
-		case <-timeoutChan:
-			break pingResponsesLoop
+			numResp++
+			fmt.Printf("%-65s %s\n", fmt.Sprintf("%x (%s):", pong.krpc.SenderID(), pong.addr), pong.rtt)
+		case <-timeout:
+			fmt.Fprintf(os.Stderr, "timed out\n")
+			return
+		}
+	}
+	return
+}
+
+func startPings(s *dht.Server, pongChan chan pong, nodes []string) {
+	for i, addr := range nodes {
+		if i != 0 {
+			// Put a small sleep between pings to avoid network issues.
+			time.Sleep(1 * time.Millisecond)
 		}
+		ping(addr, pongChan, s)
 	}
-	// timeouts := len(pingStrAddrs) - responses
-	fmt.Printf("%d/%d responses (%f%%)\n", responses, len(pingStrAddrs), 100*float64(responses)/float64(len(pingStrAddrs)))
+}
+
+type pong struct {
+	addr  string
+	krpc  dht.Msg
+	msgOk bool
+	rtt   time.Duration
+}
+
+func ping(netloc string, pongChan chan pong, s *dht.Server) {
+	addr, err := net.ResolveUDPAddr("udp4", netloc)
+	if err != nil {
+		log.Fatal(err)
+	}
+	t, err := s.Ping(addr)
+	if err != nil {
+		log.Fatal(err)
+	}
+	start := time.Now()
+	t.SetResponseHandler(func(addr string) func(dht.Msg, bool) {
+		return func(resp dht.Msg, ok bool) {
+			pongChan <- pong{
+				addr:  addr,
+				krpc:  resp,
+				rtt:   time.Now().Sub(start),
+				msgOk: ok,
+			}
+		}
+	}(netloc))
 }
-- 
2.48.1