geo/server/internal/scheduler/control_plane.go

package scheduler

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"math"
	"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
	}
	// Empty-success suppression is only for high-frequency automatic polling jobs.
	// Suppressed runs intentionally do not create running rows and do not advance
	// LastAutoRunStartedAt, so process restarts may run them immediately once as
	// overdue. Manual, dry-run, failed, and non-empty executions still get full
	// audit rows below.
	suppressEmptySuccessRun := shouldSuppressEmptySuccessRun(triggerType, dryRun, mergedConfig)

	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)

	configSnapshot := schedulerConfigSnapshot(cfg, mergedConfig)
	startedAt := time.Now()
	var run *opsdomain.SchedulerRun
	if !suppressEmptySuccessRun {
		run, err = p.repo.StartRun(parent, opsrepo.SchedulerRunStart{
			JobKey:              job.Key,
			TriggerID:           triggerID,
			SchedulerInstanceID: p.instanceID,
			TriggerType:         triggerType,
			ConfigVersion:       cfg.Version,
			ConfigSnapshot:      configSnapshot,
		})
		if err != nil {
			if p.logger != nil {
				p.logger.Warn("scheduler run start failed", zap.String("job_key", job.Key), zap.Error(err))
			}
			return
		}
		startedAt = run.StartedAt
	}

	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
	finishedAt := time.Now()
	if suppressEmptySuccessRun && runErr == nil && isEmptySuccessRun(stats, mergedConfig) {
		if p.logger != nil {
			p.logger.Info("scheduler empty success run suppressed", zap.String("job_key", job.Key))
		}
		return
	}

	finishCtx, finishCancel := context.WithTimeout(context.Background(), schedulerFinishTimeout)
	defer finishCancel()
	var finished *opsdomain.SchedulerRun
	var finishErr error
	if run == nil {
		finished, finishErr = p.repo.RecordRun(finishCtx, opsrepo.SchedulerRunRecord{
			JobKey:              job.Key,
			TriggerID:           triggerID,
			SchedulerInstanceID: p.instanceID,
			TriggerType:         triggerType,
			ConfigVersion:       cfg.Version,
			ConfigSnapshot:      configSnapshot,
			Status:              status,
			Stats:               stats,
			ErrorMessage:        errMsg,
			StartedAt:           startedAt,
			FinishedAt:          finishedAt,
		})
	} else {
		finished, finishErr = p.repo.FinishRun(finishCtx, opsrepo.SchedulerRunFinish{
			RunID:        run.ID,
			Status:       status,
			Stats:        stats,
			ErrorMessage: errMsg,
		})
	}
	if finishErr != nil && p.logger != nil {
		fields := []zap.Field{zap.String("job_key", job.Key), zap.Error(finishErr)}
		if run != nil {
			fields = append(fields, zap.Int64("run_id", run.ID))
		}
		if finished != nil {
			fields = append(fields, zap.Int64("finished_run_id", finished.ID))
		}
		p.logger.Warn("scheduler run finish failed", fields...)
	}
	if triggerID != nil {
		triggerStatus := status
		if finishErr != nil {
			triggerStatus = opsdomain.SchedulerRunFailed
		}
		var runID *int64
		if finished != nil {
			value := finished.ID
			runID = &value
		}
		_ = 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 shouldSuppressEmptySuccessRun(triggerType string, dryRun bool, config map[string]any) bool {
	if dryRun || triggerType != opsdomain.SchedulerTriggerAuto {
		return false
	}
	return boolFromMap(config, "suppress_empty_success_runs")
}

func isEmptySuccessRun(stats map[string]any, config map[string]any) bool {
	if stats == nil || !boolFromMap(config, "suppress_empty_success_runs") {
		return false
	}
	keys := emptySuccessMetricKeys(config)
	if len(keys) == 0 {
		return false
	}
	for _, key := range keys {
		value, ok := stats[key]
		if !ok {
			return false
		}
		if numericValue(value) != 0 {
			return false
		}
	}
	return true
}

func emptySuccessMetricKeys(config map[string]any) []string {
	value, ok := config["empty_success_metric_keys"]
	if !ok || value == nil {
		return nil
	}
	switch typed := value.(type) {
	case []string:
		return typed
	case []any:
		out := make([]string, 0, len(typed))
		for _, item := range typed {
			if key := strings.TrimSpace(fmt.Sprint(item)); key != "" {
				out = append(out, key)
			}
		}
		return out
	case string:
		parts := strings.Split(typed, ",")
		out := make([]string, 0, len(parts))
		for _, part := range parts {
			if key := strings.TrimSpace(part); key != "" {
				out = append(out, key)
			}
		}
		return out
	default:
		return nil
	}
}

func numericValue(value any) float64 {
	switch typed := value.(type) {
	case int:
		return float64(typed)
	case int8:
		return float64(typed)
	case int16:
		return float64(typed)
	case int32:
		return float64(typed)
	case int64:
		return float64(typed)
	case uint:
		return float64(typed)
	case uint8:
		return float64(typed)
	case uint16:
		return float64(typed)
	case uint32:
		return float64(typed)
	case uint64:
		return float64(typed)
	case float32:
		return float64(typed)
	case float64:
		if math.IsNaN(typed) {
			return 0
		}
		return typed
	case string:
		parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64)
		if err == nil {
			return parsed
		}
	}
	return 0
}

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
}