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 ProbeSQL string PreDeleteSQL 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) totalCandidateProbe := int64(0) perTable := map[string]any{} for _, target := range retentionTargets() { tableStats := map[string]any{} if dryRun { candidateProbe, err := w.probeTarget(ctx, conn, target, cutoffDate, 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, 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 } 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 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) probeTarget(ctx context.Context, conn *pgxpool.Conn, target retentionTarget, cutoffDate string, limit int) (int64, error) { var count int64 if limit <= 0 { limit = defaultMonitoringRetentionBatch } if err := conn.QueryRow(ctx, target.ProbeSQL, cutoffDate, limit).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 { affected, err := w.deleteTargetBatch(ctx, conn, target, cutoffDate, batchSize) if err != nil { return total, batches, err } if affected == 0 { break } total += affected batches++ if affected < int64(batchSize) { break } } return total, batches, nil } func (w *MonitoringRetentionWorker) deleteTargetBatch(ctx context.Context, conn *pgxpool.Conn, target retentionTarget, cutoffDate string, batchSize int) (int64, error) { if strings.TrimSpace(target.PreDeleteSQL) == "" { tag, err := conn.Exec(ctx, target.DeleteSQL, cutoffDate, batchSize) if err != nil { return 0, err } return tag.RowsAffected(), nil } tx, err := conn.Begin(ctx) if err != nil { return 0, fmt.Errorf("begin delete batch: %w", err) } committed := false defer func() { if !committed { _ = tx.Rollback(context.Background()) } }() if _, err := tx.Exec(ctx, target.PreDeleteSQL, cutoffDate, batchSize); err != nil { return 0, fmt.Errorf("pre-delete batch: %w", err) } tag, err := tx.Exec(ctx, target.DeleteSQL, cutoffDate, batchSize) if err != nil { return 0, err } if err := tx.Commit(ctx); err != nil { return 0, fmt.Errorf("commit delete batch: %w", err) } committed = true return tag.RowsAffected(), nil } func retentionTargets() []retentionTarget { return []retentionTarget{ { Name: "monitoring_collect_dispatch_outbox", ProbeSQL: batchProbeSQL("monitoring_collect_dispatch_outbox", "id", "created_at < $1::date", "created_at ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_collect_dispatch_outbox", "id", "created_at < $1::date", "created_at ASC, id ASC"), }, { Name: "monitoring_collect_requests", ProbeSQL: batchProbeSQL("monitoring_collect_requests", "id", "created_at < $1::date", "created_at ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_collect_requests", "id", "created_at < $1::date", "created_at ASC, id ASC"), }, { Name: "monitoring_brand_platform_daily", ProbeSQL: batchProbeSQL("monitoring_brand_platform_daily", "id", "business_date < $1::date", "business_date ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_brand_platform_daily", "id", "business_date < $1::date", "business_date ASC, id ASC"), }, { Name: "monitoring_brand_daily", ProbeSQL: batchProbeSQL("monitoring_brand_daily", "id", "business_date < $1::date", "business_date ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_brand_daily", "id", "business_date < $1::date", "business_date ASC, id ASC"), }, { Name: "monitoring_platform_access_snapshots", ProbeSQL: batchProbeSQL("monitoring_platform_access_snapshots", "id", "business_date < $1::date", "business_date ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_platform_access_snapshots", "id", "business_date < $1::date", "business_date ASC, id ASC"), }, { Name: "monitoring_collect_tasks", ProbeSQL: batchProbeSQL("monitoring_collect_tasks", "id", "business_date < $1::date", "business_date ASC, id ASC"), DeleteSQL: batchDeleteSQL("monitoring_collect_tasks", "id", "business_date < $1::date", "business_date ASC, id ASC"), }, { Name: "question_monitor_runs", ProbeSQL: batchProbeSQL("question_monitor_runs", "id", "business_date < $1::date", "business_date ASC, id ASC"), DeleteSQL: batchDeleteSQL("question_monitor_runs", "id", "business_date < $1::date", "business_date ASC, id ASC"), }, { Name: "monitoring_user_marked_articles", ProbeSQL: batchProbeSQL("monitoring_user_marked_articles", "id", "$1::date IS NOT NULL AND expires_at < NOW()", "expires_at ASC, id ASC"), PreDeleteSQL: markedArticleAttributionCleanupSQL(), DeleteSQL: batchDeleteSQL("monitoring_user_marked_articles", "id", "$1::date IS NOT NULL AND expires_at < NOW()", "expires_at ASC, id ASC"), }, } } func batchProbeSQL(table, idColumn, predicate, orderBy string) string { parts := []string{ "SELECT COUNT(*)::bigint", "FROM (", "SELECT " + idColumn, "FROM " + table, "WHERE " + predicate, "ORDER BY " + orderBy, "LIMIT $2", ") candidates", } return strings.Join(parts, "\n") } func batchDeleteSQL(table, idColumn, predicate, orderBy string) string { parts := []string{ "WITH doomed AS (", "SELECT " + idColumn, "FROM " + table, "WHERE " + predicate, "ORDER BY " + orderBy, "LIMIT $2", "FOR UPDATE SKIP LOCKED", ")", "DELETE FROM " + table + " t", "USING doomed", "WHERE t." + idColumn + " = doomed." + idColumn, } return strings.Join(parts, "\n") } func markedArticleAttributionCleanupSQL() string { parts := []string{ "WITH doomed AS (", "SELECT tenant_id, id", "FROM monitoring_user_marked_articles", "WHERE $1::date IS NOT NULL AND expires_at < NOW()", "ORDER BY expires_at ASC, id ASC", "LIMIT $2", "FOR UPDATE SKIP LOCKED", ")", "UPDATE monitoring_citation_facts cf", "SET content_source_type = NULL,", " content_source_id = NULL", "FROM doomed", "WHERE cf.tenant_id = doomed.tenant_id", " AND cf.content_source_type = 'manual_mark'", " AND cf.content_source_id = doomed.id", } return strings.Join(parts, "\n") }