Add worker health endpoints

cmd/worker/main.go

@@ -2,10 +2,16 @@ package main
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
 	"log/slog"
+	"net/http"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
+	"time"
 
 	"ai-service/internal/config"
 	"ai-service/internal/llm"
@@ -44,6 +50,7 @@ func main() {
 	llmClient := llm.New(cfg.LLMBaseURL, cfg.LLMAPIKey, cfg.LLMModel, cfg.LLMTimeout)
 	transcriber := transcription.New(cfg.WhisperXURL, cfg.WhisperXTimeout, cfg.FfmpegPath, cfg.WhisperXLeadSilence)
 	w := worker.New(db, llmClient, transcriber, cfg.WorkerID, cfg.LLMModel, cfg.WorkerTaskTypes, cfg.WorkerModelProfiles, cfg.WorkerPollInterval, cfg.WorkerLeaseTimeout, cfg.WorkerClaimLimit)
+	healthSrv := startHealthServer(ctx, db, cfg)
 
 	slog.Info("ai_worker_started",
 		"worker_id", cfg.WorkerID,
@@ -57,4 +64,82 @@ func main() {
 		"claim_limit", cfg.WorkerClaimLimit,
 	)
 	w.Run(ctx)
+
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_ = healthSrv.Shutdown(shutdownCtx)
+}
+
+type workerHealth struct {
+	store *store.Store
+	cfg   config.Config
+}
+
+func startHealthServer(ctx context.Context, db *store.Store, cfg config.Config) *http.Server {
+	srv := &http.Server{
+		Addr:              fmt.Sprintf("%s:%d", cfg.WorkerHTTPHost, cfg.WorkerHTTPPort),
+		Handler:           workerHealth{store: db, cfg: cfg},
+		ReadHeaderTimeout: 5 * time.Second,
+	}
+	go func() {
+		slog.Info("ai_worker_health_started", "addr", srv.Addr)
+		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			slog.Error("ai_worker_health_failed", "error", err)
+		}
+	}()
+	go func() {
+		<-ctx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = srv.Shutdown(shutdownCtx)
+	}()
+	return srv
+}
+
+func (h workerHealth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	path := strings.TrimSuffix(r.URL.Path, "/")
+	if path == "" {
+		path = "/"
+	}
+	switch {
+	case r.Method == http.MethodGet && path == "/healthz":
+		writeWorkerJSON(w, http.StatusOK, map[string]any{
+			"status":    "ok",
+			"worker_id": h.cfg.WorkerID,
+		})
+	case r.Method == http.MethodGet && path == "/readyz":
+		ctx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
+		defer cancel()
+		if err := h.store.Ping(ctx); err != nil {
+			writeWorkerJSON(w, http.StatusServiceUnavailable, map[string]any{
+				"status": "not_ready",
+				"error":  err.Error(),
+			})
+			return
+		}
+		writeWorkerJSON(w, http.StatusOK, map[string]any{
+			"status":    "ready",
+			"worker_id": h.cfg.WorkerID,
+		})
+	case r.Method == http.MethodGet && path == "/worker/status":
+		writeWorkerJSON(w, http.StatusOK, map[string]any{
+			"status":         "running",
+			"worker_id":      h.cfg.WorkerID,
+			"task_types":     h.cfg.WorkerTaskTypes,
+			"model_profiles": h.cfg.WorkerModelProfiles,
+			"claim_limit":    h.cfg.WorkerClaimLimit,
+			"poll_interval":  h.cfg.WorkerPollInterval.String(),
+			"lease_timeout":  h.cfg.WorkerLeaseTimeout.String(),
+		})
+	default:
+		writeWorkerJSON(w, http.StatusNotFound, map[string]any{"error": "not found"})
+	}
+}
+
+func writeWorkerJSON(w http.ResponseWriter, code int, body any) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(code)
+	if err := json.NewEncoder(w).Encode(body); err != nil {
+		slog.Warn("worker_health_write_failed", "error", err)
+	}
 }