Add health checks

proxy/internal/debug/handler.go

@@ -46,15 +46,15 @@ func formatDuration(d time.Duration) string {
 	}
 }
 
-// ClientProvider provides access to NetBird clients.
-type ClientProvider interface {
+// clientProvider provides access to NetBird clients.
+type clientProvider interface {
 	GetClient(accountID types.AccountID) (*nbembed.Client, bool)
 	ListClientsForDebug() map[types.AccountID]roundtrip.ClientDebugInfo
 }
 
 // Handler provides HTTP debug endpoints.
 type Handler struct {
-	provider   ClientProvider
+	provider   clientProvider
 	logger     *log.Logger
 	startTime  time.Time
 	templates  *template.Template
@@ -62,7 +62,7 @@ type Handler struct {
 }
 
 // NewHandler creates a new debug handler.
-func NewHandler(provider ClientProvider, logger *log.Logger) *Handler {
+func NewHandler(provider clientProvider, logger *log.Logger) *Handler {
 	if logger == nil {
 		logger = log.StandardLogger()
 	}

proxy/internal/health/health.go

@@ -0,0 +1,340 @@
+// Package health provides health probes for the proxy server.
+package health
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/embed"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+const (
+	maxConcurrentChecks    = 3
+	maxClientCheckTimeout  = 5 * time.Minute
+)
+
+// clientProvider provides access to NetBird clients for health checks.
+type clientProvider interface {
+	ListClientsForStartup() map[types.AccountID]*embed.Client
+}
+
+// Checker tracks health state and provides probe endpoints.
+type Checker struct {
+	logger   *log.Logger
+	provider clientProvider
+
+	mu                  sync.RWMutex
+	managementConnected bool
+	initialSyncComplete bool
+
+	// checkSem limits concurrent client health checks.
+	checkSem chan struct{}
+}
+
+// ClientHealth represents the health status of a single NetBird client.
+type ClientHealth struct {
+	Healthy             bool   `json:"healthy"`
+	ManagementConnected bool   `json:"management_connected"`
+	SignalConnected     bool   `json:"signal_connected"`
+	RelaysConnected     int    `json:"relays_connected"`
+	RelaysTotal         int    `json:"relays_total"`
+	Error               string `json:"error,omitempty"`
+}
+
+// ProbeResponse represents the JSON response for health probes.
+type ProbeResponse struct {
+	Status  string                           `json:"status"`
+	Checks  map[string]bool                  `json:"checks,omitempty"`
+	Clients map[types.AccountID]ClientHealth `json:"clients,omitempty"`
+}
+
+// Server runs the health probe HTTP server on a dedicated port.
+type Server struct {
+	server  *http.Server
+	logger  *log.Logger
+	checker *Checker
+}
+
+// SetManagementConnected updates the management connection state.
+func (c *Checker) SetManagementConnected(connected bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.managementConnected = connected
+}
+
+// SetInitialSyncComplete marks that the initial mapping sync has completed.
+func (c *Checker) SetInitialSyncComplete() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.initialSyncComplete = true
+}
+
+// CheckClientsConnected verifies all clients are connected to management/signal/relay.
+// Uses the provided context for timeout/cancellation, with a maximum bound of maxClientCheckTimeout.
+// Limits concurrent checks via semaphore.
+func (c *Checker) CheckClientsConnected(ctx context.Context) (bool, map[types.AccountID]ClientHealth) {
+	// Apply upper bound timeout in case parent context has no deadline
+	ctx, cancel := context.WithTimeout(ctx, maxClientCheckTimeout)
+	defer cancel()
+
+	clients := c.provider.ListClientsForStartup()
+
+	// No clients yet means not ready
+	if len(clients) == 0 {
+		return false, make(map[types.AccountID]ClientHealth)
+	}
+
+	type result struct {
+		accountID types.AccountID
+		health    ClientHealth
+	}
+
+	resultsCh := make(chan result, len(clients))
+	var wg sync.WaitGroup
+
+	for accountID, client := range clients {
+		wg.Add(1)
+		go func(id types.AccountID, cl *embed.Client) {
+			defer wg.Done()
+
+			// Acquire semaphore
+			select {
+			case c.checkSem <- struct{}{}:
+				defer func() { <-c.checkSem }()
+			case <-ctx.Done():
+				resultsCh <- result{id, ClientHealth{Healthy: false, Error: ctx.Err().Error()}}
+				return
+			}
+
+			resultsCh <- result{id, checkClientHealth(cl)}
+		}(accountID, client)
+	}
+
+	go func() {
+		wg.Wait()
+		close(resultsCh)
+	}()
+
+	results := make(map[types.AccountID]ClientHealth)
+	allHealthy := true
+	for r := range resultsCh {
+		results[r.accountID] = r.health
+		if !r.health.Healthy {
+			allHealthy = false
+		}
+	}
+
+	return allHealthy, results
+}
+
+// LivenessProbe returns true if the process is alive.
+// This should always return true if we can respond.
+func (c *Checker) LivenessProbe() bool {
+	return true
+}
+
+// ReadinessProbe returns true if the server can accept traffic.
+func (c *Checker) ReadinessProbe() bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.managementConnected
+}
+
+// StartupProbe checks if initial startup is complete.
+// Checks management connection, initial sync, and all client health directly.
+// Uses the provided context for timeout/cancellation.
+func (c *Checker) StartupProbe(ctx context.Context) bool {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	sync := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	if !mgmt || !sync {
+		return false
+	}
+
+	// Check all clients are connected to management/signal/relay
+	allHealthy, _ := c.CheckClientsConnected(ctx)
+	return allHealthy
+}
+
+// Handler returns an http.Handler for health probe endpoints.
+func (c *Checker) Handler() http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz/live", c.handleLiveness)
+	mux.HandleFunc("/healthz/ready", c.handleReadiness)
+	mux.HandleFunc("/healthz/startup", c.handleStartup)
+	mux.HandleFunc("/healthz", c.handleFull)
+	return mux
+}
+
+func (c *Checker) handleLiveness(w http.ResponseWriter, r *http.Request) {
+	if c.LivenessProbe() {
+		c.writeProbeResponse(w, http.StatusOK, "ok", nil, nil)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", nil, nil)
+}
+
+func (c *Checker) handleReadiness(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	checks := map[string]bool{
+		"management_connected": c.managementConnected,
+	}
+	c.mu.RUnlock()
+
+	if c.ReadinessProbe() {
+		c.writeProbeResponse(w, http.StatusOK, "ok", checks, nil)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, nil)
+}
+
+func (c *Checker) handleStartup(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	sync := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	// Check clients directly using request context
+	allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
+
+	checks := map[string]bool{
+		"management_connected":  mgmt,
+		"initial_sync_complete": sync,
+		"all_clients_healthy":   allClientsHealthy,
+	}
+
+	if c.StartupProbe(r.Context()) {
+		c.writeProbeResponse(w, http.StatusOK, "ok", checks, clientHealth)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, clientHealth)
+}
+
+func (c *Checker) handleFull(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	sync := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
+
+	checks := map[string]bool{
+		"management_connected":  mgmt,
+		"initial_sync_complete": sync,
+		"all_clients_healthy":   allClientsHealthy,
+	}
+
+	status := "ok"
+	statusCode := http.StatusOK
+	if !c.ReadinessProbe() {
+		status = "fail"
+		statusCode = http.StatusServiceUnavailable
+	}
+
+	c.writeProbeResponse(w, statusCode, status, checks, clientHealth)
+}
+
+func (c *Checker) writeProbeResponse(w http.ResponseWriter, statusCode int, status string, checks map[string]bool, clients map[types.AccountID]ClientHealth) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+
+	resp := ProbeResponse{
+		Status:  status,
+		Checks:  checks,
+		Clients: clients,
+	}
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		c.logger.Debugf("write health response: %v", err)
+	}
+}
+
+// ListenAndServe starts the health probe server.
+func (s *Server) ListenAndServe() error {
+	s.logger.Infof("starting health probe server on %s", s.server.Addr)
+	return s.server.ListenAndServe()
+}
+
+// Serve starts the health probe server on the given listener.
+func (s *Server) Serve(l net.Listener) error {
+	s.logger.Infof("starting health probe server on %s", l.Addr())
+	return s.server.Serve(l)
+}
+
+// Shutdown gracefully shuts down the health probe server.
+func (s *Server) Shutdown(ctx context.Context) error {
+	return s.server.Shutdown(ctx)
+}
+
+// NewChecker creates a new health checker.
+func NewChecker(logger *log.Logger, provider clientProvider) *Checker {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &Checker{
+		logger:   logger,
+		provider: provider,
+		checkSem: make(chan struct{}, maxConcurrentChecks),
+	}
+}
+
+// NewServer creates a new health probe server.
+func NewServer(addr string, checker *Checker, logger *log.Logger) *Server {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &Server{
+		server: &http.Server{
+			Addr:         addr,
+			Handler:      checker.Handler(),
+			ReadTimeout:  5 * time.Second,
+			WriteTimeout: 5 * time.Second,
+		},
+		logger:  logger,
+		checker: checker,
+	}
+}
+
+func checkClientHealth(client *embed.Client) ClientHealth {
+	status, err := client.Status()
+	if err != nil {
+		return ClientHealth{
+			Healthy: false,
+			Error:   err.Error(),
+		}
+	}
+
+	// Count only rel:// and rels:// relays (not stun/turn)
+	var relayCount, relaysConnected int
+	for _, relay := range status.Relays {
+		if !strings.HasPrefix(relay.URI, "rel://") && !strings.HasPrefix(relay.URI, "rels://") {
+			continue
+		}
+		relayCount++
+		if relay.Err == nil {
+			relaysConnected++
+		}
+	}
+
+	// Client is healthy if connected to management, signal, and at least one relay (if any are defined)
+	healthy := status.ManagementState.Connected &&
+		status.SignalState.Connected &&
+		(relayCount == 0 || relaysConnected > 0)
+
+	return ClientHealth{
+		Healthy:             healthy,
+		ManagementConnected: status.ManagementState.Connected,
+		SignalConnected:     status.SignalState.Connected,
+		RelaysConnected:     relaysConnected,
+		RelaysTotal:         relayCount,
+	}
+}

