worker/app/models/check_aggregator.go

package models

import (
	"context"
	"log"
	"time"

	"gorm.io/gorm"
)

// Phase 3 of docs/todo.md — result aggregation for multi-region checks.
//
// When a check has RequireQuorum > 1, ApplyRemoteCheckResult does not
// write Check.State directly. Instead it appends a CheckRegionResult row
// and leaves AggregatedAt NULL. This file owns the background goroutine
// that walks those pending rows once their aggregation window has
// elapsed, decides OK/ERR/DEGRADED per the documented rule, writes the
// aggregate state onto Check, stamps AggregatedAt on the contributing
// rows, and triggers Monitor.UpdateStatusFromChecks so the monitor's own
// status follows.
//
// Aggregation rule (see docs/todo.md Phase 3 + checkSeverityRank in
// monitor.go for the corresponding severity order):
//
//   - Aggregate only rows whose created_at is older than
//     NOW() - Check.AggregationWindowSeconds. This is the "watermark"
//     pattern: a row is eligible only when no fresher regional result
//     could still arrive and tip the vote. The window is per-check so
//     noisy checks can use a longer wait than fast ones.
//   - If zero eligible rows exist for a check, leave Check.State
//     untouched (the special case called out in the spec).
//   - Otherwise count OK vs not-OK among the eligible rows:
//       OK      >= RequireQuorum → Check.State = OK
//       not-OK  >= RequireQuorum → Check.State = ERR
//       neither side reaches quorum → Check.State = DEGRADED
//   - Stamp AggregatedAt = NOW() on every contributing row so the next
//     tick skips them. One transaction per check; per-row failures do
//     not poison other checks.

// AggregatorTickInterval is the default cadence of StartCheckAggregator
// when the caller passes interval <= 0. Mirrors the 30s default used by
// the other reapers in this package so the three reapers all tick on
// the same wall clock cadence — easier to grep, easier to reason about
// in incident timelines.
const AggregatorTickInterval = 30 * time.Second

// EnsureCheckAggregatorIndexes adds the partial indexes the aggregator
// relies on. AutoMigrate creates AggregatedAt as a regular btree column,
// but the per-tick SELECT filters on `aggregated_at IS NULL` over what
// grows to be a busy table; a partial index keeps the working set
// tiny. Idempotent so it is safe to call from Migrate() and from tests.
func EnsureCheckAggregatorIndexes() error {
	return DB().Exec(`
		CREATE INDEX IF NOT EXISTS check_region_results_pending_idx
		    ON check_region_results (check_id, created_at)
		    WHERE aggregated_at IS NULL
	`).Error
}

// aggregateCheckState holds the per-check aggregation inputs we need to
// keep the rule readable. Rows is the set of CheckRegionResult rows
// eligible for the current decision; quorum is Check.RequireQuorum.
type aggregateCheckState struct {
	CheckID int64
	Quorum  int
	Rows    []CheckRegionResult
}

// decideAggregateState encodes the OK/ERR/DEGRADED rule described in
// the package doc. Pure function — no DB, no time — so it is trivially
// unit-testable from the test file.
func decideAggregateState(in aggregateCheckState) (string, bool) {
	if len(in.Rows) == 0 || in.Quorum <= 1 {
		// Zero eligible rows in the window OR a misconfigured check
		// (QuorumEnabled false). Caller must leave Check.State alone
		// in both cases.
		return "", false
	}
	okCount := 0
	badCount := 0
	for i := range in.Rows {
		if in.Rows[i].State == stateOK {
			okCount++
		} else {
			badCount++
		}
	}
	switch {
	case okCount >= in.Quorum:
		return stateOK, true
	case badCount >= in.Quorum:
		return stateERR, true
	default:
		return stateDegraded, true
	}
}

