geo/server/internal/scheduler/monitoring_retention_worker.go

package scheduler

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

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

const (
	defaultMonitoringRetentionDays    = 30
	defaultMonitoringRetentionBatch   = 5000
	defaultMonitoringRetentionBatches = 200
)

type MonitoringRetentionWorker struct {
	monitoringPool *pgxpool.Pool
	logger         *zap.Logger
}

type retentionTarget struct {
	Name      string
	DeleteSQL string
	CountSQL  string
}

func NewMonitoringRetentionWorker(monitoringPool *pgxpool.Pool, logger *zap.Logger) *MonitoringRetentionWorker {
	return &MonitoringRetentionWorker{monitoringPool: monitoringPool, logger: logger}
}

func (w *MonitoringRetentionWorker) RunOnce(ctx context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.monitoringPool == nil {
		return map[string]any{"skipped": true, "reason": "worker_not_ready"}, nil
	}
	startedAt := time.Now()
	retentionDays := AsInt(run.Config, "retention_days", defaultMonitoringRetentionDays)
	if retentionDays < defaultMonitoringRetentionDays {
		retentionDays = defaultMonitoringRetentionDays
	}
	batchSize := AsInt(run.Config, "batch_size", defaultMonitoringRetentionBatch)
	if run.Job != nil && run.Job.BatchSize != nil {
		batchSize = *run.Job.BatchSize
	}
	if batchSize <= 0 {
		batchSize = defaultMonitoringRetentionBatch
	}
	maxBatches := AsInt(run.Config, "max_batches_per_run", defaultMonitoringRetentionBatches)
	if maxBatches <= 0 {
		maxBatches = defaultMonitoringRetentionBatches
	}
	statementTimeout := AsDuration(run.Config, "statement_timeout", 5*time.Second)
	cutoffDate := monitoringRetentionCutoffDate(time.Now(), retentionDays)
	dryRun := run.DryRun

	stats := map[string]any{
		"retention_days":      retentionDays,
		"cutoff_date":         cutoffDate,
		"batch_size":          batchSize,
		"max_batches_per_run": maxBatches,
		"dry_run":             dryRun,
	}

	conn, err := w.monitoringPool.Acquire(ctx)
	if err != nil {
		return stats, fmt.Errorf("acquire monitoring 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)
	totalCandidates := int64(0)
	perTable := map[string]any{}
	for _, target := range retentionTargets() {
		candidateCount, err := w.countTarget(ctx, conn, target, cutoffDate)
		if err != nil {
			return stats, fmt.Errorf("count %s: %w", target.Name, err)
		}
		totalCandidates += candidateCount
		tableStats := map[string]any{
			"candidate_count": candidateCount,
		}
		if !dryRun && candidateCount > 0 {
			deleted, batches, err := w.deleteTarget(ctx, conn, target, cutoffDate, 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
	}

	stats["candidate_count"] = totalCandidates
	stats["deleted_count"] = totalDeleted
	stats["tables"] = perTable
	stats["duration_ms"] = time.Since(startedAt).Milliseconds()
	return stats, nil
}

func monitoringRetentionCutoffDate(now time.Time, retentionDays int) string {
	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)
	cutoff := today.AddDate(0, 0, -(retentionDays - 1))
	return cutoff.Format("2006-01-02")
}

func (w *MonitoringRetentionWorker) countTarget(ctx context.Context, conn *pgxpool.Conn, target retentionTarget, cutoffDate string) (int64, error) {
	var count int64
	if err := conn.QueryRow(ctx, target.CountSQL, cutoffDate).Scan(&count); err != nil {
		return 0, err
	}
	return count, nil
}

func (w *MonitoringRetentionWorker) deleteTarget(ctx context.Context, conn *pgxpool.Conn, target retentionTarget, cutoffDate string, batchSize, maxBatches int) (int64, int, error) {
	total := int64(0)
	batches := 0
	for batches < maxBatches {
		tag, err := conn.Exec(ctx, target.DeleteSQL, cutoffDate, 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 retentionTargets() []retentionTarget {
	return []retentionTarget{
		{
			Name:      "monitoring_collect_dispatch_outbox",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_collect_dispatch_outbox WHERE created_at < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_collect_dispatch_outbox", "id", "created_at < $1::date"),
		},
		{
			Name:      "monitoring_collect_requests",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_collect_requests WHERE created_at < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_collect_requests", "id", "created_at < $1::date"),
		},
		{
			Name:      "monitoring_brand_platform_daily",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_brand_platform_daily WHERE business_date < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_brand_platform_daily", "id", "business_date < $1::date"),
		},
		{
			Name:      "monitoring_brand_daily",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_brand_daily WHERE business_date < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_brand_daily", "id", "business_date < $1::date"),
		},
		{
			Name:      "monitoring_platform_access_snapshots",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_platform_access_snapshots WHERE business_date < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_platform_access_snapshots", "id", "business_date < $1::date"),
		},
		{
			Name:      "monitoring_collect_tasks",
			CountSQL:  `SELECT COUNT(*) FROM monitoring_collect_tasks WHERE business_date < $1::date`,
			DeleteSQL: batchDeleteSQL("monitoring_collect_tasks", "id", "business_date < $1::date"),
		},
		{
			Name:      "question_monitor_runs",
			CountSQL:  `SELECT COUNT(*) FROM question_monitor_runs WHERE business_date < $1::date`,
			DeleteSQL: batchDeleteSQL("question_monitor_runs", "id", "business_date < $1::date"),
		},
	}
}

func batchDeleteSQL(table, idColumn, predicate string) string {
	parts := []string{
		"WITH doomed AS (",
		"SELECT " + idColumn,
		"FROM " + table,
		"WHERE " + predicate,
		"ORDER BY " + idColumn + " ASC",
		"LIMIT $2",
		")",
		"DELETE FROM " + table,
		"WHERE " + idColumn + " IN (SELECT " + idColumn + " FROM doomed)",
	}
	return strings.Join(parts, "\n")
}