From b3ab6317fe58d565a71e8dde09abbcdd5b8a8e50 Mon Sep 17 00:00:00 2001
From: tiennm99 <tiennm99@outlook.com>
Date: Fri, 26 Jun 2026 15:48:35 +0700
Subject: [PATCH] fix(health): report readiness instead of static 200

The health endpoint returned 200 unconditionally from before dependencies were ready and never reflected MongoDB state. Handler now takes a readiness check; the endpoint returns 503 while starting or when MongoDB is unreachable and 200 only when ready.
---
 cmd/openai-status-bot/main.go  | 19 +++++++++++++++++--
 internal/health/server.go      | 24 ++++++++++++++++++++----
 internal/health/server_test.go | 28 ++++++++++++++++++++++++++--
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/cmd/openai-status-bot/main.go b/cmd/openai-status-bot/main.go
index 12a4c1c..6ad552f 100644
--- a/cmd/openai-status-bot/main.go
+++ b/cmd/openai-status-bot/main.go
@@ -2,9 +2,11 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log/slog"
 	"os"
 	"os/signal"
+	"sync/atomic"
 	"syscall"
 
 	"github.com/tiennm99/openai-status-bot/internal/bot"
@@ -33,9 +35,21 @@ func main() {
 	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
 
-	go health.Run(ctx, logger)
+	// The health endpoint comes up immediately but reports 503 until startup
+	// finishes, then pings MongoDB on each probe so a dependency outage is
+	// reflected instead of a static 200.
+	var (
+		ready       atomic.Bool
+		mongoClient *mongo.Client
+	)
+	go health.Run(ctx, logger, func(ctx context.Context) error {
+		if !ready.Load() {
+			return errors.New("bot is still starting")
+		}
+		return mongoClient.Ping(ctx, nil)
+	})
 
-	mongoClient, err := mongo.Connect(options.Client().ApplyURI(cfg.MongoURI))
+	mongoClient, err = mongo.Connect(options.Client().ApplyURI(cfg.MongoURI))
 	if err != nil {
 		logger.Error("connect mongodb", "database", cfg.MongoDatabase, "error", err)
 		os.Exit(1)
@@ -77,6 +91,7 @@ func main() {
 
 	go statusPoller.Run(ctx)
 
+	ready.Store(true)
 	logger.Info("openai status bot started", "poll_interval", cfg.PollInterval.String())
 	if err := commandBot.Run(ctx); err != nil && ctx.Err() == nil {
 		logger.Error("telegram bot stopped", "error", err)
diff --git a/internal/health/server.go b/internal/health/server.go
index 4d049b3..ba92a94 100644
--- a/internal/health/server.go
+++ b/internal/health/server.go
@@ -13,19 +13,35 @@ const (
 	Path    = "/healthz"
 )
 
-func Handler() http.Handler {
+// Check reports whether the bot is ready to serve. It returns nil when healthy
+// and an error describing why not otherwise. A nil Check is always healthy.
+type Check func(ctx context.Context) error
+
+// Handler serves the health endpoint. It returns 200 only when check passes,
+// and 503 while the bot is still starting or a dependency is unreachable, so an
+// orchestrator probe reflects real readiness instead of a static 200.
+func Handler(check Check) http.Handler {
 	mux := http.NewServeMux()
-	mux.HandleFunc(Path, func(w http.ResponseWriter, _ *http.Request) {
+	mux.HandleFunc(Path, func(w http.ResponseWriter, r *http.Request) {
+		if check != nil {
+			ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
+			defer cancel()
+			if err := check(ctx); err != nil {
+				w.WriteHeader(http.StatusServiceUnavailable)
+				_, _ = w.Write([]byte("unavailable\n"))
+				return
+			}
+		}
 		w.WriteHeader(http.StatusOK)
 		_, _ = w.Write([]byte("ok\n"))
 	})
 	return mux
 }
 
-func Run(ctx context.Context, logger *slog.Logger) {
+func Run(ctx context.Context, logger *slog.Logger, check Check) {
 	server := &http.Server{
 		Addr:              Address,
-		Handler:           Handler(),
+		Handler:           Handler(check),
 		ReadHeaderTimeout: 3 * time.Second,
 	}
 
diff --git a/internal/health/server_test.go b/internal/health/server_test.go
index bbd0c22..d1cb7f4 100644
--- a/internal/health/server_test.go
+++ b/internal/health/server_test.go
@@ -1,16 +1,18 @@
 package health
 
 import (
+	"context"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 )
 
-func TestHandlerReturnsOK(t *testing.T) {
+func TestHandlerReturnsOKWhenCheckPasses(t *testing.T) {
 	req := httptest.NewRequest(http.MethodGet, Path, nil)
 	res := httptest.NewRecorder()
 
-	Handler().ServeHTTP(res, req)
+	Handler(func(context.Context) error { return nil }).ServeHTTP(res, req)
 
 	if res.Code != http.StatusOK {
 		t.Fatalf("status = %d, want %d", res.Code, http.StatusOK)
@@ -19,3 +21,25 @@ func TestHandlerReturnsOK(t *testing.T) {
 		t.Fatalf("body = %q, want ok", res.Body.String())
 	}
 }
+
+func TestHandlerNilCheckIsHealthy(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, Path, nil)
+	res := httptest.NewRecorder()
+
+	Handler(nil).ServeHTTP(res, req)
+
+	if res.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", res.Code, http.StatusOK)
+	}
+}
+
+func TestHandlerReturnsUnavailableWhenCheckFails(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, Path, nil)
+	res := httptest.NewRecorder()
+
+	Handler(func(context.Context) error { return errors.New("not ready") }).ServeHTTP(res, req)
+
+	if res.Code != http.StatusServiceUnavailable {
+		t.Fatalf("status = %d, want %d", res.Code, http.StatusServiceUnavailable)
+	}
+}