// CheckAggregatorTick performs one pass of the aggregator. It is the
// per-tick body StartCheckAggregator calls. Exported so the test suite
// can call it directly without spinning up the goroutine; production
// always goes through StartCheckAggregator.
//
// The returned (aggregated, err) tuple lets the caller log a metric:
// aggregated counts how many Check rows had their State written this
// tick. The function is idempotent — a second call with no new
// unaggregated rows is a no-op that returns (0, nil).
func CheckAggregatorTick() (aggregated int, err error) {
	// Step 1: collect candidate check IDs. The JOIN to checks is needed
	// to read each check's window length and to filter on
	// require_quorum > 1 (so we never aggregate the legacy path).
	rows, err := DB().Raw(`
		SELECT DISTINCT crr.check_id
		FROM check_region_results crr
		JOIN checks c ON c.id = crr.check_id
		WHERE crr.aggregated_at IS NULL
		  AND c.require_quorum > 1
		  AND crr.created_at < NOW() - make_interval(secs => GREATEST(c.aggregation_window_seconds, 1))
		ORDER BY crr.check_id
	`).Rows()
	if err != nil {
		return 0, err
	}
	defer func() { _ = rows.Close() }()

	var checkIDs []int64
	for rows.Next() {
		var id int64
		if scanErr := rows.Scan(&id); scanErr != nil {
			return 0, scanErr
		}
		checkIDs = append(checkIDs, id)
	}
	if scanErr := rows.Err(); scanErr != nil {
		return 0, scanErr
	}
	if len(checkIDs) == 0 {
		return 0, nil
	}

	for _, checkID := range checkIDs {
		n, err := aggregateOneCheck(checkID)
		if err != nil {
			// Log and continue: one bad check must not stop the loop.
			log.Printf("check_aggregator: check_id=%d error: %v", checkID, err)
			continue
		}
		aggregated += n
	}
	return aggregated, nil
}

// aggregateOneCheck runs the aggregation logic for a single check inside
// a transaction. The transaction holds a FOR UPDATE row lock on the
// check so concurrent aggregator instances (multiple web processes) can
// not race on the same check — the second one waits for the first to
// commit, then sees AggregatedAt IS NOT NULL on every row and the
// candidate SELECT below returns an empty set.
func aggregateOneCheck(checkID int64) (int, error) {
	tx := DB().Begin()
	if tx.Error != nil {
		return 0, tx.Error
	}
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
			panic(r)
		}
	}()

	var check Check
	if err := tx.Clauses(SkipLockedClause).First(&check, checkID).Error; err != nil {
		tx.Rollback()
		if err == gorm.ErrRecordNotFound {
			// Check was deleted between candidate SELECT and lock; not
			// an error, just nothing to do.
			return 0, nil
		}
		return 0, err
	}
	// Defensive: only aggregate quorum-enabled checks. The candidate
	// SELECT already filters on this, but a stale row that flipped off
	// quorum between calls must not be re-aggregated here.
	if !check.QuorumEnabled() {
		tx.Rollback()
		return 0, nil
	}

	var results []CheckRegionResult
	if err := tx.
		Where("check_id = ? AND aggregated_at IS NULL", checkID).
		Order("created_at ASC").
		Find(&results).Error; err != nil {
		tx.Rollback()
		return 0, err
	}

	decision, ok := decideAggregateState(aggregateCheckState{
		CheckID: checkID,
		Quorum:  check.RequireQuorum,
		Rows:    results,
	})
	if !ok {
		// Zero eligible rows — leave Check.State alone. There is also
		// nothing to stamp, so just rollback and move on.
		tx.Rollback()
		return 0, nil
	}

	now := time.Now()
	// Pull the latest error string from the contributing rows so the
	// monitor event / notifier pipeline has something to show. Prefer
	// the most recent ERR row's message; fall back to the most recent
	// any-row message. UNK / empty stays NULL.
	var lastError *string
	for i := len(results) - 1; i >= 0; i-- {
		if results[i].Error != nil && *results[i].Error != "" {
			lastError = results[i].Error
			break
		}
	}

	upd := map[string]interface{}{
		colState:   decision,
		colLastEnd: now,
	}
	if decision == stateOK {
		// OK resets error — mirrors the legacy ApplyRemoteCheckResult
		// path that sets `error = gorm.Expr("NULL")` when state==OK.
		upd["error"] = gorm.Expr("NULL")
		upd["last_ok"] = now
		upd["fails"] = 0
		upd["was_up"] = now
	} else {
		// Non-OK: bump the fail counter and only overwrite the error
		// when one of the contributing rows actually carries a
		// message. If none do, leave whatever was there before —
		// mirrors the legacy `if report.Error != nil` branch.
		upd["last_fail"] = now
		upd["fails"] = gorm.Expr("fails + 1")
		if lastError != nil {
			upd["error"] = *lastError
		}
	}

	if err := tx.Model(&Check{}).Where("id = ?", checkID).UpdateColumns(upd).Error; err != nil {
		tx.Rollback()
		return 0, err
	}

	if err := tx.Model(&CheckRegionResult{}).
		Where("check_id = ? AND aggregated_at IS NULL", checkID).
		UpdateColumns(map[string]interface{}{
			"aggregated_at": now,
		}).Error; err != nil {
		tx.Rollback()
		return 0, err
	}

	if err := tx.Commit().Error; err != nil {
		return 0, err
	}

	// Mirror ApplyRemoteCheckResult: propagate the aggregate decision
	// up to the monitor. We do this AFTER commit so a rollback does
	// not leave the monitor in a state whose corresponding check is
	// still pre-aggregate. The goroutine keeps the failure path of
	// UpdateStatusFromChecks isolated from the aggregator's hot loop.
	if check.MonitorID != 0 {
		var mon Monitor
		if err := DB().First(&mon, check.MonitorID).Error; err == nil {
			go mon.UpdateStatusFromChecks()
		} else {
			log.Printf("check_aggregator: monitor lookup failed for check_id=%d: %v", checkID, err)
		}
	}
	return 1, nil
}

