worker/app/models/check_jobs.go

package models

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/google/uuid"
	"github.com/lib/pq"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"rocketgit.ru/rsmon/worker/internal/influx"
	"rocketgit.ru/rsmon/worker/internal/wire"
)

// defaultRegionCode is the historical default region seeded by Migrate()
// (see app/models/migrate.go) and used as a catch-all bucket for results
// reported without a region code. Aliased to the exported Region
// constant (DefaultRegionCode) so admin endpoints and the in-process
// job router share one source of truth.
const defaultRegionCode = DefaultRegionCode

// ChecksForWorker returns checks that need to be executed by a distributed worker.
// It uses FOR UPDATE SKIP LOCKED to prevent race conditions between concurrent workers.
// The worker specifies which check kinds it can handle via the kinds parameter.
//
// Phase 2 of docs/plans/worker-notifier-mvp.md adds regional job routing: when
// worker is non-nil, the candidate monitor set is filtered by
// applyRegionRouting so a worker only sees checks that explicitly allow its
// region. Pass nil for the legacy "no region scoping" path used by
// diagnostics/dashboard tooling.
func ChecksForWorker(worker *WorkerNode, kinds []string, limit int) []*Check {
	tx := DB().Begin()
	q := tx.Joins("JOIN monitors ON checks.monitor_id = monitors.id").
		Where("monitors.enabled").
		Where("checks.enabled AND checks.kind IN (?)", kinds)

	if worker != nil {
		q = applyRegionRouting(q, worker)
		if accounts := worker.AccessibleAccountIDs(); len(accounts) > 0 {
			q = q.Joins("JOIN groups ON monitors.group_id = groups.id").Where("groups.account_id IN (?)", accounts)
		}
		// Flagged workers remain visible for audit/history but never receive new work.
		if worker.NetworkProblemActive(time.Now()) {
			tx.Rollback()
			return nil
		}
	}

	// Allow faster retry for failed http/dns checks
	notOk := ""
	hasHTTPOrDNS := false
	for _, k := range kinds {
		if k == kindHTTP || k == kindDNS {
			hasHTTPOrDNS = true
			break
		}
	}
	if hasHTTPOrDNS {
		notOk = `OR (checks.state != 'OK' AND checks.last_start + '120 second'::interval < now())`
	}

	whereClause := `
		(checks.last_start IS NULL) OR
		(checks.last_start + (checks.interval || ' second')::interval < now())
	`
	if notOk != "" {
		whereClause += notOk
	}
	rq := q.Where(whereClause)

	// Use SKIP LOCKED to avoid contention between workers. The same
	// FOR UPDATE SKIP LOCKED clause also gives us implicit load balancing
	// across workers in the same region: each concurrent worker call
	// grabs a disjoint slice of the pending checks and a row leased by
	// worker A is invisible to worker B until A's transaction commits
	// (or rolls back / lease expires).
	var checks []*Check
	rq.Clauses(SkipLockedClause).
		Limit(limit).
		Preload("Monitor").
		Find(&checks)

	for _, c := range checks {
		log.Println("worker: assigned remote check:", c.ID, c.Kind)
		tx.Model(&c).Where("id = ?", c.ID).Update(colLastStart, time.Now())
	}

	tx.Commit()
	return checks
}

