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

// Package probe is the master-side history store of item 7 (probe merge):
// it persists the per-minute metric rows and the whole-machine traffic
// ledger reported by agents (B1) into a dedicated SQLite file, probe.db,
// beside the panel database. The model, the booking rules and the
// roll-up/prune discipline are a port of monitor-probe's hub (src/db.rs,
// src/agent_ws.rs); see ITEM7-SCOPE.md §3.2/§3.3.
//
// probe.db is deliberately a separate file from x-ui.db: history writes,
// roll-ups and prunes then never contend for the panel database's write
// lock (ITEM7-SCOPE §7 risk 4). Everything here is best-effort by design:
// a broken or missing probe.db must never fail a heartbeat or a panel
// request, so the facade entry points log and swallow errors.
package probe

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/config"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"

	_ "github.com/mattn/go-sqlite3"
)

// DetailDays is how long minute rows are kept before they may be pruned;
// it only ever bounds the minute tier — the hourly tier and the traffic
// ledger live by the retention window itself. Mirrors hub DETAIL_DAYS.
const DetailDays = 7

// lateSeconds is the grace a just-closed hour gets before roll_up folds
// it, so rows landing slightly late still make it into their hour.
// Mirrors hub LATE (3600 + 60).
const lateSeconds = 3600 + 60

// rolledKey is the probe_setting key holding the roll-up watermark: the
// start (unix seconds) of the next hour not yet folded into the hourly
// tier. Mirrors hub ROLLED ("history_rolled").
const rolledKey = "history_rolled"

