geo/server/internal/scheduler/scheduler_run_retention_worker.go

package scheduler

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"
)

const (
	defaultSchedulerRunRetentionDays    = 15
	defaultSchedulerRunRetentionBatch   = 5000
	defaultSchedulerRunRetentionBatches = 100
)

type SchedulerRunRetentionWorker struct {
	pool   *pgxpool.Pool
	logger *zap.Logger
}

type schedulerRunRetentionTarget struct {
	Name      string
	DeleteSQL string
	ProbeSQL  string
}

func NewSchedulerRunRetentionWorker(pool *pgxpool.Pool, logger *zap.Logger) *SchedulerRunRetentionWorker {
	return &SchedulerRunRetentionWorker{pool: pool, logger: logger}
}

func (w *SchedulerRunRetentionWorker) RunOnce(ctx context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.pool == nil {
		return map[string]any{"skipped": true, "reason": "worker_not_ready"}, nil
	}

	startedAt := time.Now()
	requestedRetentionDays := AsInt(run.Config, "retention_days", defaultSchedulerRunRetentionDays)
	retentionDays := requestedRetentionDays
	retentionDaysClamped := false
	if requestedRetentionDays < defaultSchedulerRunRetentionDays {
		retentionDays = defaultSchedulerRunRetentionDays
		retentionDaysClamped = true
		if w.logger != nil {
			w.logger.Warn("scheduler run retention_days below hard floor",
				zap.Int("requested_retention_days", requestedRetentionDays),
				zap.Int("effective_retention_days", retentionDays),
			)
		}
	}
	batchSize := AsInt(run.Config, "batch_size", defaultSchedulerRunRetentionBatch)
	if run.Job != nil && run.Job.BatchSize != nil {
		batchSize = *run.Job.BatchSize
	}
	if batchSize <= 0 {
		batchSize = defaultSchedulerRunRetentionBatch
	}
	maxBatches := AsInt(run.Config, "max_batches_per_run", defaultSchedulerRunRetentionBatches)
	if maxBatches <= 0 {
		maxBatches = defaultSchedulerRunRetentionBatches
	}
	statementTimeout := AsDuration(run.Config, "statement_timeout", 5*time.Second)
	cutoff := schedulerRunRetentionCutoff(time.Now(), retentionDays)
	dryRun := run.DryRun

	stats := map[string]any{
		"retention_days":           retentionDays,
		"requested_retention_days": requestedRetentionDays,
		"retention_days_clamped":   retentionDaysClamped,
		"cutoff_at":                cutoff,
		"batch_size":               batchSize,
		"max_batches_per_run":      maxBatches,
		"dry_run":                  dryRun,
	}

	conn, err := w.pool.Acquire(ctx)
	if err != nil {
		return stats, fmt.Errorf("acquire ops connection: %w", err)
	}
	defer conn.Release()

	if statementTimeout > 0 {
		if _, err := conn.Exec(ctx, `SELECT set_config('statement_timeout', $1, false)`, fmt.Sprintf("%dms", statementTimeout.Milliseconds())); err != nil {
			return stats, fmt.Errorf("set statement timeout: %w", err)
		}
		defer conn.Exec(context.Background(), `RESET statement_timeout`)
	}

	totalDeleted := int64(0)
	totalCandidateProbe := int64(0)
	perTable := map[string]any{}
	for _, target := range schedulerRunRetentionTargets() {
		tableStats := map[string]any{}
		if dryRun {
			candidateProbe, err := w.probeTarget(ctx, conn, target, cutoff, batchSize)
			if err != nil {
				return stats, fmt.Errorf("probe %s: %w", target.Name, err)
			}
			tableStats["candidate_probe_count"] = candidateProbe
			tableStats["candidate_probe_limited"] = candidateProbe >= int64(batchSize)
			totalCandidateProbe += candidateProbe
		} else {
			deleted, batches, err := w.deleteTarget(ctx, conn, target, cutoff, batchSize, maxBatches)
			if err != nil {
				return stats, fmt.Errorf("delete %s: %w", target.Name, err)
			}
			tableStats["deleted_count"] = deleted
			tableStats["batch_count"] = batches
			totalDeleted += deleted
		}
		perTable[target.Name] = tableStats
	}

	if dryRun {
		stats["candidate_probe_count"] = totalCandidateProbe
	}
	stats["deleted_count"] = totalDeleted
	stats["tables"] = perTable
	stats["duration_ms"] = time.Since(startedAt).Milliseconds()
	return stats, nil
}