// EnqueueDueCheckTasks atomically turns due normal checks into durable generic
// task envelopes. ChecksForWorker remains for the HTTP polling compatibility
// endpoint, while websocket scheduling uses this task-producing path.
func EnqueueDueCheckTasks(worker *WorkerNode, kinds []string, limit int) error {
	if worker == nil || len(kinds) == 0 || limit <= 0 || worker.NetworkProblemActive(time.Now()) {
		return nil
	}
	return DB().Transaction(func(tx *gorm.DB) error {
		q := tx.Joins("JOIN monitors ON checks.monitor_id = monitors.id").
			Where("monitors.enabled").Where("checks.enabled AND checks.kind IN (?)", kinds)
		if accounts := worker.AccessibleAccountIDs(); len(accounts) > 0 {
			q = q.Joins("JOIN groups ON monitors.group_id = groups.id").Where("groups.account_id IN (?)", accounts)
		}
		q = applyRegionRouting(q, worker)
		where := `(checks.last_start IS NULL) OR (checks.last_start + (checks.interval || ' second')::interval < now())`
		for _, kind := range kinds {
			if kind == kindHTTP || kind == kindDNS {
				where += ` OR (checks.state != 'OK' AND checks.last_start + '120 second'::interval < now())`
				break
			}
		}
		var checks []*Check
		if err := q.Where(where).Clauses(SkipLockedClause).Limit(limit).Preload("Monitor.Group").Find(&checks).Error; err != nil {
			return err
		}
		now := time.Now()
		for _, check := range checks {
			if check.Monitor == nil || check.Monitor.Group == nil {
				continue
			}
			job := JobForCheck(check)
			payload, err := json.Marshal(job)
			if err != nil {
				return err
			}
			checkID, monitorID := check.ID, check.MonitorID
			bucket := now.UTC().Unix() / int64(check.Interval)
			task := Task{
				JobID: job.JobID, Kind: TaskKindCheck, State: TaskStateQueued,
				AccountID: check.Monitor.Group.AccountID, CheckID: &checkID, MonitorID: &monitorID,
				Payload: payload, NotBefore: now, MaxAttempts: DefaultTaskMaxAttempts,
				IdempotencyKey: fmt.Sprintf("check:%d:%d", check.ID, bucket),
			}
			if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&task).Error; err != nil {
				return err
			}
			if err := tx.Model(check).Update(colLastStart, now).Error; err != nil {
				return err
			}
		}
		return nil
	})
}

// applyRegionRouting narrows the monitor JOIN in ChecksForWorker to the
// subset whose routing rule matches the calling worker's region. The filter
// is applied at the SQL layer so the SKIP LOCKED page only scans/leases
// rows that this worker is allowed to run, instead of leasing and then
// discarding forbidden checks.
//
// The function intentionally mirrors Monitor.WantsRegion so the helper can
// be reused from non-SQL callers (UI preview, plan validation, etc.).
//
// SQL form:
//
//   - monitors.region_mode IN ('any', 'all', ”)
//     → unconditional match; legacy / Phase 3 placeholder behavior.
//   - monitors.region_mode = 'specific' AND
//     (monitors.preferred_regions IS NULL OR
//     cardinality(monitors.preferred_regions) = 0 OR
//     ? = ANY(monitors.preferred_regions))
//     → empty array falls back to "any"; otherwise the worker code must
//     be in the whitelist.
func applyRegionRouting(q *gorm.DB, worker *WorkerNode) *gorm.DB {
	if worker == nil || worker.RegionCode == "" {
		return q
	}
	if worker.RegionCode == defaultRegionCode {
		// The default "local" region is the historical catch-all; the
		// in-process scheduler (not ChecksForWorker) handles those
		// monitors. Skip regional filtering entirely so we don't leak
		// Phase 1 in-process workers through the new router.
		return q
	}
	// TODO(phase3): split RegionMode="all" into N assignments, one per
	// preferred region, so the result aggregator can build a quorum.
	// Today it is treated as "any" so existing checks keep flowing.
	return q.Where(
		`(monitors.region_mode IN ('any', 'all', '') OR `+
			`(monitors.region_mode = 'specific' AND `+
			`(monitors.preferred_regions IS NULL OR `+
			`coalesce(array_length(monitors.preferred_regions, 1), 0) = 0 OR `+
			`? = ANY(monitors.preferred_regions))))`,
		worker.RegionCode,
	)
}

// JobForCheck creates a CheckJob from a Check model for sending to a worker
func JobForCheck(c *Check) wire.CheckJob {
	jobID := uuid.New().String()
	var urlStr *string
	if c.URL != nil {
		urlStr = c.URL
	}
	return wire.CheckJob{
		JobID:     jobID,
		CheckID:   c.ID,
		MonitorID: c.MonitorID,
		Kind:      c.Kind,
		Host:      c.Monitor.Host,
		URL:       urlStr,
		Interval:  c.Interval,
		Settings:  json.RawMessage(c.Settings),
	}
}

