This commit is contained in:
@@ -0,0 +1,248 @@
|
||||
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")
|
||||
}
|
||||
Reference in New Issue
Block a user