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()) } } }