// QueueMonitorChecks makes enabled checks for a monitor immediately eligible for remote assignment.
func QueueMonitorChecks(monitorID int64) error {
	return DB().Model(&Check{}).
		Where("monitor_id = ? AND enabled", monitorID).
		Updates(map[string]interface{}{
			colLastStart: nil,
			colLastEnd:   nil,
		}).Error
}

// QueueMonitorChecksKind makes enabled checks of one kind immediately eligible for remote assignment.
func QueueMonitorChecksKind(monitorID int64, kind string) error {
	return DB().Model(&Check{}).
		Where("monitor_id = ? AND kind = ? AND enabled", monitorID, kind).
		Updates(map[string]interface{}{
			colLastStart: nil,
			colLastEnd:   nil,
		}).Error
}

// ApplyRemoteCheckResult applies a check result reported by a distributed worker.
// It updates the check state in the database and triggers monitor status aggregation.
//
// Phase 3 of docs/todo.md (multi-region quorum aggregation): when the
// check has RequireQuorum > 1, the per-region result is recorded in
// check_region_results but Check.State is NOT touched here — that is
// the job of app/models/check_aggregator.go, which decides OK/ERR/
// DEGRADED once enough regional results have arrived or the aggregation
// window has elapsed. QuorumEnabled() == false preserves the legacy
// direct-update path so single-region / non-aggregated monitors keep
// the same behavior.
func ApplyRemoteCheckResult(report wire.CheckResultReport, regionCode string) error { //nolint:gocritic,lll // hugeParam: accepted for interface compatibility
	return ApplyRemoteCheckResultFromWorker(report, regionCode, nil)
}

// ApplyRemoteCheckResultFromWorker persists worker attribution before changing
// legacy check state. A confirmation result is consumed exactly once and never
// overwrites the original check result.
func ApplyRemoteCheckResultFromWorker(report wire.CheckResultReport, regionCode string, worker *WorkerNode) error { //nolint:gocritic,lll // hugeParam: wire compatibility
	var monitor *Monitor
	err := DB().Transaction(func(tx *gorm.DB) error {
		var err error
		monitor, err = ApplyRemoteCheckResultFromWorkerTx(tx, report, regionCode, worker)
		return err
	})
	if err != nil {
		return err
	}
	if monitor != nil {
		monitor.UpdateStatusFromChecks()
	}
	// VictoriaMetrics is outside PostgreSQL and is deliberately post-commit.
	// A caller that retries after this error will not duplicate relational state;
	// metric points are external at-least-once observations and need TSDB repair
	// if the write remains unavailable.
	return StoreRemoteCheckMetrics(report.Metrics)
}