proxy/internal/health/health_test.go

@@ -0,0 +1,155 @@
+package health
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/embed"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+type mockClientProvider struct {
+	clients map[types.AccountID]*embed.Client
+}
+
+func (m *mockClientProvider) ListClientsForStartup() map[types.AccountID]*embed.Client {
+	return m.clients
+}
+
+func TestChecker_LivenessProbe(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// Liveness should always return true if we can respond.
+	assert.True(t, checker.LivenessProbe())
+}
+
+func TestChecker_ReadinessProbe(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// Initially not ready (management not connected).
+	assert.False(t, checker.ReadinessProbe())
+
+	// After management connects, should be ready.
+	checker.SetManagementConnected(true)
+	assert.True(t, checker.ReadinessProbe())
+
+	// If management disconnects, should not be ready.
+	checker.SetManagementConnected(false)
+	assert.False(t, checker.ReadinessProbe())
+}
+
+func TestChecker_StartupProbe_NoClients(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// Initially startup not complete.
+	assert.False(t, checker.StartupProbe(context.Background()))
+
+	// Just management connected is not enough.
+	checker.SetManagementConnected(true)
+	assert.False(t, checker.StartupProbe(context.Background()))
+
+	// Management + initial sync but no clients = not ready
+	checker.SetInitialSyncComplete()
+	assert.False(t, checker.StartupProbe(context.Background()))
+}
+
+func TestChecker_Handler_Liveness(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+}
+
+func TestChecker_Handler_Readiness_NotReady(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+	assert.False(t, resp.Checks["management_connected"])
+}
+
+func TestChecker_Handler_Readiness_Ready(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.True(t, resp.Checks["management_connected"])
+}
+
+func TestChecker_Handler_Startup_NotComplete(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+}
+
+func TestChecker_Handler_Full(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.NotNil(t, resp.Checks)
+	// Clients may be empty map when no clients exist.
+	assert.Empty(t, resp.Clients)
+}
+
+func TestChecker_StartupProbe_RespectsContext(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+
+	// Cancelled context should return false quickly
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	result := checker.StartupProbe(ctx)
+	assert.False(t, result)
+}