// StartCheckAggregator launches a goroutine that calls
// CheckAggregatorTick on the given interval until ctx is canceled.
// Mirrors StartTaskReaper / StartDeadWorkerReaper in this package — same
// ticker shape, same default-interval fall-back, same per-tick recover
// so a malformed row cannot crash the web process.
//
// The default interval is AggregatorTickInterval (30s); values <= 0
// fall back to the default so the helper is safe to call from any call
// site without a guard. A nil context falls back to context.Background()
// the same way StartDeadWorkerReaper does, so main.init() and tests
// can both call it without ceremony.
//
// Wire from main.init() once per process. The aggregator is cheap in
// steady state (one indexed SELECT for candidates + a per-check
// transaction over a handful of unaggregated rows). Under load it
// scales horizontally — multiple web processes can each run their own
// StartCheckAggregator goroutine because FOR UPDATE SKIP LOCKED on the
// per-check transaction guarantees at-most-one winner per check.
func StartCheckAggregator(ctx context.Context, interval time.Duration) {
	if interval <= 0 {
		interval = AggregatorTickInterval
	}
	if ctx == nil {
		ctx = context.Background()
	}
	// Best-effort index bootstrap. AutoMigrate declares AggregatedAt as
	// a regular btree column; the partial index speeds up the per-tick
	// candidate SELECT. Idempotent — safe to call on every boot.
	if err := EnsureCheckAggregatorIndexes(); err != nil {
		log.Printf("check_aggregator: ensure index: %v", err)
	}
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				func() {
					defer func() {
						if r := recover(); r != nil {
							log.Printf("check_aggregator: panic recovered: %v", r)
						}
					}()
					n, err := CheckAggregatorTick()
					if err != nil {
						log.Printf("check_aggregator: error: %v", err)
						return
					}
					if n > 0 {
						log.Printf("check_aggregator: aggregated=%d", n)
					}
				}()
			}
		}
	}()
}