package scheduler import ( "context" "errors" "fmt" "hash/fnv" "os" "strconv" "strings" "sync" "time" "github.com/jackc/pgx/v5/pgxpool" "go.uber.org/zap" opsdomain "github.com/geo-platform/tenant-api/internal/ops/domain" opsrepo "github.com/geo-platform/tenant-api/internal/ops/repository" sharedschedule "github.com/geo-platform/tenant-api/internal/shared/schedule" ) type ControlledJob struct { Key string DefaultName string Run func(context.Context, JobRunContext) (map[string]any, error) } type JobRunContext struct { Job *opsdomain.SchedulerJob TriggerID *int64 Trigger string DryRun bool Config map[string]any } const ( schedulerManualPollInterval = 15 * time.Second schedulerConfigRetryInterval = 30 * time.Second schedulerFinishTimeout = 10 * time.Second schedulerMaxManualTriggersDrain = 100 ) type ControlPlane struct { pool *pgxpool.Pool repo *opsrepo.SchedulerRepository logger *zap.Logger instanceID string hostname string startedAt time.Time jobs map[string]ControlledJob } func NewControlPlane(pool *pgxpool.Pool, logger *zap.Logger, jobs []ControlledJob) *ControlPlane { hostname, _ := os.Hostname() if hostname == "" { hostname = "unknown" } instanceID := hostname + ":" + strconv.Itoa(os.Getpid()) + ":" + strconv.FormatInt(time.Now().UnixNano(), 36) byKey := make(map[string]ControlledJob, len(jobs)) for _, job := range jobs { if job.Key == "" || job.Run == nil { continue } byKey[job.Key] = job } return &ControlPlane{ pool: pool, repo: opsrepo.NewSchedulerRepository(pool), logger: logger, instanceID: instanceID, hostname: hostname, startedAt: time.Now().UTC(), jobs: byKey, } } func (p *ControlPlane) Run(ctx context.Context) { if p == nil || p.pool == nil || len(p.jobs) == 0 { return } p.heartbeat(context.Background()) heartbeatCtx, heartbeatCancel := context.WithCancel(ctx) defer heartbeatCancel() go p.runHeartbeat(heartbeatCtx) var wg sync.WaitGroup for _, job := range p.jobs { job := job wg.Add(1) go func() { defer wg.Done() p.runJobLoop(ctx, job) }() } wg.Wait() } func (p *ControlPlane) InstanceID() string { if p == nil { return "" } return p.instanceID } func (p *ControlPlane) runHeartbeat(ctx context.Context) { ticker := time.NewTicker(15 * time.Second) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: p.heartbeat(context.Background()) } } } func (p *ControlPlane) heartbeat(ctx context.Context) { if p == nil || p.repo == nil { return } buildVersion := os.Getenv("BUILD_VERSION") if buildVersion == "" { buildVersion = os.Getenv("GIT_SHA") } var buildVersionPtr *string if buildVersion != "" { buildVersionPtr = &buildVersion } if err := p.repo.UpsertInstance(ctx, opsdomain.SchedulerInstance{ InstanceID: p.instanceID, Hostname: p.hostname, Process: "scheduler", BuildVersion: buildVersionPtr, StartedAt: p.startedAt, LastSeenAt: time.Now().UTC(), Metadata: map[string]any{ "pid": os.Getpid(), }, }); err != nil && p.logger != nil { p.logger.Warn("scheduler control heartbeat failed", zap.Error(err)) } } func (p *ControlPlane) runJobLoop(ctx context.Context, job ControlledJob) { nextAutoAt := p.nextAutoAt(ctx, job.Key, time.Now(), true) for { if ctx.Err() != nil { return } p.drainManualTriggers(ctx, job) now := time.Now() if !now.Before(nextAutoAt) { p.runJobOnce(ctx, job, opsdomain.SchedulerTriggerAuto, nil, nil) nextAutoAt = p.nextAutoAt(ctx, job.Key, time.Now(), false) continue } sleepFor := schedulerLoopSleep(time.Now(), nextAutoAt) timer := time.NewTimer(sleepFor) select { case <-ctx.Done(): timer.Stop() return case <-timer.C: } } } func (p *ControlPlane) drainManualTriggers(parent context.Context, job ControlledJob) { for claimed := 0; claimed < schedulerMaxManualTriggersDrain; claimed++ { if parent.Err() != nil { return } trigger, ok, err := p.repo.ClaimPendingTrigger(parent, job.Key, p.instanceID) if err != nil { if p.logger != nil { p.logger.Warn("scheduler trigger claim failed", zap.String("job_key", job.Key), zap.Error(err)) } return } if !ok { return } triggerID := trigger.ID p.runJobOnce(parent, job, trigger.TriggerType, &triggerID, trigger.ConfigOverride) } if p.logger != nil { p.logger.Warn("scheduler manual trigger drain reached batch limit", zap.String("job_key", job.Key), zap.Int("limit", schedulerMaxManualTriggersDrain), ) } } func (p *ControlPlane) nextAutoAt(ctx context.Context, key string, now time.Time, initial bool) time.Time { job, err := p.repo.GetJobForRuntime(ctx, key) if err != nil || job == nil { return now.Add(schedulerConfigRetryInterval) } if job.ScheduleType == "manual" { return now.Add(schedulerManualPollInterval) } if job.ScheduleType == "cron" && job.CronExpr != nil { loc, locErr := time.LoadLocation(job.Timezone) if locErr != nil { loc = time.Local } schedule, parseErr := sharedschedule.Parse(*job.CronExpr) if parseErr != nil { if p.logger != nil { p.logger.Warn("scheduler controlled job has invalid cron", zap.String("job_key", key), zap.String("cron_expr", *job.CronExpr), zap.Error(parseErr), ) } return now.Add(schedulerConfigRetryInterval) } base := now.In(loc) return schedule.Next(base) } if runAt := parseRunAtLocal(job.Config); runAt != "" { loc, locErr := time.LoadLocation(job.Timezone) if locErr != nil { loc = time.Local } hour, minute, ok := parseHourMinute(runAt) if ok { localNow := now.In(loc) candidate := time.Date(localNow.Year(), localNow.Month(), localNow.Day(), hour, minute, 0, 0, loc) if !initial || !localNow.Before(candidate) { candidate = candidate.AddDate(0, 0, 1) } return candidate } } interval := time.Duration(job.IntervalSeconds) * time.Second if interval <= 0 { interval = schedulerConfigRetryInterval } var lastStartedAt *time.Time if initial { startedAt, ok, lastErr := p.repo.LastAutoRunStartedAt(ctx, key) if lastErr != nil { if p.logger != nil { p.logger.Warn("scheduler last auto run lookup failed", zap.String("job_key", key), zap.Error(lastErr), ) } } else if ok { lastStartedAt = &startedAt } } return nextIntervalAutoAt(now, interval, initial, lastStartedAt) } func (p *ControlPlane) runJobOnce(parent context.Context, job ControlledJob, triggerType string, triggerID *int64, configOverride map[string]any) { cfg, err := p.repo.GetJobForRuntime(parent, job.Key) if err != nil { if p.logger != nil { p.logger.Warn("scheduler job config load failed", zap.String("job_key", job.Key), zap.Error(err)) } return } isManual := triggerType == opsdomain.SchedulerTriggerManual || triggerType == opsdomain.SchedulerTriggerDryRun if !cfg.Enabled && !isManual { return } if cfg.ScheduleType == "manual" && !isManual { return } mergedConfig := cloneMap(cfg.Config) for key, value := range configOverride { mergedConfig[key] = value } dryRun := triggerType == opsdomain.SchedulerTriggerDryRun || boolFromMap(mergedConfig, "dry_run") if dryRun { mergedConfig["dry_run"] = true } lockKey := advisoryLockKey(job.Key) lockConn, locked := p.tryAdvisoryLock(parent, lockKey) if !locked { if p.logger != nil { p.logger.Info("scheduler job skipped because another instance owns lock", zap.String("job_key", job.Key)) } if triggerID != nil { _ = p.repo.FinishTrigger(parent, *triggerID, nil, opsdomain.SchedulerRunSkipped, controlStringPtr("another scheduler instance is running this job")) } return } defer p.advisoryUnlock(context.Background(), lockConn, lockKey) run, err := p.repo.StartRun(parent, opsrepo.SchedulerRunStart{ JobKey: job.Key, TriggerID: triggerID, SchedulerInstanceID: p.instanceID, TriggerType: triggerType, ConfigVersion: cfg.Version, ConfigSnapshot: schedulerConfigSnapshot(cfg, mergedConfig), }) if err != nil { if p.logger != nil { p.logger.Warn("scheduler run start failed", zap.String("job_key", job.Key), zap.Error(err)) } return } timeout := time.Duration(cfg.TimeoutSeconds) * time.Second if timeout <= 0 { timeout = 30 * time.Second } ctx, cancel := context.WithTimeout(parent, timeout) defer cancel() stats, runErr := job.Run(ctx, JobRunContext{ Job: cfg, TriggerID: triggerID, Trigger: triggerType, DryRun: dryRun, Config: mergedConfig, }) status := opsdomain.SchedulerRunSuccess var errMsg *string if runErr != nil { status = opsdomain.SchedulerRunFailed message := runErr.Error() errMsg = &message } if stats == nil { stats = map[string]any{} } stats["dry_run"] = dryRun finishCtx, finishCancel := context.WithTimeout(context.Background(), schedulerFinishTimeout) defer finishCancel() finished, finishErr := p.repo.FinishRun(finishCtx, opsrepo.SchedulerRunFinish{ RunID: run.ID, Status: status, Stats: stats, ErrorMessage: errMsg, }) if finishErr != nil && p.logger != nil { p.logger.Warn("scheduler run finish failed", zap.String("job_key", job.Key), zap.Int64("run_id", run.ID), zap.Error(finishErr)) } if triggerID != nil { triggerStatus := status if finishErr != nil { triggerStatus = opsdomain.SchedulerRunFailed } runID := run.ID _ = p.repo.FinishTrigger(finishCtx, *triggerID, &runID, triggerStatus, errMsg) } if runErr != nil { if p.logger != nil { p.logger.Warn("scheduler controlled job failed", zap.String("job_key", job.Key), zap.Error(runErr)) } return } if p.logger != nil { duration := int64(0) if finished != nil && finished.DurationMilliseconds != nil { duration = *finished.DurationMilliseconds } p.logger.Info("scheduler controlled job completed", zap.String("job_key", job.Key), zap.Int64("duration_ms", duration)) } } func (p *ControlPlane) tryAdvisoryLock(ctx context.Context, key int64) (*pgxpool.Conn, bool) { conn, err := p.pool.Acquire(ctx) if err != nil { if p.logger != nil { p.logger.Warn("scheduler advisory lock acquire failed", zap.Int64("lock_key", key), zap.Error(err)) } return nil, false } var ok bool if err := conn.QueryRow(ctx, `SELECT pg_try_advisory_lock($1)`, key).Scan(&ok); err != nil { conn.Release() if p.logger != nil { p.logger.Warn("scheduler advisory lock failed", zap.Int64("lock_key", key), zap.Error(err)) } return nil, false } if !ok { conn.Release() return nil, false } return conn, true } func (p *ControlPlane) advisoryUnlock(ctx context.Context, conn *pgxpool.Conn, key int64) { if conn == nil { return } defer conn.Release() _, _ = conn.Exec(ctx, `SELECT pg_advisory_unlock($1)`, key) } func advisoryLockKey(value string) int64 { hash := fnv.New64a() _, _ = hash.Write([]byte("scheduler:" + value)) return int64(hash.Sum64() & 0x7fffffffffffffff) } func schedulerLoopSleep(now, nextAutoAt time.Time) time.Duration { sleepFor := nextAutoAt.Sub(now) if sleepFor > schedulerManualPollInterval { sleepFor = schedulerManualPollInterval } if sleepFor <= 0 { sleepFor = time.Second } return sleepFor } func nextIntervalAutoAt(now time.Time, interval time.Duration, initial bool, lastStartedAt *time.Time) time.Time { if interval <= 0 { interval = schedulerConfigRetryInterval } if initial { if lastStartedAt != nil && !lastStartedAt.IsZero() { dueAt := lastStartedAt.Add(interval) if dueAt.After(now) { return dueAt } } return now } return now.Add(interval) } func schedulerConfigSnapshot(job *opsdomain.SchedulerJob, config map[string]any) map[string]any { out := map[string]any{ "enabled": job.Enabled, "schedule_type": job.ScheduleType, "interval_seconds": job.IntervalSeconds, "timezone": job.Timezone, "timeout_seconds": job.TimeoutSeconds, "max_concurrency": job.MaxConcurrency, "config": config, } if job.BatchSize != nil { out["batch_size"] = *job.BatchSize } if job.CronExpr != nil { out["cron_expr"] = *job.CronExpr } return out } func cloneMap(in map[string]any) map[string]any { out := make(map[string]any, len(in)) for key, value := range in { out[key] = value } return out } func boolFromMap(in map[string]any, key string) bool { value, ok := in[key] if !ok { return false } switch typed := value.(type) { case bool: return typed case string: parsed, _ := strconv.ParseBool(typed) return parsed default: return false } } func parseRunAtLocal(config map[string]any) string { value, ok := config["run_at_local"] if !ok || value == nil { return "" } return fmt.Sprint(value) } func parseHourMinute(value string) (int, int, bool) { parts := strings.Split(strings.TrimSpace(value), ":") if len(parts) != 2 { return 0, 0, false } hour, hourErr := strconv.Atoi(parts[0]) minute, minuteErr := strconv.Atoi(parts[1]) if hourErr != nil || minuteErr != nil || hour < 0 || hour > 23 || minute < 0 || minute > 59 { return 0, 0, false } return hour, minute, true } func controlStringPtr(value string) *string { return &value } func AsInt(config map[string]any, key string, fallback int) int { value, ok := config[key] if !ok || value == nil { return fallback } switch typed := value.(type) { case int: return typed case int64: return int(typed) case float64: return int(typed) case string: parsed, err := strconv.Atoi(typed) if err == nil { return parsed } } return fallback } func AsDuration(config map[string]any, key string, fallback time.Duration) time.Duration { value, ok := config[key] if !ok || value == nil { return fallback } switch typed := value.(type) { case time.Duration: return typed case string: parsed, err := time.ParseDuration(typed) if err == nil { return parsed } case float64: return time.Duration(typed) * time.Second } return fallback } func ErrorFromContext(ctx context.Context) error { if err := ctx.Err(); err != nil && !errors.Is(err, context.Canceled) { return fmt.Errorf("scheduler job context finished: %w", err) } return nil }