nmon/internal/checks/runner.go

package checks

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/v10networks/nmon/internal/clock"
	"github.com/v10networks/nmon/internal/metrics"
	"github.com/v10networks/nmon/internal/netpolicy"
	"github.com/v10networks/nmon/internal/selfmon"
	"github.com/v10networks/nmon/internal/store"
)

// Source is the part of store.Store the runner reads.
type Source interface {
	ListChecks(ctx context.Context, includeDisabled bool) ([]store.Check, error)
	GetDeviceByID(ctx context.Context, id int64) (*store.Device, error)
}

// Config tunes a Runner. Zero values take the defaults shown.
type Config struct {
	// Workers bounds checks running at once (default 32). A check that finds
	// the pool full waits for the next tick.
	Workers int `yaml:"workers" json:"workers"`
	// Timeout is the default deadline of one run (default 10 s, never above the check's interval).
	Timeout time.Duration `yaml:"timeout" json:"timeout"`
	// Tick is how often due checks are looked for (default 5 s).
	Tick time.Duration `yaml:"tick" json:"tick"`
}

// Runner schedules checks and writes their results.
type Runner struct {
	cfg    Config
	clk    clock.Clock
	src    Source
	pol    *netpolicy.Policy
	pinger Pinger
	out    metrics.Writer
	obs    Observer
	log    *slog.Logger
	reg    *selfmon.Registry
	http   *http.Client

	sem chan struct{}
	wg  sync.WaitGroup

	mu       sync.Mutex
	next     map[int64]time.Time
	inflight map[int64]bool
	latest   map[int64]map[int64]Probe // device id -> check id -> latest probe
	last     map[int64]LastResult
}

// LastResult is the latest outcome of one check, for the API.
type LastResult struct {
	CheckID int64
	Device  string
	At      time.Time
	Result  Result
}

// New builds a Runner. pol is the only way out to the network; pinger sends
// ICMP echo (an *icmp.Pinger); out receives samples; obs gets ObserveCheck.
func New(cfg Config, pol *netpolicy.Policy, clk clock.Clock, src Source, pinger Pinger,
	out metrics.Writer, obs Observer, log *slog.Logger) (*Runner, error) {
	if cfg.Workers <= 0 {
		cfg.Workers = 32
	}
	if cfg.Timeout <= 0 {
		cfg.Timeout = 10 * time.Second
	}
	if cfg.Tick <= 0 {
		cfg.Tick = 5 * time.Second
	}
	if log == nil {
		log = slog.Default()
	}
	hc := pol.HTTPClient(2 * cfg.Timeout) // the run's own deadline fires first; this is a backstop
	return &Runner{cfg: cfg, clk: clk, src: src, pol: pol, pinger: pinger, out: out, obs: obs, log: log, reg: selfmon.Default, http: hc,
		sem: make(chan struct{}, cfg.Workers), next: map[int64]time.Time{}, inflight: map[int64]bool{},
		latest: map[int64]map[int64]Probe{}, last: map[int64]LastResult{}}, nil
}

// Run looks for due checks every Tick until ctx ends, then waits for running checks.
func (r *Runner) Run(ctx context.Context) {
	defer r.Wait()
	t := r.clk.NewTicker(r.cfg.Tick)
	defer t.Stop()
	for {
		r.RunDue(ctx)
		select {
		case <-ctx.Done():
			return
		case <-t.C():
		}
	}
}

// Wait blocks until every started check has finished.
func (r *Runner) Wait() { r.wg.Wait() }

// Results returns the latest result of each check that has run.
func (r *Runner) Results() []LastResult {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]LastResult, 0, len(r.last))
	for _, l := range r.last {
		out = append(out, l)
	}
	return out
}

// RunDue starts every enabled check whose time has come and returns without
// waiting for them, so one slow target never holds up the rest. A check still
// running is not started again.
func (r *Runner) RunDue(ctx context.Context) {
	cs, err := r.src.ListChecks(ctx, false)
	if err != nil {
		r.log.Warn("checks: list failed", "err", err)
		return
	}
	now := r.clk.Now()
	r.forget(cs)
	for _, c := range cs {
		if !r.claim(c, now) {
			continue
		}
		slug, ok := r.deviceSlug(ctx, c)
		if !ok {
			r.release(c.ID)
			continue
		}
		r.wg.Add(1)
		go func() {
			defer func() { <-r.sem; r.wg.Done() }()
			r.run(ctx, c, slug, now)
			r.release(c.ID)
		}()
	}
}