const schema = `
CREATE TABLE IF NOT EXISTS probe_setting (
  key   TEXT PRIMARY KEY,
  value TEXT NOT NULL
);

-- Per-node probe metadata (ITEM7-SCOPE §5.2). B2 creates the table and
-- reads traffic_reset_day; the editing surface lands in B3.
CREATE TABLE IF NOT EXISTS probe_meta (
  agent_id          INTEGER PRIMARY KEY,
  price             REAL    NOT NULL DEFAULT 0,
  currency          TEXT    NOT NULL DEFAULT 'USD',
  billing_cycle     TEXT    NOT NULL DEFAULT 'monthly',
  expires_at        TEXT,
  traffic_limit     INTEGER NOT NULL DEFAULT 0,
  traffic_mode      TEXT    NOT NULL DEFAULT 'sum',
  traffic_reset_day INTEGER NOT NULL DEFAULT 1,
  public            INTEGER NOT NULL DEFAULT 1,
  sort              INTEGER NOT NULL DEFAULT 0,
  group_name        TEXT    NOT NULL DEFAULT '',
  public_remark     TEXT    NOT NULL DEFAULT '',
  notify            INTEGER NOT NULL DEFAULT 0,
  country_pin       TEXT    NOT NULL DEFAULT '',
  ipv4_pin          TEXT    NOT NULL DEFAULT '',
  ipv6_pin          TEXT    NOT NULL DEFAULT '',
  -- V2.0.1 probe batch: operator traffic calibration (GB, either
  -- sign) and the counted-interface pin ('' = automatic selection).
  traffic_cal_up_gb   REAL NOT NULL DEFAULT 0,
  traffic_cal_down_gb REAL NOT NULL DEFAULT 0,
  stat_iface          TEXT NOT NULL DEFAULT ''
);

-- Monotonic byte counters that survive agent reboots and master
-- restarts; the whole-machine ledger (hub "traffic"). One row per agent.
CREATE TABLE IF NOT EXISTS probe_traffic (
  agent_id    INTEGER PRIMARY KEY,
  boot_id     TEXT    NOT NULL DEFAULT '',
  last_rx     INTEGER NOT NULL DEFAULT 0,
  last_tx     INTEGER NOT NULL DEFAULT 0,
  total_rx    INTEGER NOT NULL DEFAULT 0,
  total_tx    INTEGER NOT NULL DEFAULT 0,
  month_rx    INTEGER NOT NULL DEFAULT 0,
  month_tx    INTEGER NOT NULL DEFAULT 0,
  month_start TEXT    NOT NULL DEFAULT '',
  day_rx      INTEGER NOT NULL DEFAULT 0,
  day_tx      INTEGER NOT NULL DEFAULT 0,
  day_start   TEXT    NOT NULL DEFAULT ''
);

-- One row per agent per minute (hub "metric"). net_rx/net_tx are bytes
-- per second derived from the kernel counters, the *_max columns are the
-- minute's peaks of the agent-reported figures.
CREATE TABLE IF NOT EXISTS probe_metric (
  agent_id   INTEGER NOT NULL,
  ts         INTEGER NOT NULL,
  cpu        REAL    NOT NULL,
  mem_used   INTEGER NOT NULL,
  swap_used  INTEGER NOT NULL,
  disk_used  INTEGER NOT NULL,
  net_rx     INTEGER NOT NULL,
  net_tx     INTEGER NOT NULL,
  tcp        INTEGER NOT NULL,
  udp        INTEGER NOT NULL,
  procs      INTEGER NOT NULL,
  net_rx_max INTEGER NOT NULL DEFAULT 0,
  net_tx_max INTEGER NOT NULL DEFAULT 0,
  cpu_max    REAL    NOT NULL DEFAULT 0,
  PRIMARY KEY (agent_id, ts)
) WITHOUT ROWID;

-- The hourly tier (hub "metric_hour"): every metric column plus the
-- count of minute rows folded in, which weights later merges.
CREATE TABLE IF NOT EXISTS probe_metric_hour (
  agent_id   INTEGER NOT NULL,
  ts         INTEGER NOT NULL,
  minutes    INTEGER NOT NULL,
  cpu        REAL    NOT NULL,
  cpu_max    REAL    NOT NULL,
  mem_used   INTEGER NOT NULL,
  swap_used  INTEGER NOT NULL,
  disk_used  INTEGER NOT NULL,
  net_rx     INTEGER NOT NULL,
  net_tx     INTEGER NOT NULL,
  net_rx_max INTEGER NOT NULL,
  net_tx_max INTEGER NOT NULL,
  tcp        INTEGER NOT NULL,
  udp        INTEGER NOT NULL,
  procs      INTEGER NOT NULL,
  PRIMARY KEY (agent_id, ts)
) WITHOUT ROWID;

-- Probe results (hub "ping_record" / "ping_hour"). B2 creates the
-- tables so roll_up/prune cover them uniformly; B4 writes the rows.
-- latency < 0 is a lost probe; the hourly latency is the median of the
-- hour's answers and NULL when none arrived.
CREATE TABLE IF NOT EXISTS probe_ping_record (
  agent_id INTEGER NOT NULL,
  ts       INTEGER NOT NULL,
  task_id  INTEGER NOT NULL,
  latency  INTEGER NOT NULL,
  PRIMARY KEY (agent_id, ts, task_id)
) WITHOUT ROWID;

CREATE TABLE IF NOT EXISTS probe_ping_hour (
  agent_id INTEGER NOT NULL,
  ts       INTEGER NOT NULL,
  task_id  INTEGER NOT NULL,
  answered INTEGER NOT NULL,
  lost     INTEGER NOT NULL,
  latency  INTEGER,
  lo       INTEGER,
  hi       INTEGER,
  PRIMARY KEY (agent_id, ts, task_id)
) WITHOUT ROWID;

-- Probe tasks and their per-agent assignment (hub "ping_task" /
-- "ping_node"). B4 adds both. A task is one operator-maintained
-- target — there is no preset list; region is the operator's own
-- grouping label (empty = ungrouped). Ids come from a high-water
-- mark in probe_setting so a deleted task's id is never reissued
-- (hub next_id discipline): in-flight results and an editor opened
-- before the delete must not land on a different task.
CREATE TABLE IF NOT EXISTS probe_ping_task (
  id       INTEGER PRIMARY KEY,
  name     TEXT    NOT NULL,
  target   TEXT    NOT NULL,
  interval INTEGER NOT NULL DEFAULT 60,
  auto_join INTEGER NOT NULL DEFAULT 0,
  sort     INTEGER NOT NULL DEFAULT 0,
  region   TEXT    NOT NULL DEFAULT ''
);

CREATE TABLE IF NOT EXISTS probe_ping_node (
  task_id  INTEGER NOT NULL,
  agent_id INTEGER NOT NULL,
  PRIMARY KEY (task_id, agent_id)
);
`

