Make router health checks bounded and non-blocking

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})

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")
+	}
+}