diff --git a/cmd/atenet/internal/router/health.go b/cmd/atenet/internal/router/health.go index d9e1fddab..d5bb845d4 100644 --- a/cmd/atenet/internal/router/health.go +++ b/cmd/atenet/internal/router/health.go @@ -16,6 +16,7 @@ package router import ( "context" + "encoding/json" "fmt" "io" "log/slog" @@ -26,9 +27,12 @@ import ( "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" + "k8s.io/apimachinery/pkg/version" "k8s.io/client-go/kubernetes" ) +const dependencyHealthCheckTimeout = 500 * time.Millisecond + type ComponentHealth struct { Healthy bool `json:"healthy"` Message string `json:"message,omitempty"` @@ -89,65 +93,49 @@ func (rh *routerHealth) Start(ctx context.Context) { } func (rh *routerHealth) check(ctx context.Context) { - rh.mu.Lock() - defer rh.mu.Unlock() - slog.InfoContext(ctx, "Checking health") - // 1. Check Envoy - { - healthy, msg := rh.checkEnvoy(ctx) - if healthy { - rh.report.Envoy.Healthy = true - rh.report.Envoy.Message = msg - rh.report.Envoy.LastSuccess = time.Now() - rh.report.Envoy.SuccessCount++ - } else { - rh.report.Envoy.Healthy = false - rh.report.Envoy.Message = msg - rh.report.Envoy.LastFailure = time.Now() - rh.report.Envoy.FailureCount++ - slog.ErrorContext(ctx, "Envoy health check failed", slog.String("msg", msg)) - } + // Run network checks without holding the report mutex, so status requests + // can continue serving the last completed report while a dependency is slow. + envoyHealthy, envoyMsg := rh.checkEnvoy(ctx) + envoyCheckedAt := time.Now() + if !envoyHealthy { + slog.ErrorContext(ctx, "Envoy health check failed", slog.String("msg", envoyMsg)) } - // 2. Check Kubernetes API - { - healthy, msg := rh.checkK8s() - if healthy { - rh.report.K8sAPI.Healthy = true - rh.report.K8sAPI.Message = msg - rh.report.K8sAPI.LastSuccess = time.Now() - rh.report.K8sAPI.SuccessCount++ - } else { - rh.report.K8sAPI.Healthy = false - rh.report.K8sAPI.Message = msg - rh.report.K8sAPI.LastFailure = time.Now() - rh.report.K8sAPI.FailureCount++ - slog.ErrorContext(ctx, "Kubernetes API health check failed", slog.String("msg", msg)) - } + k8sHealthy, k8sMsg := rh.checkK8s(ctx) + k8sCheckedAt := time.Now() + if !k8sHealthy { + slog.ErrorContext(ctx, "Kubernetes API health check failed", slog.String("msg", k8sMsg)) } - // 3. Check ATE API gRPC - { - healthy, msg := rh.checkAteAPI(ctx) - if healthy { - rh.report.AteAPI.Healthy = true - rh.report.AteAPI.Message = msg - rh.report.AteAPI.LastSuccess = time.Now() - rh.report.AteAPI.SuccessCount++ - } else { - rh.report.AteAPI.Healthy = false - rh.report.AteAPI.Message = msg - rh.report.AteAPI.LastFailure = time.Now() - rh.report.AteAPI.FailureCount++ - slog.ErrorContext(ctx, "ATE API gRPC health check failed", slog.String("msg", msg)) - } + ateHealthy, ateMsg := rh.checkAteAPI(ctx) + ateCheckedAt := time.Now() + if !ateHealthy { + slog.ErrorContext(ctx, "ATE API gRPC health check failed", slog.String("msg", ateMsg)) + } + + rh.mu.Lock() + defer rh.mu.Unlock() + updateComponentHealth(&rh.report.Envoy, envoyHealthy, envoyMsg, envoyCheckedAt) + updateComponentHealth(&rh.report.K8sAPI, k8sHealthy, k8sMsg, k8sCheckedAt) + updateComponentHealth(&rh.report.AteAPI, ateHealthy, ateMsg, ateCheckedAt) +} + +func updateComponentHealth(health *ComponentHealth, healthy bool, msg string, checkedAt time.Time) { + health.Healthy = healthy + health.Message = msg + if healthy { + health.LastSuccess = checkedAt + health.SuccessCount++ + } else { + health.LastFailure = checkedAt + health.FailureCount++ } } func (rh *routerHealth) checkEnvoy(ctx context.Context) (bool, string) { - timeoutCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + timeoutCtx, cancel := context.WithTimeout(ctx, dependencyHealthCheckTimeout) defer cancel() req, err := http.NewRequestWithContext(timeoutCtx, "GET", "http://127.0.0.1:9901/ready", nil) @@ -178,15 +166,26 @@ func (rh *routerHealth) checkEnvoy(ctx context.Context) (bool, string) { return true, "LIVE" } -func (rh *routerHealth) checkK8s() (bool, string) { +func (rh *routerHealth) checkK8s(ctx context.Context) (bool, string) { if rh.clientset == nil { return true, "Skipped (standalone/file store)" } - ver, err := rh.clientset.Discovery().ServerVersion() + timeoutCtx, cancel := context.WithTimeout(ctx, dependencyHealthCheckTimeout) + defer cancel() + + restClient := rh.clientset.Discovery().RESTClient() + if restClient == nil { + return false, "Kubernetes discovery REST client is unavailable" + } + body, err := restClient.Get().AbsPath("/version").Do(timeoutCtx).Raw() if err != nil { return false, err.Error() } + ver := &version.Info{} + if err := json.Unmarshal(body, ver); err != nil { + return false, fmt.Sprintf("decoding Kubernetes version: %v", err) + } return true, fmt.Sprintf("Version: %s", ver.GitVersion) } @@ -196,7 +195,7 @@ func (rh *routerHealth) checkAteAPI(ctx context.Context) (bool, string) { return false, "No client" } - timeoutCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + timeoutCtx, cancel := context.WithTimeout(ctx, dependencyHealthCheckTimeout) defer cancel() _, err := rh.apiClient.ListActors(timeoutCtx, &ateapipb.ListActorsRequest{PageSize: 1}) diff --git a/cmd/atenet/internal/router/health_test.go b/cmd/atenet/internal/router/health_test.go new file mode 100644 index 000000000..fb4a5bdd4 --- /dev/null +++ b/cmd/atenet/internal/router/health_test.go @@ -0,0 +1,195 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "k8s.io/client-go/kubernetes" + kubernetesfake "k8s.io/client-go/kubernetes/fake" + "k8s.io/client-go/rest" +) + +type healthRoundTripFunc func(*http.Request) (*http.Response, error) + +func (f healthRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func newHealthTestClientset(t *testing.T, server *httptest.Server) kubernetes.Interface { + t.Helper() + clientset, err := kubernetes.NewForConfig(&rest.Config{Host: server.URL}) + if err != nil { + t.Fatalf("creating Kubernetes client: %v", err) + } + return clientset +} + +func setHealthyEnvoyClient(rh *routerHealth) { + rh.envoyClient = &http.Client{Transport: healthRoundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("LIVE")), + }, nil + })} +} + +func TestCheckK8sTimesOut(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, req *http.Request) { + <-req.Context().Done() + })) + defer server.Close() + + rh := newRouterHealth(time.Second, newHealthTestClientset(t, server), nil, RouterConfig{}) + startedAt := time.Now() + healthy, msg := rh.checkK8s(context.Background()) + elapsed := time.Since(startedAt) + + if healthy { + t.Fatal("checkK8s returned healthy for a stalled request") + } + if !strings.Contains(msg, context.DeadlineExceeded.Error()) { + t.Fatalf("checkK8s message = %q, want context deadline exceeded", msg) + } + if elapsed > 3*dependencyHealthCheckTimeout { + t.Fatalf("checkK8s took %v, want a bounded timeout near %v", elapsed, dependencyHealthCheckTimeout) + } +} + +func TestCheckK8sWithoutRESTClient(t *testing.T) { + rh := newRouterHealth(time.Second, kubernetesfake.NewSimpleClientset(), nil, RouterConfig{}) + healthy, msg := rh.checkK8s(context.Background()) + if healthy { + t.Fatal("checkK8s returned healthy without a discovery REST client") + } + if msg != "Kubernetes discovery REST client is unavailable" { + t.Fatalf("checkK8s message = %q, want unavailable REST client", msg) + } +} + +func TestHealthCheckDoesNotBlockReportOrStatusz(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseRequest := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(releaseRequest) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + close(started) + select { + case <-release: + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"gitVersion":"v1.36.1"}`) + case <-req.Context().Done(): + } + })) + defer server.Close() + + rh := newRouterHealth(time.Second, newHealthTestClientset(t, server), nil, RouterConfig{}) + setHealthyEnvoyClient(rh) + checkDone := make(chan struct{}) + go func() { + rh.check(context.Background()) + close(checkDone) + }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("Kubernetes health request did not start") + } + + reportDone := make(chan struct{}) + go func() { + _ = rh.Report() + close(reportDone) + }() + select { + case <-reportDone: + case <-time.After(100 * time.Millisecond): + t.Fatal("Report blocked while a dependency health request was in flight") + } + + statusServer := httptest.NewServer(http.HandlerFunc((&RouterServer{ + cfg: RouterConfig{Standalone: true}, + health: rh, + }).handleStatusz)) + defer statusServer.Close() + statusClient := &http.Client{Timeout: 200 * time.Millisecond} + resp, err := statusClient.Get(statusServer.URL + "/statusz?format=json") + if err != nil { + t.Fatalf("GET /statusz while health check was in flight: %v", err) + } + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("GET /statusz status = %s, want 200 OK", resp.Status) + } + + releaseRequest() + select { + case <-checkDone: + case <-time.After(time.Second): + t.Fatal("health check did not finish after the Kubernetes response was released") + } + + report := rh.Report() + if !report.Envoy.Healthy || report.Envoy.SuccessCount != 1 || report.Envoy.LastSuccess.IsZero() { + t.Errorf("Envoy health = %+v, want one successful check", report.Envoy) + } + if !report.K8sAPI.Healthy || report.K8sAPI.SuccessCount != 1 || report.K8sAPI.LastSuccess.IsZero() { + t.Errorf("Kubernetes health = %+v, want one successful check", report.K8sAPI) + } + if report.K8sAPI.Message != "Version: v1.36.1" { + t.Errorf("Kubernetes health message = %q, want %q", report.K8sAPI.Message, "Version: v1.36.1") + } + if report.AteAPI.Healthy || report.AteAPI.FailureCount != 1 || report.AteAPI.LastFailure.IsZero() { + t.Errorf("ATE API health = %+v, want one failed check", report.AteAPI) + } +} + +func TestHealthStartStopsWhenK8sCheckIsCanceled(t *testing.T) { + started := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, req *http.Request) { + close(started) + <-req.Context().Done() + })) + defer server.Close() + + rh := newRouterHealth(time.Hour, newHealthTestClientset(t, server), nil, RouterConfig{}) + setHealthyEnvoyClient(rh) + ctx, cancel := context.WithCancel(context.Background()) + startDone := make(chan struct{}) + go func() { + rh.Start(ctx) + close(startDone) + }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("Kubernetes health request did not start") + } + cancel() + select { + case <-startDone: + case <-time.After(time.Second): + t.Fatal("router health loop did not stop after cancellation") + } +}