topfans/backend/services/statisticService/worker/partitioner.go

package worker

import (
	"context"
	"database/sql"
	"fmt"
	"sync"
	"time"

	"go.uber.org/zap"

	"github.com/topfans/backend/pkg/logger"
	"github.com/topfans/backend/services/statisticService/metrics"
)

// Partitioner events 分区自动管理
//  - 启动时创建未来 N 天分区
//  - 每天 00:05 滚动预创建
//  - 每天 00:30 清理超过保留期的旧分区
type Partitioner struct {
	db            *sql.DB
	schema        string
	preCreateDays int
	retentionDays int

	mu      sync.Mutex
	running bool
	stop    chan struct{}
}

// NewPartitioner 构造
func NewPartitioner(db *sql.DB, schema string, preCreateDays, retentionDays int) *Partitioner {
	return &Partitioner{
		db:            db,
		schema:        schema,
		preCreateDays: preCreateDays,
		retentionDays: retentionDays,
		stop:          make(chan struct{}),
	}
}

// EnsureFuturePartitions 创建未来 days 天的分区（含今天）
func (p *Partitioner) EnsureFuturePartitions(ctx context.Context, days int) error {
	now := time.Now().In(time.FixedZone("Asia/Shanghai", 8*3600))
	for i := 0; i <= days; i++ {
		d := now.AddDate(0, 0, i)
		next := d.AddDate(0, 0, 1)
		name := fmt.Sprintf("events_%s", d.Format("2006_01_02"))
		sqlStr := fmt.Sprintf(`
			CREATE TABLE IF NOT EXISTS %s.%s PARTITION OF %s.events
			FOR VALUES FROM ('%s 00:00:00+08') TO ('%s 00:00:00+08')
		`, p.schema, name, p.schema, d.Format("2006-01-02"), next.Format("2006-01-02"))
		if _, err := p.db.ExecContext(ctx, sqlStr); err != nil {
			return fmt.Errorf("create partition %s: %w", name, err)
		}
		metrics.EventsPartitionCount.Inc()
	}
	return nil
}

// CleanupOldPartitions 删除超过 retentionDays 的旧分区
func (p *Partitioner) CleanupOldPartitions(ctx context.Context) error {
	cutoff := time.Now().In(time.FixedZone("Asia/Shanghai", 8*3600)).AddDate(0, 0, -p.retentionDays)
	cutoffName := fmt.Sprintf("events_%s", cutoff.Format("2006_01_02"))

	// LIKE 'events_%'（% 作为 LIKE 通配符匹配任何后缀）
	// 用字符串拼接避免与 fmt.Sprintf 的 % 转义冲突
	rows, err := p.db.QueryContext(ctx,
		"SELECT tablename FROM pg_tables WHERE schemaname = $1 AND tablename LIKE 'events_%' AND tablename < $2",
		p.schema, cutoffName)
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			logger.Logger.Warn("scan partition name failed", zap.Error(err))
			continue
		}
		if _, err := p.db.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s.%s CASCADE", p.schema, name)); err != nil {
			logger.Logger.Warn("drop partition failed", zap.String("name", name), zap.Error(err))
		} else {
			logger.Logger.Info("dropped old partition", zap.String("name", name))
		}
	}
	return nil
}

// Start 启动 worker（阻塞）
//  - 启动时确保未来 N 天分区
//  - 每小时检查是否到了 00:05（创建）或 00:30（清理）窗口
func (p *Partitioner) Start(ctx context.Context) {
	p.mu.Lock()
	p.running = true
	p.mu.Unlock()
	metrics.WorkerRunningCount.WithLabelValues("partitioner").Set(1)
	defer metrics.WorkerRunningCount.WithLabelValues("partitioner").Set(0)

	// 启动时跑一次
	if err := p.EnsureFuturePartitions(ctx, p.preCreateDays); err != nil {
		logger.Logger.Error("EnsureFuturePartitions failed at startup", zap.Error(err))
	}

	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-p.stop:
			return
		case <-ticker.C:
			hour := time.Now().Hour()
			min := time.Now().Minute()
			if hour == 0 && min < 10 {
				if err := p.EnsureFuturePartitions(ctx, p.preCreateDays); err != nil {
					logger.Logger.Error("EnsureFuturePartitions failed", zap.Error(err))
				}
			}
			if hour == 0 && min >= 30 && min < 40 {
				if err := p.CleanupOldPartitions(ctx); err != nil {
					logger.Logger.Error("CleanupOldPartitions failed", zap.Error(err))
				}
			}
		}
	}
}

// Stop 停止 worker
func (p *Partitioner) Stop() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.running {
		close(p.stop)
		p.running = false
	}
}