Health Aggregator Service

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

// --- Health Check Types ---

type Status string

const (
	StatusUp       Status = "up"
	StatusDown     Status = "down"
	StatusDegraded Status = "degraded"
)

type CheckResult struct {
	Name     string        `json:"name"`
	Status   Status        `json:"status"`
	Duration time.Duration `json:"duration_ms"`
	Error    string        `json:"error,omitempty"`
}

type HealthResponse struct {
	Status  Status        `json:"status"`
	Checks  []CheckResult `json:"checks,omitempty"`
	Version string        `json:"version,omitempty"`
}

type Check struct {
	Name    string
	Fn      func() error
	Timeout time.Duration
}

type Checker struct {
	checks []Check
	mu     sync.RWMutex
}

func NewChecker() *Checker {
	return &Checker{}
}

func (c *Checker) AddCheck(name string, fn func() error, timeout time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.checks = append(c.checks, Check{Name: name, Fn: fn, Timeout: timeout})
}

func (c *Checker) RunChecks() HealthResponse {
	c.mu.RLock()
	checks := make([]Check, len(c.checks))
	copy(checks, c.checks)
	c.mu.RUnlock()

	results := make([]CheckResult, len(checks))
	var wg sync.WaitGroup

	for i, check := range checks {
		wg.Add(1)
		go func(idx int, chk Check) {
			defer wg.Done()
			start := time.Now()

			errCh := make(chan error, 1)
			go func() { errCh <- chk.Fn() }()

			select {
			case err := <-errCh:
				result := CheckResult{
					Name:     chk.Name,
					Duration: time.Since(start),
					Status:   StatusUp,
				}
				if err != nil {
					result.Status = StatusDown
					result.Error = err.Error()
				}
				results[idx] = result
			case <-time.After(chk.Timeout):
				results[idx] = CheckResult{
					Name:     chk.Name,
					Status:   StatusDown,
					Duration: chk.Timeout,
					Error:    "timeout",
				}
			}
		}(i, check)
	}

	wg.Wait()

	overall := StatusUp
	for _, r := range results {
		if r.Status == StatusDown {
			overall = StatusDown
			break
		}
	}

	return HealthResponse{Status: overall, Checks: results}
}

func (c *Checker) RunChecksWithDegradation(critical []string) HealthResponse {
	resp := c.RunChecks()

	criticalSet := make(map[string]bool)
	for _, name := range critical {
		criticalSet[name] = true
	}

	overall := StatusUp
	for _, check := range resp.Checks {
		if check.Status == StatusDown {
			if criticalSet[check.Name] {
				overall = StatusDown
				break
			}
			overall = StatusDegraded
		}
	}
	resp.Status = overall
	return resp
}

// --- Dependency Checks ---

func DatabaseCheck(db *sql.DB) func() error {
	return func() error {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		return db.PingContext(ctx)
	}
}

func RedisCheck(client *redis.Client) func() error {
	return func() error {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		return client.Ping(ctx).Err()
	}
}

func HTTPCheck(url string, expectedStatus int) func() error {
	client := &http.Client{Timeout: 5 * time.Second}
	return func() error {
		resp, err := client.Get(url)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if resp.StatusCode != expectedStatus {
			return fmt.Errorf("expected %d, got %d", expectedStatus, resp.StatusCode)
		}
		return nil
	}
}

// --- Circuit Breaker ---

type State int

const (
	Closed   State = iota // Normal operation: requests pass through
	Open                  // Failing: requests blocked
	HalfOpen              // Testing: one request allowed through
)

var ErrCircuitOpen = errors.New("circuit breaker is open")

type CircuitBreaker struct {
	mu           sync.Mutex
	state        State
	failures     int
	threshold    int
	lastFailure  time.Time
	resetTimeout time.Duration
}

func NewCircuitBreaker(threshold int, resetTimeout time.Duration) *CircuitBreaker {
	return &CircuitBreaker{
		state:        Closed,
		threshold:    threshold,
		resetTimeout: resetTimeout,
	}
}

func (cb *CircuitBreaker) Execute(fn func() error) error {
	cb.mu.Lock()

	switch cb.state {
	case Open:
		if time.Since(cb.lastFailure) > cb.resetTimeout {
			cb.state = HalfOpen
			cb.mu.Unlock()
			return cb.tryExec(fn)
		}
		cb.mu.Unlock()
		return ErrCircuitOpen

	case HalfOpen:
		cb.mu.Unlock()
		return cb.tryExec(fn)

	default: // Closed
		cb.mu.Unlock()
		return cb.tryExec(fn)
	}
}

func (cb *CircuitBreaker) tryExec(fn func() error) error {
	err := fn()

	cb.mu.Lock()
	defer cb.mu.Unlock()

	if err != nil {
		cb.failures++
		cb.lastFailure = time.Now()
		if cb.failures >= cb.threshold {
			cb.state = Open
		}
		return err
	}

	// Success: reset
	cb.failures = 0
	cb.state = Closed
	return nil
}

// --- Health Aggregation ---

type ServiceHealth struct {
	Name   string `json:"name"`
	URL    string `json:"-"`
	Status Status `json:"status"`
	Error  string `json:"error,omitempty"`
}

func AggregateHealth(services []ServiceHealth) Status {
	downCount := 0
	for _, svc := range services {
		if svc.Status == StatusDown {
			downCount++
		}
	}

	switch {
	case downCount == 0:
		return StatusUp
	case downCount == len(services):
		return StatusDown
	default:
		return StatusDegraded
	}
}

// --- Main: Wire It All Together ---

func main() {
	checker := NewChecker()

	// Example: wrap database check with circuit breaker
	dbBreaker := NewCircuitBreaker(3, 30*time.Second)
	checker.AddCheck("postgres", func() error {
		return dbBreaker.Execute(func() error {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			// In production: return db.PingContext(ctx)
			_ = ctx
			return nil
		})
	}, 3*time.Second)

	mux := http.NewServeMux()

	// Liveness: is the process alive?
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"status":"up"}`))
	})

	// Readiness: can we serve traffic?
	mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
		resp := checker.RunChecksWithDegradation([]string{"postgres"})
		status := http.StatusOK
		if resp.Status == StatusDown {
			status = http.StatusServiceUnavailable
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(status)
		json.NewEncoder(w).Encode(resp)
	})

	// Aggregation endpoint for multi-service monitoring
	mux.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		services := []ServiceHealth{
			{Name: "api", Status: StatusUp},
			{Name: "worker", Status: StatusUp},
			{Name: "scheduler", Status: StatusDegraded},
		}
		overall := AggregateHealth(services)
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"status":   overall,
			"services": services,
		})
	})

	fmt.Println("Health service listening on :8080")
	http.ListenAndServe(":8080", mux)
}