diff --git a/state/prefix_health.go b/state/prefix_health.go index 3a1f90d..c34a370 100644 --- a/state/prefix_health.go +++ b/state/prefix_health.go @@ -2,6 +2,7 @@ package state import ( "fmt" + "io" "log/slog" "net" "net/http" @@ -196,24 +197,39 @@ func (h *HTTPPrefixHealth) Start(log *slog.Logger, t *RouterTunables) { go func() { ticker := time.NewTicker(*h.Delay) defer ticker.Stop() + client := &http.Client{Timeout: *h.Delay} for h.running.Load() { // TODO: add a way to interrupt this sleep, if ticker has a high delay <-ticker.C - // HTTP probe logic would go here - startTime := time.Now() - resp, err := http.Get(h.URL) - if err != nil || resp.StatusCode != http.StatusOK { - // failed - h.lastMetric.Store(INF) - log.Debug("prefix healthcheck failed", "prefix", h.Prefix.String(), "url", h.URL, "error", err) - } else { - // success - rtt := time.Since(startTime) - h.lastMetric.Store(DurationToMetric(rtt)) - } + h.checkHTTP(log, client) } }() } +func (h *HTTPPrefixHealth) checkHTTP(log *slog.Logger, client *http.Client) { + startTime := time.Now() + resp, err := client.Get(h.URL) + if err != nil { + h.lastMetric.Store(INF) + log.Debug("prefix healthcheck failed", "prefix", h.Prefix.String(), "url", h.URL, "error", err) + return + } + _, drainErr := io.Copy(io.Discard, resp.Body) + closeErr := resp.Body.Close() + if drainErr != nil { + log.Debug("failed to drain prefix healthcheck response", "prefix", h.Prefix.String(), "url", h.URL, "error", drainErr) + } + if closeErr != nil { + log.Debug("failed to close prefix healthcheck response", "prefix", h.Prefix.String(), "url", h.URL, "error", closeErr) + } + if resp.StatusCode != http.StatusOK { + h.lastMetric.Store(INF) + log.Debug("prefix healthcheck failed", "prefix", h.Prefix.String(), "url", h.URL, "status", resp.StatusCode) + return + } + rtt := time.Since(startTime) + h.lastMetric.Store(DurationToMetric(rtt)) +} + func (h *HTTPPrefixHealth) sameConfig(other PrefixHealth, tunables *RouterTunables) bool { o, ok := other.(*HTTPPrefixHealth) return ok && diff --git a/state/prefix_health_test.go b/state/prefix_health_test.go index 658c95a..e6e58a5 100644 --- a/state/prefix_health_test.go +++ b/state/prefix_health_test.go @@ -1,7 +1,11 @@ package state import ( + "io" + "log/slog" + "net/http" "net/netip" + "strings" "testing" "time" @@ -110,3 +114,46 @@ delay: 5s }) } } + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type closeTrackingBody struct { + *strings.Reader + closed bool +} + +func (b *closeTrackingBody) Close() error { + b.closed = true + return nil +} + +func TestHTTPPrefixHealthClosesResponseBody(t *testing.T) { + for _, statusCode := range []int{http.StatusOK, http.StatusServiceUnavailable} { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + body := &closeTrackingBody{Reader: strings.NewReader("ok")} + client := &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: statusCode, + Body: body, + Header: make(http.Header), + Request: req, + }, nil + }), + } + health := &HTTPPrefixHealth{ + Prefix: netip.MustParsePrefix("172.16.0.0/16"), + URL: "http://example.com/health", + } + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + + health.checkHTTP(logger, client) + + assert.True(t, body.closed) + }) + } +}