upstream-centrifugo/internal/pgoutbox/partitioner.go

package pgoutbox

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

	"github.com/centrifugal/centrifugo/v6/internal/pgschema"

	"github.com/jackc/pgerrcode"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

// isShutdownErr reports whether err is a normal-shutdown signal — context
// cancellation surfacing because the broker is being closed mid-tick. Those
// aren't real failures and shouldn't be logged at error level, since they
// would otherwise pollute CI logs and mask genuine cleanup errors.
func isShutdownErr(err error) bool {
	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
}

// isDuplicateTableErr reports whether err is PostgreSQL's 42P07, raised when the
// relation already exists. CREATE TABLE IF NOT EXISTS checks the catalog before
// it takes the lock, so two nodes creating the same partition at the same moment
// both pass the check and the loser gets a hard error rather than the notice a
// sequential re-run produces. The partition exists either way, which is the
// outcome the statement was asking for.
func isDuplicateTableErr(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && pgErr.Code == pgerrcode.DuplicateTable
}

// Partitioner manages daily time-based partition maintenance for an
// existing PARTITION BY RANGE (created_at) parent table.
//
// It does NOT create the parent partitioned table itself — the parent's
// column list is schema-specific and must be created by the calling
// broker before any Partitioner methods are called. Partitioner only
// handles the ongoing maintenance:
//
//   - EnsureLookaheadPartitions creates today's partition and up to
//     LookaheadDays future daily partitions via CREATE TABLE IF NOT EXISTS
//     ... PARTITION OF. Safe to call repeatedly. Used both by the caller
//     at schema init time and by Run on a periodic ticker.
//
//   - DropOldPartitions lists child partitions via pg_inherits, parses
//     the date from the partition name suffix, and drops any partition
//     older than (now - RetentionDays). Errors on individual drops are
//     logged via ErrorFn rather than returned, so a failed drop does not
//     prevent the cleanup pass from attempting the remaining partitions.
//
//   - Run starts a ticker that calls both methods on each tick.
//
// Partition naming convention: {ParentTable}_{YYYY}_{MM}_{DD}. Partitions
// that don't match this convention are silently skipped by DropOldPartitions
// — they are assumed to be user-managed and are not subject to automatic
// cleanup.
type Partitioner struct {
	// Pool is the pgxpool used for all partition maintenance SQL.
	// Typically, the primary pool (DDL must target the primary).
	Pool *pgxpool.Pool

	// ParentTable is the unquoted name of the parent partitioned table
	// (e.g. "cf_map_stream").
	ParentTable string

	// CleanupInterval is the ticker interval for Run.
	CleanupInterval time.Duration

	// LookaheadDays controls how many future daily partitions to pre-create.
	// 0 means only today's partition.
	LookaheadDays int

	// RetentionDays controls how old a partition must be before
	// DropOldPartitions removes it. Partitions dated strictly before
	// (now - RetentionDays) are dropped.
	RetentionDays int

	// ErrorFn is called for non-ctx errors that do not abort the pass
	// (e.g. per-partition drop failures, listing failures in the
	// cleanup path). Must be safe for concurrent use.
	ErrorFn func(msg string, err error)
}

// EnsureLookaheadPartitions creates today's daily partition and up to
// LookaheadDays future daily partitions as PARTITION OF ParentTable.
// The CREATE TABLE IF NOT EXISTS form makes this safe to call repeatedly, and
// tolerating 42P07 (plus retrying the other concurrent-DDL conflicts) makes it
// safe to call concurrently: every node runs the partitioner, so the same
// partition is routinely created by two of them at once.
//
// Boundaries are written as fully-qualified UTC timestamps (e.g.
// '2026-04-23 00:00:00+00'). Using a bare DATE/date-text literal would
// bind the boundary to the PostgreSQL session's timezone: a server
// configured for Asia/Nicosia would treat '2026-04-23' as
// 2026-04-22 21:00 UTC, producing partitions that are offset from the
// UTC calendar day and don't cover rows inserted in the last hours of a
// UTC day.
func (p *Partitioner) EnsureLookaheadPartitions(ctx context.Context) error {
	now := time.Now().UTC()
	for d := 0; d <= p.LookaheadDays; d++ {
		day := now.AddDate(0, 0, d)
		nextDay := day.AddDate(0, 0, 1)
		partName := fmt.Sprintf("%s_%s", p.ParentTable, day.Format("2006_01_02"))
		stmt := fmt.Sprintf(
			`CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES FROM ('%s') TO ('%s')`,
			partName, p.ParentTable,
			day.Format("2006-01-02 00:00:00+00"),
			nextDay.Format("2006-01-02 00:00:00+00"),
		)
		// Nodes that start together run this at the same moment, so the same
		// concurrent-DDL conflicts the schema batch retries apply here too —
		// and this runs on the startup path (EnsureSchema pre-creates the
		// lookahead), where a bare failure means the node refuses to start.
		// 42P07 is the one conflict that needs no retry: the partition exists,
		// which is the outcome the statement was asking for.
		err := pgschema.RetrySchemaExec(ctx, func(ctx context.Context) error {
			_, err := p.Pool.Exec(ctx, stmt)
			if isDuplicateTableErr(err) {
				return nil
			}
			return err
		})
		if err != nil {
			return fmt.Errorf("create partition %s: %w", partName, err)
		}
	}
	return nil
}

