Merge pull request #626 from libp2p/fix/ping

ping: return a stream of results

p2p/protocol/ping/ping.go

@@ -71,42 +71,62 @@ func (p *PingService) PingHandler(s inet.Stream) {
 	}
 }
 
-func (ps *PingService) Ping(ctx context.Context, p peer.ID) (<-chan time.Duration, error) {
+// Result is a result of a ping attempt, either an RTT or an error.
+type Result struct {
+	RTT   time.Duration
+	Error error
+}
+
+func (ps *PingService) Ping(ctx context.Context, p peer.ID) <-chan Result {
 	return Ping(ctx, ps.Host, p)
 }
 
-func Ping(ctx context.Context, h host.Host, p peer.ID) (<-chan time.Duration, error) {
+// Ping pings the remote peer until the context is canceled, returning a stream
+// of RTTs or errors.
+func Ping(ctx context.Context, h host.Host, p peer.ID) <-chan Result {
 	s, err := h.NewStream(ctx, p, ID)
 	if err != nil {
-		return nil, err
+		ch := make(chan Result, 1)
+		ch <- Result{Error: err}
+		close(ch)
+		return ch
 	}
 
-	out := make(chan time.Duration)
+	ctx, cancel := context.WithCancel(ctx)
+
+	out := make(chan Result)
 	go func() {
 		defer close(out)
-		defer s.Reset()
-		for {
+		defer cancel()
+
+		for ctx.Err() == nil {
+			var res Result
+			res.RTT, res.Error = ping(s)
+
+			// canceled, ignore everything.
+			if ctx.Err() != nil {
+				return
+			}
+
+			// No error, record the RTT.
+			if res.Error == nil {
+				h.Peerstore().RecordLatency(p, res.RTT)
+			}
+
 			select {
+			case out <- res:
 			case <-ctx.Done():
 				return
-			default:
-				t, err := ping(s)
-				if err != nil {
-					log.Debugf("ping error: %s", err)
-					return
-				}
-
-				h.Peerstore().RecordLatency(p, t)
-				select {
-				case out <- t:
-				case <-ctx.Done():
-					return
-				}
 			}
 		}
 	}()
+	go func() {
+		// forces the ping to abort.
+		<-ctx.Done()
+		s.Reset()
+	}()
 
-	return out, nil
+	return out
 }
 
 func ping(s inet.Stream) (time.Duration, error) {

p2p/protocol/ping/ping_test.go

@@ -37,15 +37,15 @@ func TestPing(t *testing.T) {
 func testPing(t *testing.T, ps *ping.PingService, p peer.ID) {
 	pctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	ts, err := ps.Ping(pctx, p)
-	if err != nil {
-		t.Fatal(err)
-	}
+	ts := ps.Ping(pctx, p)
 
 	for i := 0; i < 5; i++ {
 		select {
-		case took := <-ts:
-			t.Log("ping took: ", took)
+		case res := <-ts:
+			if res.Error != nil {
+				t.Fatal(res.Error)
+			}
+			t.Log("ping took: ", res.RTT)
 		case <-time.After(time.Second * 4):
 			t.Fatal("failed to receive ping")
 		}