From c0253c98f9e9ef8f7dee0f76d0554edb3d8d0b4b Mon Sep 17 00:00:00 2001 From: liangxu Date: Fri, 24 Apr 2026 22:20:43 +0800 Subject: [PATCH] feat(scheduler): expose metrics HTTP and add graceful shutdown Scheduler now runs a token-protected HTTP server on scheduler.http_port (default 8081, -1 disables) that exposes Prometheus /metrics and the daily-task JSON snapshot endpoint. Workers gain a blocking Run(ctx) and are supervised by a WaitGroup: on SIGINT/SIGTERM the metrics server shuts down first, worker ticks stop scheduling but ongoing runOnce calls keep their own context and finish, and the process waits up to 60s before exiting. Adds scheduler.http_host/http_port/internal_metrics_token to config with SCHEDULER_* env overrides and exposes port 8081 in the Docker image. Co-Authored-By: Claude Opus 4.7 (1M context) --- deploy/config.yaml | 3 + server/Dockerfile | 2 +- server/cmd/scheduler/main.go | 198 +++++++++++++++++- server/cmd/scheduler/main_test.go | 95 +++++++++ server/configs/config.yaml | 3 + .../knowledge_deleted_cleanup_worker.go | 6 +- .../monitoring_lease_recovery_worker.go | 10 +- .../monitoring_received_inspection_worker.go | 10 +- .../monitoring_result_recovery_worker.go | 10 +- .../scheduler/schedule_dispatch_worker.go | 10 +- 10 files changed, 328 insertions(+), 19 deletions(-) create mode 100644 server/cmd/scheduler/main_test.go diff --git a/deploy/config.yaml b/deploy/config.yaml index 9341cb5..14d5615 100644 --- a/deploy/config.yaml +++ b/deploy/config.yaml @@ -28,6 +28,9 @@ rabbitmq: publish_channel_pool_size: 32 scheduler: + http_host: 127.0.0.1 + http_port: 8081 + internal_metrics_token: "" dispatch_interval: 15s dispatch_timeout: 30s dispatch_batch_size: 100 diff --git a/server/Dockerfile b/server/Dockerfile index 543f417..e5cd043 100644 --- a/server/Dockerfile +++ b/server/Dockerfile @@ -37,5 +37,5 @@ WORKDIR /app COPY --from=builder /bin/service /app/service COPY configs/ /app/configs/ -EXPOSE 8080 +EXPOSE 8080 8081 ENTRYPOINT ["/app/service"] diff --git a/server/cmd/scheduler/main.go b/server/cmd/scheduler/main.go index e990a3a..06b80b0 100644 --- a/server/cmd/scheduler/main.go +++ b/server/cmd/scheduler/main.go @@ -2,16 +2,28 @@ package main import ( "context" + "crypto/subtle" + "errors" "log" + "net" + "net/http" "os" "os/signal" + "strconv" + "strings" + "sync" "syscall" + "time" "github.com/geo-platform/tenant-api/internal/bootstrap" internalscheduler "github.com/geo-platform/tenant-api/internal/scheduler" + "github.com/geo-platform/tenant-api/internal/shared/response" tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" + "github.com/gin-gonic/gin" ) +const schedulerWorkerShutdownGracePeriod = 60 * time.Second + func main() { configPath := "configs/config.yaml" if p := os.Getenv("CONFIG_PATH"); p != "" { @@ -57,21 +69,197 @@ func main() { app.LLM, app.Logger, ) + monitoringService := tenantapp.NewMonitoringService( + app.DB, + app.MonitoringDB, + app.RabbitMQ, + app.Config.MonitoringDispatch, + app.Config.BrandLibrary, + app.Logger, + ) ctx, cancel := context.WithCancel(context.Background()) defer cancel() - internalscheduler.NewScheduleDispatchWorker(app.DB, app.RabbitMQ, promptRuleSvc, app.Cache, app.Logger, app.Config.Scheduler).Start(ctx) - internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger).Start(ctx) - internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, app.Logger).Start(ctx) - internalscheduler.NewMonitoringReceivedInspectionWorker(app.MonitoringDB, app.Logger).Start(ctx) - internalscheduler.NewKnowledgeDeletedCleanupWorker(knowledgeSvc).Start(ctx) + var workerWG sync.WaitGroup + scheduleDispatchWorker := internalscheduler.NewScheduleDispatchWorker(app.DB, app.RabbitMQ, promptRuleSvc, app.Cache, app.Logger, app.Config.Scheduler) + monitoringDailyTaskWorker := tenantapp.NewMonitoringDailyTaskWorker(monitoringService, monitoringCallbackService, app.Logger) + monitoringResultRecoveryWorker := internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger) + monitoringLeaseRecoveryWorker := internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, app.Logger) + monitoringReceivedInspectionWorker := internalscheduler.NewMonitoringReceivedInspectionWorker(app.MonitoringDB, app.Logger) + knowledgeDeletedCleanupWorker := internalscheduler.NewKnowledgeDeletedCleanupWorker(knowledgeSvc) + + startSchedulerWorker(ctx, &workerWG, "schedule_dispatch", app.Logger.Sugar(), scheduleDispatchWorker.Run) + startSchedulerWorker(ctx, &workerWG, "monitoring_daily_task", app.Logger.Sugar(), monitoringDailyTaskWorker.Run) + startSchedulerWorker(ctx, &workerWG, "monitoring_result_recovery", app.Logger.Sugar(), monitoringResultRecoveryWorker.Run) + startSchedulerWorker(ctx, &workerWG, "monitoring_lease_recovery", app.Logger.Sugar(), monitoringLeaseRecoveryWorker.Run) + startSchedulerWorker(ctx, &workerWG, "monitoring_received_inspection", app.Logger.Sugar(), monitoringReceivedInspectionWorker.Run) + startSchedulerWorker(ctx, &workerWG, "knowledge_deleted_cleanup", app.Logger.Sugar(), knowledgeDeletedCleanupWorker.Run) + + var metricsServer *http.Server + if app.Config.Scheduler.HTTPPort > 0 { + metricsServer = startSchedulerMetricsServer(app) + } else { + app.Logger.Sugar().Infof("scheduler metrics http server disabled by http_port=%d", app.Config.Scheduler.HTTPPort) + } app.Logger.Info("scheduler started") quit := make(chan os.Signal, 1) signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) <-quit + app.Logger.Info("scheduler shutdown signal received") + shutdownSchedulerMetricsServer(metricsServer, app.Logger.Sugar()) cancel() + waitSchedulerWorkers(&workerWG, schedulerWorkerShutdownGracePeriod, app.Logger.Sugar()) app.Logger.Info("scheduler shutting down...") } + +func startSchedulerWorker( + ctx context.Context, + wg *sync.WaitGroup, + name string, + logger interface { + Infof(template string, args ...interface{}) + Errorf(template string, args ...interface{}) + }, + run func(context.Context), +) { + if wg == nil || run == nil { + return + } + + wg.Add(1) + go func() { + defer wg.Done() + logger.Infof("scheduler worker %s started", name) + run(ctx) + logger.Infof("scheduler worker %s stopped", name) + }() +} + +func waitSchedulerWorkers(wg *sync.WaitGroup, timeout time.Duration, logger interface { + Infof(template string, args ...interface{}) + Errorf(template string, args ...interface{}) +}) bool { + if wg == nil { + return true + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + logger.Infof("scheduler workers stopped") + return true + case <-time.After(timeout): + logger.Errorf("scheduler worker shutdown timed out after %s", timeout) + return false + } +} + +func startSchedulerMetricsServer(app *bootstrap.App) *http.Server { + token := strings.TrimSpace(app.Config.Scheduler.InternalMetricsToken) + if token == "" { + app.Logger.Warn("scheduler metrics http server not started: scheduler.internal_metrics_token is required") + return nil + } + + authMiddleware := schedulerMetricsAuthMiddleware(token) + app.Engine.GET("/api/internal/metrics/monitoring/daily-tasks", authMiddleware, func(c *gin.Context) { + response.Success(c, tenantapp.MonitoringDailyTaskMetricsSnapshotValue()) + }) + prometheusHandler := tenantapp.MonitoringSchedulerMetricsPrometheusHandler() + app.Engine.GET("/metrics", authMiddleware, func(c *gin.Context) { + prometheusHandler.ServeHTTP(c.Writer, c.Request) + }) + + addr := schedulerHTTPAddr(app.Config.Scheduler.HTTPHost, app.Config.Scheduler.HTTPPort) + server := &http.Server{ + Addr: addr, + Handler: app.Engine, + ReadHeaderTimeout: 5 * time.Second, + } + + go func() { + app.Logger.Sugar().Infof("scheduler metrics http listening on %s", addr) + if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + app.Logger.Sugar().Errorf("scheduler metrics http server stopped: %v", err) + } + }() + return server +} + +func shutdownSchedulerMetricsServer(server *http.Server, logger interface { + Infof(template string, args ...interface{}) + Errorf(template string, args ...interface{}) +}) { + if server == nil { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := server.Shutdown(ctx); err != nil { + logger.Errorf("scheduler metrics http shutdown failed: %v", err) + return + } + logger.Infof("scheduler metrics http server stopped") +} + +func schedulerHTTPAddr(host string, port int) string { + host = strings.TrimSpace(host) + if host == "" { + host = "127.0.0.1" + } + return net.JoinHostPort(host, strconv.Itoa(port)) +} + +func schedulerMetricsAuthMiddleware(token string) gin.HandlerFunc { + expected := strings.TrimSpace(token) + return func(c *gin.Context) { + if expected == "" || !schedulerMetricsTokenMatches(c, expected) { + response.Error(c, response.ErrUnauthorized(40181, "internal_metrics_unauthorized", "valid internal metrics token is required")) + c.Abort() + return + } + c.Next() + } +} + +func schedulerMetricsTokenMatches(c *gin.Context, expected string) bool { + if c == nil { + return false + } + + candidates := []string{ + c.GetHeader("X-Internal-Metrics-Token"), + bearerToken(c.GetHeader("Authorization")), + } + for _, candidate := range candidates { + candidate = strings.TrimSpace(candidate) + if candidate == "" || len(candidate) != len(expected) { + continue + } + if subtle.ConstantTimeCompare([]byte(candidate), []byte(expected)) == 1 { + return true + } + } + return false +} + +func bearerToken(header string) string { + header = strings.TrimSpace(header) + if header == "" { + return "" + } + parts := strings.Fields(header) + if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { + return "" + } + return parts[1] +} diff --git a/server/cmd/scheduler/main_test.go b/server/cmd/scheduler/main_test.go new file mode 100644 index 0000000..71ffc43 --- /dev/null +++ b/server/cmd/scheduler/main_test.go @@ -0,0 +1,95 @@ +package main + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestSchedulerHTTPAddrDefaultsToLoopback(t *testing.T) { + if got := schedulerHTTPAddr("", 8081); got != "127.0.0.1:8081" { + t.Fatalf("schedulerHTTPAddr() = %q, want 127.0.0.1:8081", got) + } + if got := schedulerHTTPAddr("0.0.0.0", 8081); got != "0.0.0.0:8081" { + t.Fatalf("schedulerHTTPAddr() = %q, want 0.0.0.0:8081", got) + } +} + +func TestSchedulerMetricsAuthMiddleware(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + headerName string + header string + wantStatus int + }{ + {name: "missing token", wantStatus: http.StatusUnauthorized}, + {name: "bad bearer token", headerName: "Authorization", header: "Bearer bad", wantStatus: http.StatusUnauthorized}, + {name: "valid bearer token", headerName: "Authorization", header: "Bearer secret-token", wantStatus: http.StatusNoContent}, + {name: "valid internal token header", headerName: "X-Internal-Metrics-Token", header: "secret-token", wantStatus: http.StatusNoContent}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + router := gin.New() + router.GET("/metrics", schedulerMetricsAuthMiddleware("secret-token"), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + req := httptest.NewRequest(http.MethodGet, "/metrics", nil) + if tt.headerName != "" { + req.Header.Set(tt.headerName, tt.header) + } + rec := httptest.NewRecorder() + + router.ServeHTTP(rec, req) + if rec.Code != tt.wantStatus { + t.Fatalf("status = %d, want %d", rec.Code, tt.wantStatus) + } + }) + } +} + +func TestStartSchedulerWorkerWaitsForRunToReturn(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var wg sync.WaitGroup + started := make(chan struct{}) + startSchedulerWorker(ctx, &wg, "test_worker", testSchedulerLogger{}, func(ctx context.Context) { + close(started) + <-ctx.Done() + }) + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("worker did not start") + } + + cancel() + if ok := waitSchedulerWorkers(&wg, time.Second, testSchedulerLogger{}); !ok { + t.Fatal("expected worker shutdown wait to complete") + } +} + +func TestWaitSchedulerWorkersTimesOut(t *testing.T) { + var wg sync.WaitGroup + wg.Add(1) + defer wg.Done() + + if ok := waitSchedulerWorkers(&wg, time.Nanosecond, testSchedulerLogger{}); ok { + t.Fatal("expected worker shutdown wait to time out") + } +} + +type testSchedulerLogger struct{} + +func (testSchedulerLogger) Infof(string, ...interface{}) {} +func (testSchedulerLogger) Errorf(string, ...interface{}) {} diff --git a/server/configs/config.yaml b/server/configs/config.yaml index 991fe54..b2ed809 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -63,6 +63,9 @@ rabbitmq: publish_channel_pool_size: 32 scheduler: + http_host: 127.0.0.1 + http_port: 8081 + internal_metrics_token: dev_scheduler_metrics_token dispatch_interval: 15s dispatch_timeout: 30s dispatch_batch_size: 100 diff --git a/server/internal/scheduler/knowledge_deleted_cleanup_worker.go b/server/internal/scheduler/knowledge_deleted_cleanup_worker.go index 60ef9c0..3965053 100644 --- a/server/internal/scheduler/knowledge_deleted_cleanup_worker.go +++ b/server/internal/scheduler/knowledge_deleted_cleanup_worker.go @@ -20,10 +20,14 @@ func NewKnowledgeDeletedCleanupWorker(service *tenantapp.KnowledgeService) *Know } func (w *KnowledgeDeletedCleanupWorker) Start(ctx context.Context) { + go w.Run(ctx) +} + +func (w *KnowledgeDeletedCleanupWorker) Run(ctx context.Context) { if w == nil || w.service == nil { return } - go w.run(ctx) + w.run(ctx) } func (w *KnowledgeDeletedCleanupWorker) run(ctx context.Context) { diff --git a/server/internal/scheduler/monitoring_lease_recovery_worker.go b/server/internal/scheduler/monitoring_lease_recovery_worker.go index 65f38ec..342bfa9 100644 --- a/server/internal/scheduler/monitoring_lease_recovery_worker.go +++ b/server/internal/scheduler/monitoring_lease_recovery_worker.go @@ -32,14 +32,18 @@ func NewMonitoringLeaseRecoveryWorker(monitoringPool *pgxpool.Pool, logger *zap. } func (w *MonitoringLeaseRecoveryWorker) Start(ctx context.Context) { + go w.Run(ctx) +} + +func (w *MonitoringLeaseRecoveryWorker) Run(ctx context.Context) { if w == nil || w.monitoringPool == nil { return } - go w.run(ctx) + w.run(ctx) } func (w *MonitoringLeaseRecoveryWorker) run(ctx context.Context) { - w.runOnce(ctx) + w.runOnce(context.Background()) ticker := time.NewTicker(w.interval) defer ticker.Stop() @@ -49,7 +53,7 @@ func (w *MonitoringLeaseRecoveryWorker) run(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - w.runOnce(ctx) + w.runOnce(context.Background()) } } } diff --git a/server/internal/scheduler/monitoring_received_inspection_worker.go b/server/internal/scheduler/monitoring_received_inspection_worker.go index 994695c..18481fb 100644 --- a/server/internal/scheduler/monitoring_received_inspection_worker.go +++ b/server/internal/scheduler/monitoring_received_inspection_worker.go @@ -38,14 +38,18 @@ func NewMonitoringReceivedInspectionWorker(monitoringPool *pgxpool.Pool, logger } func (w *MonitoringReceivedInspectionWorker) Start(ctx context.Context) { + go w.Run(ctx) +} + +func (w *MonitoringReceivedInspectionWorker) Run(ctx context.Context) { if w == nil || w.monitoringPool == nil { return } - go w.run(ctx) + w.run(ctx) } func (w *MonitoringReceivedInspectionWorker) run(ctx context.Context) { - w.runOnce(ctx) + w.runOnce(context.Background()) ticker := time.NewTicker(w.interval) defer ticker.Stop() @@ -55,7 +59,7 @@ func (w *MonitoringReceivedInspectionWorker) run(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - w.runOnce(ctx) + w.runOnce(context.Background()) } } } diff --git a/server/internal/scheduler/monitoring_result_recovery_worker.go b/server/internal/scheduler/monitoring_result_recovery_worker.go index c4930a2..8f03efe 100644 --- a/server/internal/scheduler/monitoring_result_recovery_worker.go +++ b/server/internal/scheduler/monitoring_result_recovery_worker.go @@ -48,14 +48,18 @@ func NewMonitoringResultRecoveryWorker( } func (w *MonitoringResultRecoveryWorker) Start(ctx context.Context) { + go w.Run(ctx) +} + +func (w *MonitoringResultRecoveryWorker) Run(ctx context.Context) { if w == nil || w.monitoringPool == nil || w.service == nil { return } - go w.run(ctx) + w.run(ctx) } func (w *MonitoringResultRecoveryWorker) run(ctx context.Context) { - w.runOnce(ctx) + w.runOnce(context.Background()) ticker := time.NewTicker(w.interval) defer ticker.Stop() @@ -65,7 +69,7 @@ func (w *MonitoringResultRecoveryWorker) run(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - w.runOnce(ctx) + w.runOnce(context.Background()) } } } diff --git a/server/internal/scheduler/schedule_dispatch_worker.go b/server/internal/scheduler/schedule_dispatch_worker.go index 0a88316..96a0095 100644 --- a/server/internal/scheduler/schedule_dispatch_worker.go +++ b/server/internal/scheduler/schedule_dispatch_worker.go @@ -80,14 +80,18 @@ func NewScheduleDispatchWorker( } func (w *ScheduleDispatchWorker) Start(ctx context.Context) { + go w.Run(ctx) +} + +func (w *ScheduleDispatchWorker) Run(ctx context.Context) { if w == nil || w.pool == nil || w.promptRuleService == nil { return } - go w.run(ctx) + w.run(ctx) } func (w *ScheduleDispatchWorker) run(ctx context.Context) { - w.runOnce(ctx) + w.runOnce(context.Background()) ticker := time.NewTicker(w.interval) defer ticker.Stop() @@ -97,7 +101,7 @@ func (w *ScheduleDispatchWorker) run(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - w.runOnce(ctx) + w.runOnce(context.Background()) } } }