// DropOldPartitions lists child partitions of ParentTable via pg_inherits,
// parses the date from each partition name, and drops partitions dated
// before (now - RetentionDays). Errors on individual drops (or the list
// query itself) are logged via ErrorFn rather than returned.
//
// When RetentionDays is 0 or negative, this is a no-op: the caller wants
// "lookahead only, never drop" semantics. Without this guard, RetentionDays=0
// would compute cutoff = now, dropping every partition dated strictly before
// today (i.e. yesterday and older) — equivalent to a 1-day retention, not
// "no retention".
func (p *Partitioner) DropOldPartitions(ctx context.Context) {
	if p.RetentionDays <= 0 {
		return
	}
	cutoff := time.Now().UTC().AddDate(0, 0, -p.RetentionDays)

	// Filter by current_schema() so that two Centrifugo clusters (or other
	// software) sharing the same PostgreSQL instance but using different
	// search_path schemas don't accidentally drop each other's partitions.
	rows, err := p.Pool.Query(ctx, `
		SELECT n.nspname, c.relname
		FROM pg_inherits i
		JOIN pg_class c ON c.oid = i.inhrelid
		JOIN pg_class pp ON pp.oid = i.inhparent
		JOIN pg_namespace n ON n.oid = c.relnamespace
		WHERE pp.relname = $1 AND n.nspname = current_schema()
	`, p.ParentTable)
	if err != nil {
		// Context cancellation (shutdown mid-cleanup) is normal; only genuine
		// query failures should surface as errors.
		if !isShutdownErr(err) {
			p.ErrorFn("error listing partitions for cleanup", err)
		}
		return
	}
	defer rows.Close()

	type partition struct{ schema, name string }
	var partitions []partition
	for rows.Next() {
		var pt partition
		if err := rows.Scan(&pt.schema, &pt.name); err != nil {
			continue
		}
		partitions = append(partitions, pt)
	}
	// Close rows before issuing further queries on the same pool.
	rows.Close()

	for _, pt := range partitions {
		partDate, ok := parsePartitionDate(pt.name, p.ParentTable)
		if !ok {
			continue
		}
		if partDate.Before(cutoff) {
			qualified := quoteIdent(pt.schema) + "." + quoteIdent(pt.name)
			if _, err := p.Pool.Exec(ctx, "DROP TABLE IF EXISTS "+qualified); err != nil {
				if !isShutdownErr(err) {
					p.ErrorFn("error dropping old partition "+qualified, err)
				}
			}
		}
	}
}

// quoteIdent returns a double-quoted PostgreSQL identifier with any embedded
// double-quotes escaped by doubling them, per the SQL standard.
func quoteIdent(s string) string {
	return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
}

// Run starts a ticker that calls EnsureLookaheadPartitions and
// DropOldPartitions on each CleanupInterval tick. Returns when ctx is
// cancelled or closeCh is closed.
func (p *Partitioner) Run(ctx context.Context, closeCh <-chan struct{}) {
	ticker := time.NewTicker(p.CleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-closeCh:
			return
		case <-ctx.Done():
			return
		case <-ticker.C:
			if err := p.EnsureLookaheadPartitions(ctx); err != nil {
				if !isShutdownErr(err) {
					p.ErrorFn("error ensuring lookahead partitions", err)
				}
			}
			p.DropOldPartitions(ctx)
		}
	}
}

// parsePartitionDate parses the date suffix from a child partition name
// of the form {ParentTable}_{YYYY}_{MM}_{DD}. Returns (time, true) on
// success; (zero, false) on any parse failure (too-short name, non-date
// suffix, invalid calendar date, etc.). Failures are silent — they
// indicate a partition that doesn't follow the expected naming convention
// and should be left alone by automatic cleanup.
//
// The parentTable argument is currently unused by the parser (the date
// suffix is interpreted positionally as the last three underscore-separated
// components). It is retained in the signature so future refinements can
// validate the prefix without a breaking change.
func parsePartitionDate(partName, _ string) (time.Time, bool) {
	parts := strings.Split(partName, "_")
	if len(parts) < 3 {
		return time.Time{}, false
	}
	dateStr := strings.Join(parts[len(parts)-3:], "-")
	t, err := time.Parse("2006-01-02", dateStr)
	if err != nil {
		return time.Time{}, false
	}
	return t, true
}