feat: unified timeout or interrupt messages

macrat · Apr 24, 2021 · ea59440 · ea59440
1 parent df8c6ea
commit ea59440
Show file tree

Hide file tree

Showing 8 changed files with 60 additions and 11 deletions.
diff --git a/probe/http.go b/probe/http.go
@@ -98,9 +98,12 @@ func (p HTTPProbe) Check(ctx context.Context) []store.Record {
 		if e, ok := errors.Unwrap(errors.Unwrap(err)).(*net.DNSError); ok && e.IsNotFound {
 			status = store.STATUS_UNKNOWN
 		}
-		if e := errors.Unwrap(err); e != nil && e.Error() == "context deadline exceeded" {
-			status = store.STATUS_UNKNOWN
-			message = "timed out or interrupted"
+		if e := errors.Unwrap(err); e != nil {
+			switch e.Error() {
+			case "context deadline exceeded", "context canceled":
+				status = store.STATUS_UNKNOWN
+				message = "timed out or interrupted"
+			}
 		}
 	} else {
 		message = resp.Status

diff --git a/probe/http_test.go b/probe/http_test.go
@@ -25,4 +25,6 @@ func TestHTTPProbe(t *testing.T) {
 		{strings.Replace(server.URL, "http", "http-options", 1) + "/only/options", store.STATUS_HEALTHY, `200 OK`},
 		{server.URL + "/slow-page", store.STATUS_UNKNOWN, `timed out or interrupted`},
 	})
+
+	AssertTimeout(t, server.URL)
 }
diff --git a/probe/ping.go b/probe/ping.go
@@ -61,7 +61,7 @@ func (p PingProbe) Check(ctx context.Context) []store.Record {
 	err = pinger.Run()
 	if err != nil {
 		return []store.Record{{
-			CheckedAt: time.Now(),
+			CheckedAt: startTime,
 			Target:    p.target,
 			Status:    store.STATUS_UNKNOWN,
 			Message:   err.Error(),
@@ -76,18 +76,27 @@ func (p PingProbe) Check(ctx context.Context) []store.Record {
 		status = store.STATUS_HEALTHY
 	}
 
-	return []store.Record{{
-		CheckedAt: startTime,
-		Target:    p.target,
-		Status:    status,
-		Message: fmt.Sprintf(
+	var message string
+	select {
+	case <-ctx.Done():
+		status = store.STATUS_UNKNOWN
+		message = "timed out or interrupted"
+	default:
+		message = fmt.Sprintf(
 			"rtt(min/avg/max)=%.2f/%.2f/%.2f send/rcv=%d/%d",
 			float64(stat.MinRtt.Microseconds())/1000,
 			float64(stat.AvgRtt.Microseconds())/1000,
 			float64(stat.MaxRtt.Microseconds())/1000,
 			pinger.PacketsSent,
 			pinger.PacketsRecv,
-		),
-		Latency: stat.AvgRtt,
+		)
+	}
+
+	return []store.Record{{
+		CheckedAt: startTime,
+		Target:    p.target,
+		Status:    status,
+		Message:   message,
+		Latency:   stat.AvgRtt,
 	}}
 }
diff --git a/probe/ping_test.go b/probe/ping_test.go
@@ -12,4 +12,6 @@ func TestPingProbe(t *testing.T) {
 	AssertProbe(t, []ProbeTest{
 		{"ping:localhost", store.STATUS_HEALTHY, `rtt\(min/avg/max\)=[0-9.]*/[0-9.]*/[0-9.]* send/rcv=4/4`},
 	})
+
+	AssertTimeout(t, "ping:localhost")
 }
diff --git a/probe/probe_test.go b/probe/probe_test.go
@@ -123,6 +123,31 @@ func AssertProbe(t *testing.T, tests []ProbeTest) {
 	}
 }
 
+func AssertTimeout(t *testing.T, target string) {
+	t.Run("timeout-or-interrupt", func(t *testing.T) {
+		p, err := probe.New(target)
+		if err != nil {
+			t.Fatalf("failed to get probe: %s", err)
+		}
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		records := p.Check(ctx)
+		if len(records) != 1 {
+			t.Fatalf("unexpected number of records: %#v", records)
+		}
+
+		if records[0].Message != "timed out or interrupted" {
+			t.Errorf("unexpected message: %s", records[0].Message)
+		}
+
+		if records[0].Status != store.STATUS_UNKNOWN {
+			t.Errorf("unexpected status: %s", records[0].Status)
+		}
+	})
+}
+
 func RunDummyHTTPServer() *httptest.Server {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/ok", func(w http.ResponseWriter, r *http.Request) {

diff --git a/probe/tcp.go b/probe/tcp.go
@@ -55,6 +55,10 @@ func (p TCPProbe) Check(ctx context.Context) []store.Record {
 		if e, ok := errors.Unwrap(err).(*net.DNSError); ok && e.IsNotFound {
 			r.Status = store.STATUS_UNKNOWN
 		}
+		if e := errors.Unwrap(err); e != nil && e.Error() == "operation was canceled" {
+			r.Status = store.STATUS_UNKNOWN
+			r.Message = "timed out or interrupted"
+		}
 	} else {
 		r.Status = store.STATUS_HEALTHY
 		r.Message = fmt.Sprintf("%s -> %s", conn.LocalAddr(), conn.RemoteAddr())

diff --git a/probe/tcp_otheros_test.go b/probe/tcp_otheros_test.go
@@ -19,4 +19,6 @@ func TestTCPProbe(t *testing.T) {
 		{strings.Replace(server.URL, "http://", "tcp:", 1), store.STATUS_HEALTHY, `(127\.0\.0\.1|\[::1\]):[0-9]+ -> (127\.0\.0\.1|\[::1\]):[0-9]+`},
 		{"tcp:localhost:56789", store.STATUS_FAILURE, `dial tcp (127\.0\.0\.1|\[::1\]):56789: connect: connection refused`},
 	})
+
+	AssertTimeout(t, strings.Replace(server.URL, "http://", "tcp:", 1))
 }
diff --git a/probe/tcp_windows_test.go b/probe/tcp_windows_test.go
@@ -19,4 +19,6 @@ func TestTCPProbe(t *testing.T) {
 		{strings.Replace(server.URL, "http://", "tcp:", 1), store.STATUS_HEALTHY, `(127\.0\.0\.1|\[::1\]):[0-9]+ -> (127\.0\.0\.1|\[::1\]):[0-9]+`},
 		{"tcp:localhost:56789", store.STATUS_FAILURE, `dial tcp (127\.0\.0\.1|\[::1\]):56789: connectex: No connection could be made because the target machine actively refused it.`},
 	})
+
+	AssertTimeout(t, strings.Replace(server.URL, "http://", "tcp:", 1))
 }