diff --git a/cmd/atenet/internal/router/health.go b/cmd/atenet/internal/router/health.go index d9e1fddab..21332c409 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"` @@ -44,6 +48,12 @@ type RouterHealthReport struct { AteAPI ComponentHealth `json:"ate_api"` } +type componentHealthCheckResult struct { + healthy bool + message string + checkedAt time.Time +} + // routerHealth periodically checks the dependent services of router to track health // status of this component. type routerHealth struct { @@ -89,65 +99,65 @@ func (rh *routerHealth) Start(ctx context.Context) { } func (rh *routerHealth) check(ctx context.Context) { + slog.InfoContext(ctx, "Checking health") + + // Run network checks concurrently and without holding the report mutex, so + // the cycle is bounded by the slowest dependency and status requests can + // continue serving the last completed report. + var envoyResult, k8sResult, ateResult componentHealthCheckResult + var wg sync.WaitGroup + wg.Add(3) + go func() { + defer wg.Done() + envoyResult = runComponentHealthCheck(ctx, "Envoy health check failed", rh.checkEnvoy) + }() + go func() { + defer wg.Done() + k8sResult = runComponentHealthCheck(ctx, "Kubernetes API health check failed", rh.checkK8s) + }() + go func() { + defer wg.Done() + ateResult = runComponentHealthCheck(ctx, "ATE API gRPC health check failed", rh.checkAteAPI) + }() + wg.Wait() + rh.mu.Lock() defer rh.mu.Unlock() + updateComponentHealth(&rh.report.Envoy, envoyResult.healthy, envoyResult.message, envoyResult.checkedAt) + updateComponentHealth(&rh.report.K8sAPI, k8sResult.healthy, k8sResult.message, k8sResult.checkedAt) + updateComponentHealth(&rh.report.AteAPI, ateResult.healthy, ateResult.message, ateResult.checkedAt) +} - 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)) - } +func runComponentHealthCheck( + ctx context.Context, + failureMessage string, + check func(context.Context) (bool, string), +) componentHealthCheckResult { + healthy, message := check(ctx) + if !healthy { + slog.ErrorContext(ctx, failureMessage, slog.String("msg", message)) } - - // 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)) - } + return componentHealthCheckResult{ + healthy: healthy, + message: message, + checkedAt: time.Now(), } +} - // 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)) - } +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 +188,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 +217,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..29d410f7b --- /dev/null +++ b/cmd/atenet/internal/router/health_test.go @@ -0,0 +1,283 @@ +// 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" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" + "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) +} + +type healthControlClient struct { + ateapipb.ControlClient + listActorsFn func(context.Context, *ateapipb.ListActorsRequest, ...grpc.CallOption) (*ateapipb.ListActorsResponse, error) +} + +func (c *healthControlClient) ListActors( + ctx context.Context, + req *ateapipb.ListActorsRequest, + opts ...grpc.CallOption, +) (*ateapipb.ListActorsResponse, error) { + return c.listActorsFn(ctx, req, opts...) +} + +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 TestHealthChecksRunConcurrently(t *testing.T) { + started := make(chan string, 3) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseChecks := func() { releaseOnce.Do(func() { close(release) }) } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + started <- "k8s" + select { + case <-release: + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"gitVersion":"v1.36.1"}`) + case <-req.Context().Done(): + } + })) + defer server.Close() + + apiClient := &healthControlClient{ + listActorsFn: func(ctx context.Context, _ *ateapipb.ListActorsRequest, _ ...grpc.CallOption) (*ateapipb.ListActorsResponse, error) { + started <- "ateapi" + select { + case <-release: + return &ateapipb.ListActorsResponse{}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + rh := newRouterHealth(time.Second, newHealthTestClientset(t, server), apiClient, RouterConfig{}) + rh.envoyClient = &http.Client{Transport: healthRoundTripFunc(func(*http.Request) (*http.Response, error) { + started <- "envoy" + <-release + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("LIVE")), + }, nil + })} + + checkDone := make(chan struct{}) + go func() { + rh.check(context.Background()) + close(checkDone) + }() + defer func() { + releaseChecks() + <-checkDone + }() + + seen := make(map[string]bool, 3) + timer := time.NewTimer(time.Second) + defer timer.Stop() + for len(seen) < 3 { + select { + case dependency := <-started: + seen[dependency] = true + case <-timer.C: + t.Fatalf("started health checks = %v, want envoy, k8s, and ateapi before any check finishes", seen) + } + } + + releaseChecks() + select { + case <-checkDone: + case <-time.After(time.Second): + t.Fatal("health check did not finish after dependencies were released") + } + + report := rh.Report() + if !report.Envoy.Healthy || !report.K8sAPI.Healthy || !report.AteAPI.Healthy { + t.Errorf("health report = %+v, want all dependencies healthy", report) + } +} + +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") + } +}