3x-ui-agent/internal/web/service/probe/maintenance.go

package probe

import (
	"database/sql"
	"math"
	"sort"
	"strconv"
	"time"
)

// watermark reads the roll-up watermark (unix seconds, hour-aligned):
// the next hour not yet folded into the hourly tier.
func (s *Store) watermark() (int64, bool, error) {
	var v string
	err := s.db.QueryRow(`SELECT value FROM probe_setting WHERE key = ?`, rolledKey).Scan(&v)
	if err == sql.ErrNoRows {
		return 0, false, nil
	}
	if err != nil {
		return 0, false, err
	}
	n, perr := strconv.ParseInt(v, 10, 64)
	if perr != nil {
		return 0, false, nil
	}
	return n, true, nil
}

func (s *Store) setWatermark(tx *sql.Tx, hour int64) error {
	_, err := tx.Exec(`INSERT OR REPLACE INTO probe_setting (key, value) VALUES (?, ?)`,
		rolledKey, strconv.FormatInt(hour, 10))
	return err
}

func hourFloor(ts int64) int64 {
	r := ts % 3600
	if r < 0 {
		r += 3600
	}
	return ts - r
}

// oldestStamp is the earliest ts across the minute-tier tables, used
// to initialize the watermark on first roll_up (hub `oldest`).
func (s *Store) oldestStamp() (int64, bool, error) {
	var oldest sql.NullInt64
	err := s.db.QueryRow(`
		SELECT MIN(ts) FROM (
		  SELECT ts FROM probe_metric
		  UNION ALL SELECT ts FROM probe_ping_record
		)`).Scan(&oldest)
	if err != nil {
		return 0, false, err
	}
	return oldest.Int64, oldest.Valid, nil
}

// RollUp folds closed hours from the minute tier into the hourly tier
// (hub roll_up): the watermark advances one hour per transaction, so a
// failure leaves the hour to be folded again rather than half folded,
// and INSERT OR REPLACE makes a second fold of the same hour a rewrite.
// Hours older than the retention floor are skipped (the watermark
// jumps to the floor), and the most recent closed hour waits out the
// late grace. Returns the number of hours folded.
func (s *Store) RollUp(now int64, keepDays int64) (int, error) {
	if _, ok, err := s.watermark(); err != nil {
		return 0, err
	} else if !ok {
		from := now
		if oldest, valid, err := s.oldestStamp(); err != nil {
			return 0, err
		} else if valid {
			from = oldest
		}
		tx, err := s.db.Begin()
		if err != nil {
			return 0, err
		}
		if err := s.setWatermark(tx, hourFloor(from)); err != nil {
			tx.Rollback()
			return 0, err
		}
		if err := tx.Commit(); err != nil {
			return 0, err
		}
	}
	floor := hourFloor(now - keepDays*86400)
	folded := 0
	for {
		ok, err := s.foldNext(floor, now-3600-lateSeconds)
		if err != nil {
			return folded, err
		}
		if !ok {
			return folded, nil
		}
		folded++
	}
}