func schedulerRunRetentionCutoff(now time.Time, retentionDays int) time.Time {
	if retentionDays < defaultSchedulerRunRetentionDays {
		retentionDays = defaultSchedulerRunRetentionDays
	}
	loc, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		loc = time.FixedZone("UTC+8", 8*60*60)
	}
	localNow := now.In(loc)
	today := time.Date(localNow.Year(), localNow.Month(), localNow.Day(), 0, 0, 0, 0, loc)
	return today.AddDate(0, 0, -(retentionDays - 1)).UTC()
}

func (w *SchedulerRunRetentionWorker) probeTarget(ctx context.Context, conn *pgxpool.Conn, target schedulerRunRetentionTarget, cutoff time.Time, limit int) (int64, error) {
	var count int64
	if limit <= 0 {
		limit = defaultSchedulerRunRetentionBatch
	}
	if err := conn.QueryRow(ctx, target.ProbeSQL, cutoff.UTC(), limit).Scan(&count); err != nil {
		return 0, err
	}
	return count, nil
}

func (w *SchedulerRunRetentionWorker) deleteTarget(ctx context.Context, conn *pgxpool.Conn, target schedulerRunRetentionTarget, cutoff time.Time, batchSize, maxBatches int) (int64, int, error) {
	total := int64(0)
	batches := 0
	for batches < maxBatches {
		tag, err := conn.Exec(ctx, target.DeleteSQL, cutoff.UTC(), batchSize)
		if err != nil {
			return total, batches, err
		}
		affected := tag.RowsAffected()
		if affected == 0 {
			break
		}
		total += affected
		batches++
		if affected < int64(batchSize) {
			break
		}
	}
	return total, batches, nil
}

func schedulerRunRetentionTargets() []schedulerRunRetentionTarget {
	return []schedulerRunRetentionTarget{
		{
			Name:      "scheduler_job_triggers",
			ProbeSQL:  schedulerRunRetentionTriggerProbeSQL(),
			DeleteSQL: schedulerRunRetentionTriggerDeleteSQL(),
		},
		{
			Name:      "scheduler_job_runs",
			ProbeSQL:  schedulerRunRetentionRunProbeSQL(),
			DeleteSQL: schedulerRunRetentionRunDeleteSQL(),
		},
	}
}

func schedulerRunRetentionTriggerProbeSQL() string {
	return schedulerRunRetentionProbeSQL(
		"ops.scheduler_job_triggers",
		"id",
		"status IN ('success', 'failed', 'skipped') AND requested_at < $1",
		"requested_at ASC, id ASC",
	)
}

func schedulerRunRetentionTriggerDeleteSQL() string {
	return schedulerRunRetentionDeleteSQL(
		"ops.scheduler_job_triggers",
		"id",
		"status IN ('success', 'failed', 'skipped') AND requested_at < $1",
		"requested_at ASC, id ASC",
	)
}

func schedulerRunRetentionRunProbeSQL() string {
	return schedulerRunRetentionProbeSQL(
		"ops.scheduler_job_runs r",
		"r.id",
		"r.status IN ('success', 'failed', 'skipped') AND r.started_at < $1 AND NOT EXISTS (SELECT 1 FROM ops.scheduler_job_triggers t WHERE t.run_id = r.id)",
		"r.started_at ASC, r.id ASC",
	)
}

func schedulerRunRetentionRunDeleteSQL() string {
	return schedulerRunRetentionDeleteSQL(
		"ops.scheduler_job_runs r",
		"r.id",
		"r.status IN ('success', 'failed', 'skipped') AND r.started_at < $1 AND NOT EXISTS (SELECT 1 FROM ops.scheduler_job_triggers t WHERE t.run_id = r.id)",
		"r.started_at ASC, r.id ASC",
	)
}

func schedulerRunRetentionProbeSQL(table, idColumn, predicate, orderBy string) string {
	return strings.Join([]string{
		"SELECT COUNT(*)::bigint",
		"FROM (",
		"SELECT " + idColumn,
		"FROM " + table,
		"WHERE " + predicate,
		"ORDER BY " + orderBy,
		"LIMIT $2",
		") candidates",
	}, "\n")
}

func schedulerRunRetentionDeleteSQL(table, idColumn, predicate, orderBy string) string {
	return strings.Join([]string{
		"WITH doomed AS (",
		"SELECT " + idColumn,
		"FROM " + table,
		"WHERE " + predicate,
		"ORDER BY " + orderBy,
		"LIMIT $2",
		"FOR UPDATE SKIP LOCKED",
		")",
		"DELETE FROM " + strings.Fields(table)[0] + " target",
		"USING doomed",
		"WHERE target.id = doomed.id",
	}, "\n")
}