cmd/dht-ping: Various improvements more like a real ping program

cmd/dht-ping/main.go

@@ -2,9 +2,11 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"log"
 	"net"
 	"os"
+	"time"
 
 	"bitbucket.org/anacrolix/go.torrent/dht"
 )
@@ -12,10 +14,12 @@ import (
 type pingResponse struct {
 	addr string
 	krpc dht.Msg
+	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 {
@@ -28,23 +32,49 @@ func main() {
 	}
 	log.Printf("dht server on %s", s.LocalAddr())
 	pingResponses := make(chan pingResponse)
-	for _, netloc := range pingStrAddrs {
-		addr, err := net.ResolveUDPAddr("udp4", netloc)
-		if err != nil {
-			log.Fatal(err)
-		}
-		t, err := s.Ping(addr)
-		if err != nil {
-			log.Fatal(err)
-		}
-		go func(addr string) {
-			pingResponses <- pingResponse{
-				addr: addr,
-				krpc: <-t.Response,
+	timeoutChan := make(chan struct{})
+	go func() {
+		for i, netloc := range pingStrAddrs {
+			if i != 0 {
+				time.Sleep(1 * time.Millisecond)
 			}
-		}(netloc)
-	}
+			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()
+			go func(addr string) {
+				resp := <-t.Response
+				pingResponses <- pingResponse{
+					addr: addr,
+					krpc: resp,
+					rtt:  time.Now().Sub(start),
+				}
+			}(netloc)
+		}
+		if *timeout >= 0 {
+			time.Sleep(*timeout)
+			close(timeoutChan)
+		}
+	}()
+	responses := 0
+pingResponses:
 	for _ = range pingStrAddrs {
-		log.Printf("%q", <-pingResponses)
+		select {
+		case resp := <-pingResponses:
+			if resp.krpc == nil {
+				break
+			}
+			responses++
+			fmt.Printf("%-65s %s\n", fmt.Sprintf("%x (%s):", resp.krpc["r"].(map[string]interface{})["id"].(string), resp.addr), resp.rtt)
+		case <-timeoutChan:
+			break pingResponses
+		}
 	}
+	// timeouts := len(pingStrAddrs) - responses
+	fmt.Printf("%d/%d responses (%f%%)\n", responses, len(pingStrAddrs), 100*float64(responses)/float64(len(pingStrAddrs)))
 }