// foldNext folds the hour at the watermark (or at floor when the
// watermark is older) and moves the watermark past it, in one
// transaction. False once the hour is past last.
func (s *Store) foldNext(floor, last int64) (bool, error) {
	tx, err := s.db.Begin()
	if err != nil {
		return false, err
	}
	defer tx.Rollback()

	var wmVal string
	err = tx.QueryRow(`SELECT value FROM probe_setting WHERE key = ?`, rolledKey).Scan(&wmVal)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	wm, perr := strconv.ParseInt(wmVal, 10, 64)
	if perr != nil {
		return false, nil
	}
	hour := wm
	if floor > hour {
		hour = floor
	}
	if hour > last {
		return false, nil
	}

	// Metric fold: means weighted naturally by COUNT, byte columns
	// truncate their mean (a loss under one byte), the count columns
	// round (hub fold_next).
	if _, err := tx.Exec(`
		INSERT OR REPLACE INTO probe_metric_hour
		  (agent_id, ts, minutes, cpu, cpu_max, mem_used, swap_used, disk_used,
		   net_rx, net_tx, net_rx_max, net_tx_max, tcp, udp, procs)
		SELECT agent_id, ?, COUNT(*), AVG(cpu), MAX(MAX(cpu, cpu_max)),
		       CAST(AVG(mem_used) AS INTEGER), CAST(AVG(swap_used) AS INTEGER),
		       CAST(AVG(disk_used) AS INTEGER), CAST(AVG(net_rx) AS INTEGER),
		       CAST(AVG(net_tx) AS INTEGER), MAX(MAX(net_rx, net_rx_max)),
		       MAX(MAX(net_tx, net_tx_max)), CAST(ROUND(AVG(tcp)) AS INTEGER),
		       CAST(ROUND(AVG(udp)) AS INTEGER), CAST(ROUND(AVG(procs)) AS INTEGER)
		FROM probe_metric WHERE ts >= ? AND ts < ? + 3600
		GROUP BY agent_id`, hour, hour, hour); err != nil {
		return false, err
	}

	// Ping fold: per agent and task, the hour's median/lo/hi over the
	// answered results and the answered/lost counts (hub Tally).
	pingRows, err := tx.Query(`
		SELECT agent_id, task_id, latency FROM probe_ping_record
		WHERE ts >= ? AND ts < ? + 3600 ORDER BY agent_id, task_id`, hour, hour)
	if err != nil {
		return false, err
	}
	type tallyKey struct{ agent, task int64 }
	tallies := make(map[tallyKey]*pingTally)
	var order []tallyKey
	for pingRows.Next() {
		var agent, task, latency int64
		if err := pingRows.Scan(&agent, &task, &latency); err != nil {
			pingRows.Close()
			return false, err
		}
		k := tallyKey{agent, task}
		t, ok := tallies[k]
		if !ok {
			t = &pingTally{}
			tallies[k] = t
			order = append(order, k)
		}
		t.add(latency)
	}
	if err := pingRows.Err(); err != nil {
		pingRows.Close()
		return false, err
	}
	pingRows.Close()
	for _, k := range order {
		t := tallies[k]
		median, lo, hi := t.medianLoHi()
		if _, err := tx.Exec(`
			INSERT OR REPLACE INTO probe_ping_hour
			  (agent_id, ts, task_id, answered, lost, latency, lo, hi)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
			k.agent, hour, k.task, t.answered, t.lost,
			nullOr(median), nullOr(lo), nullOr(hi)); err != nil {
			return false, err
		}
	}

	if err := s.setWatermark(tx, hour+3600); err != nil {
		return false, err
	}
	if err := tx.Commit(); err != nil {
		return false, err
	}
	return true, nil
}

func nullOr(v *int64) any {
	if v == nil {
		return nil
	}
	return *v
}

// pingTally accumulates one hour of probe results for one task
// (hub Tally over single results).
type pingTally struct {
	answers  []int64
	lost     int64
	answered int64
	lo, hi   *int64
}

func (t *pingTally) add(latency int64) {
	if latency < 0 {
		t.lost++
		return
	}
	t.answered++
	t.answers = append(t.answers, latency)
	if t.lo == nil || latency < *t.lo {
		v := latency
		t.lo = &v
	}
	if t.hi == nil || latency > *t.hi {
		v := latency
		t.hi = &v
	}
}

// medianLoHi returns the median of the answers (the middle two
// averaged) and the lo/hi bounds; all nil when nothing answered.
func (t *pingTally) medianLoHi() (median, lo, hi *int64) {
	if t.answered == 0 {
		return nil, nil, nil
	}
	sort.Slice(t.answers, func(i, j int) bool { return t.answers[i] < t.answers[j] })
	n := int64(len(t.answers))
	mid := (t.answers[(n-1)/2] + t.answers[n/2]) / 2
	return &mid, t.lo, t.hi
}

// Prune drops history beyond the retention window (hub prune): minute
// rows past DetailDays (or the window itself when shorter), hourly
// rows past keepDays. Minute rows are only deleted up to the roll-up
// watermark — a row outlives its window until its hour is folded, so a
// roll-up that has fallen behind costs disk rather than history.
// Traffic totals live in their own table precisely so history can be
// pruned freely. Deletes are per agent and chunked (a day per
// statement for the minute tier), releasing the write lock in between.
// Returns the number of rows deleted.
func (s *Store) Prune(keepDays int64) (int, error) {
	now := time.Now().Unix()
	detail := keepDays
	if detail > DetailDays {
		detail = DetailDays
	}
	minutesCutoff := now - detail*86400
	hoursCutoff := now - keepDays*86400

	ids, err := s.agentIDs()
	if err != nil {
		return 0, err
	}
	pruned := 0
	type tableRule struct {
		table   string
		cutoff  int64
		guarded bool
	}
	rules := []tableRule{
		{"probe_metric", minutesCutoff, true},
		{"probe_ping_record", minutesCutoff, true},
		{"probe_metric_hour", hoursCutoff, false},
		{"probe_ping_hour", hoursCutoff, false},
	}
	for _, id := range ids {
		for _, rule := range rules {
			for {
				before := rule.cutoff
				if rule.guarded {
					// Guarded by the watermark: with none set, nothing
					// is deletable (hub reads it as the minimum).
					wm, ok, err := s.watermark()
					if err != nil {
						return pruned, err
					}
					if !ok {
						before = math.MinInt64
					} else if wm < before {
						before = wm
					}
				}
				var first sql.NullInt64
				err := s.db.QueryRow(
					`SELECT MIN(ts) FROM `+rule.table+` WHERE agent_id = ?`, id).Scan(&first)
				if err != nil {
					return pruned, err
				}
				if !first.Valid || first.Int64 >= before {
					break
				}
				limit := before
				if rule.guarded && first.Int64+86400 < limit {
					limit = first.Int64 + 86400
				}
				res, err := s.db.Exec(
					`DELETE FROM `+rule.table+` WHERE agent_id = ? AND ts < ?`, id, limit)
				if err != nil {
					return pruned, err
				}
				n, _ := res.RowsAffected()
				pruned += int(n)
				if n == 0 {
					break
				}
			}
		}
	}
	return pruned, nil
}

// agentIDs lists every agent with rows in any history table.
func (s *Store) agentIDs() ([]int64, error) {
	rows, err := s.db.Query(`
		SELECT agent_id FROM probe_metric
		UNION SELECT agent_id FROM probe_ping_record
		UNION SELECT agent_id FROM probe_metric_hour
		UNION SELECT agent_id FROM probe_ping_hour
		ORDER BY 1`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var ids []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	return ids, rows.Err()
}