// claim marks c in flight and schedules its next run if it is due, a worker is
// free and it is not already running.
func (r *Runner) claim(c store.Check, now time.Time) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.inflight[c.ID] || now.Before(r.next[c.ID]) {
		return false
	}
	select {
	case r.sem <- struct{}{}:
	default:
		return false
	}
	r.inflight[c.ID] = true
	r.next[c.ID] = now.Add(interval(c))
	return true
}

func (r *Runner) release(id int64) {
	r.mu.Lock()
	delete(r.inflight, id)
	r.mu.Unlock()
}

func interval(c store.Check) time.Duration {
	if c.IntervalS <= 0 {
		return DefaultIntervalS * time.Second
	}
	return time.Duration(c.IntervalS) * time.Second
}

// forget drops state of checks that no longer exist or are disabled.
func (r *Runner) forget(cs []store.Check) {
	live := make(map[int64]bool, len(cs))
	for _, c := range cs {
		live[c.ID] = true
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	for id := range r.next {
		if !live[id] {
			delete(r.next, id)
			delete(r.last, id)
		}
	}
	for _, m := range r.latest {
		for id := range m {
			if !live[id] {
				delete(m, id)
			}
		}
	}
}

// deviceSlug returns the slug of the device the check reports for; a check
// whose device is gone or retired does not run.
func (r *Runner) deviceSlug(ctx context.Context, c store.Check) (string, bool) {
	if c.DeviceID == nil {
		return "", false
	}
	d, err := r.src.GetDeviceByID(ctx, *c.DeviceID)
	if err != nil || d.RetiredAt != nil {
		return "", false
	}
	return d.Slug, true
}

// run executes one check under its own deadline, then records the result.
func (r *Runner) run(ctx context.Context, c store.Check, slug string, at time.Time) {
	p, err := ParseParams(c.Params)
	var res Result
	if err != nil {
		res = fail(err)
	} else {
		rctx, cancel := withClockTimeout(ctx, r.clk, p.timeout(r.cfg.Timeout, interval(c)))
		res = r.probe(rctx, c, p)
		cancel()
	}
	r.record(ctx, c, slug, at, res)
}

func (r *Runner) probe(ctx context.Context, c store.Check, p Params) Result {
	var res Result
	switch c.Kind {
	case store.CheckPing:
		res = r.ping(ctx, c.Target, p.timeout(r.cfg.Timeout, interval(c)))
	case store.CheckHTTP:
		res = r.httpCheck(ctx, c.Target, p)
	case store.CheckDNS:
		res = r.dnsCheck(ctx, c.Target, p)
	case store.CheckTLS:
		res = r.tlsCheck(ctx, c.Target, p)
	default:
		res = fail(fmt.Errorf("unknown check kind %q", c.Kind))
	}
	if ctx.Err() != nil && res.Probe == ProbeFail {
		res.Err = "timeout: " + res.Err
	}
	return res
}

// record writes the samples, remembers the result and reports the device's
// combined outcome to reachability.
func (r *Runner) record(ctx context.Context, c store.Check, slug string, at time.Time, res Result) {
	name := [...]string{"ok", "fail", "unsupported"}[res.Probe]
	r.reg.Counter("nmon_check_runs_total", "kind", c.Kind, "result", name).Inc()
	if s := res.samples(slug, fmt.Sprint(c.ID), at); len(s) > 0 {
		if err := r.out.Write(ctx, s); err != nil {
			r.reg.Counter("nmon_check_write_errors_total").Inc()
			r.log.Warn("checks: write failed", "check", c.ID, "err", err)
		}
	}
	r.mu.Lock()
	r.last[c.ID] = LastResult{CheckID: c.ID, Device: slug, At: at, Result: res}
	m := r.latest[*c.DeviceID]
	if m == nil {
		m = map[int64]Probe{}
		r.latest[*c.DeviceID] = m
	}
	m[c.ID] = res.Probe
	combined := combine(m)
	r.mu.Unlock()
	r.obs.ObserveCheck(slug, combined, at)
}

// withClockTimeout is context.WithTimeout driven by clk, so deadlines follow a
// fake clock in tests.
func withClockTimeout(parent context.Context, clk clock.Clock, d time.Duration) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(parent)
	go func() {
		select {
		case <-clk.After(d):
			cancel()
		case <-ctx.Done():
		}
	}()
	return ctx, cancel
}