// Store is one open probe.db. The zero value is not usable; call Open.
// A Store is safe for concurrent use: the single connection serializes
// writes, and the ingest state is mutex-guarded.
type Store struct {
	db *sql.DB

	mu    sync.Mutex // guards nodes
	nodes map[int64]*nodeState
}

// Open creates (if needed) and initializes the probe database at path.
func Open(path string) (*Store, error) {
	if dir := filepath.Dir(path); dir != "" {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return nil, err
		}
	}
	dsn := path + "?_journal_mode=WAL&_busy_timeout=10000&_synchronous=NORMAL&_txlock=immediate"
	db, err := sql.Open("sqlite3", dsn)
	if err != nil {
		return nil, err
	}
	// One connection: writes serialize on it (hub's single write
	// connection), and WAL still lets the read side (B3) proceed.
	db.SetMaxOpenConns(1)
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("probe schema: %w", err)
	}
	if err := migrateMetaColumns(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("probe schema migration: %w", err)
	}
	return &Store{db: db, nodes: make(map[int64]*nodeState)}, nil
}

// metaColumnMigrations are the probe_meta columns added after the
// table's first release; CREATE TABLE IF NOT EXISTS never retrofits
// them onto an existing database, so they are added here, once, by
// name check (SQLite has no ADD COLUMN IF NOT EXISTS).
var metaColumnMigrations = []struct{ name, ddl string }{
	{"traffic_cal_up_gb", `ALTER TABLE probe_meta ADD COLUMN traffic_cal_up_gb REAL NOT NULL DEFAULT 0`},
	{"traffic_cal_down_gb", `ALTER TABLE probe_meta ADD COLUMN traffic_cal_down_gb REAL NOT NULL DEFAULT 0`},
	{"stat_iface", `ALTER TABLE probe_meta ADD COLUMN stat_iface TEXT NOT NULL DEFAULT ''`},
}

func migrateMetaColumns(db *sql.DB) error {
	rows, err := db.Query(`PRAGMA table_info(probe_meta)`)
	if err != nil {
		return err
	}
	have := map[string]bool{}
	for rows.Next() {
		var cid int
		var name, typ string
		var notNull int
		var dflt sql.NullString
		var pk int
		if err := rows.Scan(&cid, &name, &typ, &notNull, &dflt, &pk); err != nil {
			rows.Close()
			return err
		}
		have[name] = true
	}
	if err := rows.Close(); err != nil {
		return err
	}
	for _, m := range metaColumnMigrations {
		if have[m.name] {
			continue
		}
		if _, err := db.Exec(m.ddl); err != nil {
			return err
		}
	}
	return nil
}

// Close releases the database handle.
func (s *Store) Close() error {
	if s == nil || s.db == nil {
		return nil
	}
	return s.db.Close()
}

// DB exposes the raw handle for the read side (B3) and tests.
func (s *Store) DB() *sql.DB { return s.db }

// --- default store (lazy, failure-tolerant) ---

var def struct {
	mu      sync.Mutex
	store   *Store
	retryAt time.Time
}

// DefaultPath is where probe.db lives: beside the panel database.
func DefaultPath() string {
	return filepath.Join(config.GetDBFolderPath(), "probe.db")
}

// Default returns the process-wide store, opening it on first use. A
// failed open is retried at most once a minute and reported as nil, so
// callers degrade to "no history" instead of failing.
func Default() *Store {
	def.mu.Lock()
	defer def.mu.Unlock()
	if def.store != nil {
		return def.store
	}
	if time.Now().Before(def.retryAt) {
		return nil
	}
	st, err := Open(DefaultPath())
	if err != nil {
		logger.Warningf("probe: open %s failed (history disabled until retry): %v", DefaultPath(), err)
		def.retryAt = time.Now().Add(time.Minute)
		return nil
	}
	def.store = st
	return st
}