288 lines
8.5 KiB
Go
288 lines
8.5 KiB
Go
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"
|
|
generateworker "github.com/geo-platform/tenant-api/internal/worker/generate"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
const schedulerWorkerShutdownGracePeriod = 60 * time.Second
|
|
|
|
func main() {
|
|
configPath := "configs/config.yaml"
|
|
if p := os.Getenv("CONFIG_PATH"); p != "" {
|
|
configPath = p
|
|
}
|
|
|
|
app, err := bootstrap.New(configPath)
|
|
if err != nil {
|
|
log.Fatalf("bootstrap: %v", err)
|
|
}
|
|
defer app.Close()
|
|
|
|
generationProvider := tenantapp.NewGenerationStreamDisabledConfigProvider(app.ConfigStore)
|
|
generationCfg := app.Config().Generation
|
|
generationCfg.StreamEnabled = false
|
|
|
|
knowledgeSvc := tenantapp.NewKnowledgeService(
|
|
app.DB,
|
|
app.RetrievalProvider,
|
|
app.VectorStore,
|
|
app.ObjectStorage,
|
|
app.LLM,
|
|
app.Logger,
|
|
app.Config().Retrieval,
|
|
app.Config().LLM,
|
|
0,
|
|
0,
|
|
).WithConfigProvider(app.ConfigStore)
|
|
|
|
promptRuleSvc := tenantapp.NewPromptRuleGenerationService(
|
|
app.DB,
|
|
app.LLM,
|
|
app.RabbitMQ,
|
|
knowledgeSvc,
|
|
app.GenerationStreams,
|
|
generationCfg,
|
|
app.Config().LLM.MaxOutputTokens,
|
|
).WithCache(app.Cache).WithLogger(app.Logger).WithConfigProvider(generationProvider)
|
|
|
|
monitoringCallbackService := tenantapp.NewMonitoringCallbackService(
|
|
app.DB,
|
|
app.MonitoringDB,
|
|
app.RabbitMQ,
|
|
app.LLM,
|
|
app.Logger,
|
|
)
|
|
monitoringService := app.MonitoringService
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
app.StartConfigWatcher(ctx)
|
|
|
|
var workerWG sync.WaitGroup
|
|
scheduleDispatchWorker := internalscheduler.NewScheduleDispatchWorker(
|
|
app.DB,
|
|
app.RabbitMQ,
|
|
promptRuleSvc,
|
|
app.Cache,
|
|
app.Logger,
|
|
app.Config().Scheduler,
|
|
).WithConfigProvider(app.ConfigStore)
|
|
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)
|
|
generationStateCheckWorker := generateworker.NewGenerationTaskStateCheckWorker(app.DB, app.Logger, app.Config().Generation).
|
|
WithConfigProvider(app.ConfigStore)
|
|
|
|
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)
|
|
startSchedulerWorker(ctx, &workerWG, "generation_state_check", app.Logger.Sugar(), generationStateCheckWorker.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(app)
|
|
app.Engine.GET("/api/internal/metrics/monitoring/daily-tasks", authMiddleware, func(c *gin.Context) {
|
|
response.Success(c, tenantapp.MonitoringDailyTaskMetricsSnapshotValue())
|
|
})
|
|
app.Engine.GET("/api/internal/metrics/generation/tasks", authMiddleware, func(c *gin.Context) {
|
|
response.Success(c, tenantapp.GenerationTaskMetricsSnapshotValue())
|
|
})
|
|
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(app *bootstrap.App) gin.HandlerFunc {
|
|
return schedulerMetricsAuthMiddlewareWithToken(func() string {
|
|
if cfg := app.Config(); cfg != nil {
|
|
return cfg.Scheduler.InternalMetricsToken
|
|
}
|
|
return ""
|
|
})
|
|
}
|
|
|
|
func schedulerMetricsAuthMiddlewareWithToken(token func() string) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
expected := ""
|
|
if token != nil {
|
|
expected = strings.TrimSpace(token())
|
|
}
|
|
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]
|
|
}
|