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