// ApplyRemoteCheckResultFromWorkerTx applies all relational result effects using
// the caller's transaction. It intentionally does not write VictoriaMetrics or
// aggregate monitor state: both must happen only after the transaction commits.
// A nil monitor means the report was a consumed diagnostic attempt.
func ApplyRemoteCheckResultFromWorkerTx(tx *gorm.DB, report wire.CheckResultReport, regionCode string, worker *WorkerNode) (*Monitor, error) { //nolint:gocritic,lll // hugeParam: wire compatibility
	if tx == nil {
		return nil, fmt.Errorf("apply check result: nil transaction")
	}
	now := time.Now()
	if worker != nil {
		handled := false
		if err := ApplyDiagnosticResultTx(tx, report, worker, now, &handled); err != nil {
			return nil, err
		}
		if handled {
			return nil, nil
		}
	}
	check := Check{}
	if err := tx.Preload("Monitor").First(&check, report.CheckID).Error; err != nil {
		log.Println("worker: check not found:", report.CheckID, err)
		return nil, err
	}

	// Always persist the per-region result first so the aggregator can
	// pick it up regardless of which path we take next. We rely on
	// StoreCheckRegionResult to default AggregatedAt=NULL (the column
	// type is *time.Time, so a zero value writes SQL NULL).
	if err := StoreCheckRegionResultTx(tx, report, regionCode, now); err != nil {
		log.Println("worker: error storing region result:", report.CheckID, err)
		return nil, err
	}

	// Quorum-enabled checks: write nothing to Check.State here. The
	// aggregator will compute the aggregate state once the window has
	// elapsed (or enough regions have reported) and stamp AggregatedAt on
	// the contributing CheckRegionResult rows.
	if check.QuorumEnabled() {
		return nil, nil
	}

	update := map[string]interface{}{
		colState:    report.State,
		colLastEnd:  now,
		colWarnings: pq.StringArray(report.Warnings),
		colInfos:    pq.StringArray(report.Infos),
	}

	if report.State == "OK" {
		update["was_up"] = now
		update["last_ok"] = now
		update["fails"] = 0
		update["error"] = gorm.Expr("NULL")
	} else {
		update["last_fail"] = now
		update["fails"] = gorm.Expr("fails + 1")
		if report.Error != nil {
			update["error"] = *report.Error
		}
	}

	if report.ExpiresAt != nil {
		t, err := time.Parse(time.RFC3339, *report.ExpiresAt)
		if err == nil {
			update["expires"] = t
		}
	}

	if err := tx.Model(&check).UpdateColumns(update).Error; err != nil {
		log.Println("worker: error updating check:", report.CheckID, err)
		return nil, err
	}
	if worker != nil {
		payload, _ := json.Marshal(report)
		attempt := CheckAttempt{JobID: report.JobID, CheckID: check.ID, MonitorID: check.MonitorID, WorkerNodeID: &worker.ID, Kind: AttemptKindRegular, State: AttemptStateFinished, ResultState: report.State, Result: payload, StartedAt: &now, FinishedAt: &now, Deweighted: worker.NetworkProblemActive(now)}
		if attempt.JobID == "" {
			attempt.JobID = uuid.New().String()
		}
		// A duplicate websocket/HTTP delivery must not create another attempt.
		if err := tx.Where("job_id = ?", attempt.JobID).FirstOrCreate(&attempt).Error; err != nil {
			return nil, err
		}
		switch report.State {
		case stateERR, stateFail:
			if err := StartConfirmationTx(tx, check.ID, worker.ID, now); err != nil {
				return nil, err
			}
		case stateOK:
			if err := RecoverDiagnosticTx(tx, check.ID, now); err != nil {
				return nil, err
			}
		}
	}
	return check.Monitor, nil
}

// StoreRemoteCheckMetrics persists TSDB points reported by a distributed worker.
func StoreRemoteCheckMetrics(metrics []wire.MetricPoint) error {
	for _, metric := range metrics {
		if metric.Metric == "" || len(metric.Fields) == 0 {
			continue
		}
		if err := influx.WriteOne(metric.Metric, metric.Tags, metric.Fields); err != nil {
			return err
		}
	}
	return nil
}

// StoreCheckRegionResult stores a per-region check result for distributed monitoring analytics
func StoreCheckRegionResult(report wire.CheckResultReport, regionCode string) error { //nolint:gocritic,lll // hugeParam: accepted for interface compatibility
	return StoreCheckRegionResultTx(DB(), report, regionCode, time.Now())
}

// StoreCheckRegionResultTx stores a regional result in the caller's transaction.
func StoreCheckRegionResultTx(tx *gorm.DB, report wire.CheckResultReport, regionCode string, executedAt time.Time) error { //nolint:gocritic,lll // hugeParam: accepted for interface compatibility
	if tx == nil {
		return fmt.Errorf("store region result: nil transaction")
	}
	if regionCode == "" {
		regionCode = defaultRegionCode
	}

	result := CheckRegionResult{
		CheckID:    report.CheckID,
		RegionCode: regionCode,
		ExecutedAt: executedAt,
		State:      report.State,
		DurationMs: report.DurationMs,
		Error:      report.Error,
	}

	return tx.Create(&result).Error
}

// StaleWorkers marks workers as inactive or dead based on last_seen time
func StaleWorkers() {
	// Mark workers with no heartbeat for 2 minutes as inactive
	DB().Model(&WorkerNode{}).
		Where("status = ? AND last_seen < ?", "active", time.Now().Add(-2*time.Minute)).
		Update("status", "inactive")

	// Mark workers with no heartbeat for 5 minutes as dead
	DB().Model(&WorkerNode{}).
		Where("status IN (?, ?) AND last_seen < ?", "active", "inactive", time.Now().Add(-5*time.Minute)).
		Update("status", "dead")
}