proxy/internal/roundtrip/netbird.go

@@ -56,20 +56,18 @@ type NetBird struct {
 	statusNotifier statusNotifier
 }
 
-// NewNetBird creates a new NetBird transport.
-func NewNetBird(mgmtAddr, proxyID string, logger *log.Logger, notifier statusNotifier) *NetBird {
-	if logger == nil {
-		logger = log.StandardLogger()
-	}
-	return &NetBird{
-		mgmtAddr:       mgmtAddr,
-		proxyID:        proxyID,
-		logger:         logger,
-		clients:        make(map[types.AccountID]*clientEntry),
-		statusNotifier: notifier,
-	}
+// ClientDebugInfo contains debug information about a client.
+type ClientDebugInfo struct {
+	AccountID   types.AccountID
+	DomainCount int
+	Domains     domain.List
+	HasClient   bool
+	CreatedAt   time.Time
 }
 
+// accountIDContextKey is the context key for storing the account ID.
+type accountIDContextKey struct{}
+
 // AddPeer registers a domain for an account. If the account doesn't have a client yet,
 // one is created using the provided setup key. Multiple domains can share the same client.
 func (n *NetBird) AddPeer(ctx context.Context, accountID types.AccountID, d domain.Domain, key, reverseProxyID string) error {
@@ -379,15 +377,6 @@ func (n *NetBird) GetClient(accountID types.AccountID) (*embed.Client, bool) {
 	return entry.client, true
 }
 
-// ClientDebugInfo contains debug information about a client.
-type ClientDebugInfo struct {
-	AccountID   types.AccountID
-	DomainCount int
-	Domains     domain.List
-	HasClient   bool
-	CreatedAt   time.Time
-}
-
 // ListClientsForDebug returns information about all clients for debug purposes.
 func (n *NetBird) ListClientsForDebug() map[types.AccountID]ClientDebugInfo {
 	n.clientsMux.RLock()
@@ -410,8 +399,33 @@ func (n *NetBird) ListClientsForDebug() map[types.AccountID]ClientDebugInfo {
 	return result
 }
 
-// accountIDContextKey is the context key for storing the account ID.
-type accountIDContextKey struct{}
+// ListClientsForStartup returns all embed.Client instances for health checks.
+func (n *NetBird) ListClientsForStartup() map[types.AccountID]*embed.Client {
+	n.clientsMux.RLock()
+	defer n.clientsMux.RUnlock()
+
+	result := make(map[types.AccountID]*embed.Client)
+	for accountID, entry := range n.clients {
+		if entry.client != nil {
+			result[accountID] = entry.client
+		}
+	}
+	return result
+}
+
+// NewNetBird creates a new NetBird transport.
+func NewNetBird(mgmtAddr, proxyID string, logger *log.Logger, notifier statusNotifier) *NetBird {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &NetBird{
+		mgmtAddr:       mgmtAddr,
+		proxyID:        proxyID,
+		logger:         logger,
+		clients:        make(map[types.AccountID]*clientEntry),
+		statusNotifier: notifier,
+	}
+}
 
 // WithAccountID adds the account ID to the context.
 func WithAccountID(ctx context.Context, accountID types.AccountID) context.Context {