440 lines
8.9 KiB
Go
440 lines
8.9 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"context"
|
|
"reflect"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/monlet/agent/internal/config"
|
|
"github.com/monlet/agent/internal/eventid"
|
|
"github.com/monlet/agent/internal/runner"
|
|
)
|
|
|
|
var parseCronSchedule = config.ParseCronSchedule
|
|
|
|
// Event is a contract-shaped check result emitted by the scheduler.
|
|
type Event struct {
|
|
EventID string `json:"event_id"`
|
|
CheckID string `json:"check_id"`
|
|
ObservedAt time.Time `json:"observed_at"`
|
|
Status string `json:"status"`
|
|
ExitCode int `json:"exit_code"`
|
|
DurationMs int64 `json:"duration_ms"`
|
|
Output string `json:"output,omitempty"`
|
|
OutputTruncated bool `json:"output_truncated,omitempty"`
|
|
NotificationsEnabled *bool `json:"notifications_enabled,omitempty"`
|
|
IncidentKey string `json:"incident_key,omitempty"`
|
|
}
|
|
|
|
func (e Event) WireEvent() Event {
|
|
if e.NotificationsEnabled == nil {
|
|
enabled := true
|
|
e.NotificationsEnabled = &enabled
|
|
}
|
|
return e
|
|
}
|
|
|
|
type Hooks struct {
|
|
OnSkipped func(checkID string)
|
|
OnResult func(ev Event)
|
|
OnEventChannelFull func()
|
|
OnResourceLimitFailure func(checkID string, resource string)
|
|
OnStopped func(checkID string)
|
|
}
|
|
|
|
// Manager owns per-check workers and applies hot config updates.
|
|
type Manager struct {
|
|
ctx context.Context
|
|
agentID string
|
|
out chan<- Event
|
|
hooks Hooks
|
|
|
|
mu sync.Mutex
|
|
workers map[string]*checkWorker
|
|
draining map[string]struct{}
|
|
pending map[string]config.CheckConfig
|
|
all []*checkWorker
|
|
stopped bool
|
|
}
|
|
|
|
// NewManager creates a scheduler manager bound to the parent context.
|
|
func NewManager(ctx context.Context, agentID string, out chan<- Event, h Hooks) *Manager {
|
|
return &Manager{
|
|
ctx: ctx,
|
|
agentID: agentID,
|
|
out: out,
|
|
hooks: h,
|
|
workers: make(map[string]*checkWorker),
|
|
draining: make(map[string]struct{}),
|
|
pending: make(map[string]config.CheckConfig),
|
|
}
|
|
}
|
|
|
|
// Update applies the current check set without killing in-flight runs.
|
|
func (m *Manager) Update(checks []config.CheckConfig) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.stopped {
|
|
return
|
|
}
|
|
next := make(map[string]config.CheckConfig, len(checks))
|
|
for _, c := range checks {
|
|
next[c.ID] = c
|
|
if w, ok := m.workers[c.ID]; ok {
|
|
if _, draining := m.draining[c.ID]; draining {
|
|
m.pending[c.ID] = c
|
|
} else {
|
|
w.Update(c)
|
|
}
|
|
continue
|
|
}
|
|
id := c.ID
|
|
w := newCheckWorker(m.ctx, m.agentID, c, m.out, m.hooks, func(stopped *checkWorker) {
|
|
m.workerStopped(id, stopped)
|
|
})
|
|
m.workers[c.ID] = w
|
|
m.all = append(m.all, w)
|
|
}
|
|
for id, w := range m.workers {
|
|
if _, ok := next[id]; ok {
|
|
continue
|
|
}
|
|
m.draining[id] = struct{}{}
|
|
delete(m.pending, id)
|
|
w.Stop()
|
|
}
|
|
}
|
|
|
|
// Stop prevents new runs and lets in-flight checks finish or observe ctx cancel.
|
|
func (m *Manager) Stop() {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if m.stopped {
|
|
return
|
|
}
|
|
m.stopped = true
|
|
for _, w := range m.workers {
|
|
w.Stop()
|
|
}
|
|
}
|
|
|
|
// Wait blocks until all workers have drained.
|
|
func (m *Manager) Wait() {
|
|
m.mu.Lock()
|
|
workers := append([]*checkWorker(nil), m.all...)
|
|
m.mu.Unlock()
|
|
for _, w := range workers {
|
|
w.Wait()
|
|
}
|
|
}
|
|
|
|
func (m *Manager) workerStopped(checkID string, w *checkWorker) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if current, ok := m.workers[checkID]; !ok || current != w {
|
|
return
|
|
}
|
|
delete(m.workers, checkID)
|
|
delete(m.draining, checkID)
|
|
if m.stopped {
|
|
delete(m.pending, checkID)
|
|
return
|
|
}
|
|
if cfg, ok := m.pending[checkID]; ok {
|
|
delete(m.pending, checkID)
|
|
next := newCheckWorker(m.ctx, m.agentID, cfg, m.out, m.hooks, func(stopped *checkWorker) {
|
|
m.workerStopped(checkID, stopped)
|
|
})
|
|
m.workers[checkID] = next
|
|
m.all = append(m.all, next)
|
|
}
|
|
}
|
|
|
|
type checkWorker struct {
|
|
ctx context.Context
|
|
agentID string
|
|
out chan<- Event
|
|
hooks Hooks
|
|
onDone func(*checkWorker)
|
|
|
|
updateCh chan config.CheckConfig
|
|
stopCh chan struct{}
|
|
doneCh chan struct{}
|
|
stopOnce sync.Once
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
func newCheckWorker(ctx context.Context, agentID string, c config.CheckConfig, out chan<- Event, h Hooks, onDone func(*checkWorker)) *checkWorker {
|
|
w := &checkWorker{
|
|
ctx: ctx,
|
|
agentID: agentID,
|
|
out: out,
|
|
hooks: h,
|
|
onDone: onDone,
|
|
updateCh: make(chan config.CheckConfig, 1),
|
|
stopCh: make(chan struct{}),
|
|
doneCh: make(chan struct{}, 1),
|
|
}
|
|
w.wg.Add(1)
|
|
go w.loop(c)
|
|
return w
|
|
}
|
|
|
|
func (w *checkWorker) Update(c config.CheckConfig) {
|
|
select {
|
|
case w.updateCh <- c:
|
|
default:
|
|
select {
|
|
case <-w.updateCh:
|
|
default:
|
|
}
|
|
w.updateCh <- c
|
|
}
|
|
}
|
|
|
|
func (w *checkWorker) Stop() {
|
|
w.stopOnce.Do(func() { close(w.stopCh) })
|
|
}
|
|
|
|
func (w *checkWorker) Wait() {
|
|
w.wg.Wait()
|
|
}
|
|
|
|
func (w *checkWorker) loop(initial config.CheckConfig) {
|
|
cfg := initial
|
|
stoppedNotified := false
|
|
notifyStopped := func() {
|
|
if stoppedNotified {
|
|
return
|
|
}
|
|
stoppedNotified = true
|
|
if w.hooks.OnStopped != nil {
|
|
w.hooks.OnStopped(cfg.ID)
|
|
}
|
|
}
|
|
defer func() {
|
|
notifyStopped()
|
|
if w.onDone != nil {
|
|
w.onDone(w)
|
|
}
|
|
w.wg.Done()
|
|
}()
|
|
t := time.NewTimer(time.Hour)
|
|
stopTimer(t)
|
|
defer t.Stop()
|
|
var timerCh <-chan time.Time
|
|
|
|
running := false
|
|
stopping := false
|
|
runAfterCurrent := false
|
|
ctxDone := w.ctx.Done()
|
|
stopCh := w.stopCh
|
|
|
|
start := func(c config.CheckConfig) {
|
|
if running {
|
|
if w.hooks.OnSkipped != nil {
|
|
w.hooks.OnSkipped(c.ID)
|
|
}
|
|
return
|
|
}
|
|
running = true
|
|
go func() {
|
|
w.run(c)
|
|
w.doneCh <- struct{}{}
|
|
}()
|
|
}
|
|
|
|
arm := func(c config.CheckConfig, immediateInterval bool) bool {
|
|
delay, err := nextDelay(c, time.Now(), immediateInterval)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
resetTimer(t, delay)
|
|
timerCh = t.C
|
|
return true
|
|
}
|
|
disarm := func() {
|
|
stopTimer(t)
|
|
timerCh = nil
|
|
}
|
|
|
|
if !arm(cfg, true) {
|
|
return
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case <-ctxDone:
|
|
stopping = true
|
|
ctxDone = nil
|
|
if !running {
|
|
return
|
|
}
|
|
case <-stopCh:
|
|
stopping = true
|
|
stopCh = nil
|
|
if !running {
|
|
return
|
|
}
|
|
case next := <-w.updateCh:
|
|
if reflect.DeepEqual(cfg, next) {
|
|
continue
|
|
}
|
|
cfg = next
|
|
if cfg.UsesCron() {
|
|
runAfterCurrent = false
|
|
if !arm(cfg, false) {
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
if running {
|
|
runAfterCurrent = true
|
|
disarm()
|
|
} else {
|
|
start(cfg)
|
|
if !arm(cfg, false) {
|
|
return
|
|
}
|
|
}
|
|
case <-w.doneCh:
|
|
running = false
|
|
if !stopping {
|
|
select {
|
|
case <-stopCh:
|
|
stopping = true
|
|
stopCh = nil
|
|
default:
|
|
}
|
|
}
|
|
if stopping {
|
|
return
|
|
}
|
|
if runAfterCurrent {
|
|
runAfterCurrent = false
|
|
startAfterDrain := true
|
|
for {
|
|
select {
|
|
case next := <-w.updateCh:
|
|
if !reflect.DeepEqual(cfg, next) {
|
|
cfg = next
|
|
startAfterDrain = !cfg.UsesCron()
|
|
}
|
|
default:
|
|
if startAfterDrain {
|
|
start(cfg)
|
|
}
|
|
if !arm(cfg, false) {
|
|
return
|
|
}
|
|
goto nextLoop
|
|
}
|
|
}
|
|
nextLoop:
|
|
continue
|
|
}
|
|
case <-timerCh:
|
|
if stopping {
|
|
continue
|
|
}
|
|
start(cfg)
|
|
if !arm(cfg, false) {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *checkWorker) run(c config.CheckConfig) {
|
|
res := runner.RunCommand(w.ctx, c.Command, c.Timeout.Duration, runner.ResourceLimits{
|
|
CPUTime: c.ResourceLimits.CPUTime.Duration,
|
|
MemoryBytes: c.ResourceLimits.Memory.Bytes,
|
|
OpenFiles: c.ResourceLimits.OpenFiles,
|
|
})
|
|
if res.ResourceLimitFailure != "" && w.hooks.OnResourceLimitFailure != nil {
|
|
w.hooks.OnResourceLimitFailure(c.ID, res.ResourceLimitFailure)
|
|
}
|
|
id, err := eventid.New()
|
|
if err != nil {
|
|
return
|
|
}
|
|
ev := Event{
|
|
EventID: id,
|
|
CheckID: c.ID,
|
|
ObservedAt: time.Now().UTC(),
|
|
Status: string(res.Status),
|
|
ExitCode: res.ExitCode,
|
|
DurationMs: res.DurationMs,
|
|
Output: res.Output,
|
|
OutputTruncated: res.OutputTruncated,
|
|
NotificationsEnabled: boolPtr(c.NotificationsOn()),
|
|
IncidentKey: incidentKey(w.agentID, c),
|
|
}
|
|
if w.hooks.OnResult != nil {
|
|
w.hooks.OnResult(ev)
|
|
}
|
|
select {
|
|
case w.out <- ev:
|
|
case <-w.ctx.Done():
|
|
default:
|
|
if w.hooks.OnEventChannelFull != nil {
|
|
w.hooks.OnEventChannelFull()
|
|
}
|
|
select {
|
|
case w.out <- ev:
|
|
case <-w.ctx.Done():
|
|
case <-w.stopCh:
|
|
}
|
|
}
|
|
}
|
|
|
|
func boolPtr(v bool) *bool {
|
|
return &v
|
|
}
|
|
|
|
func incidentKey(agentID string, c config.CheckConfig) string {
|
|
if c.DedupeKey != "" {
|
|
return c.DedupeKey
|
|
}
|
|
return agentID + ":" + c.ID
|
|
}
|
|
|
|
func nextDelay(c config.CheckConfig, now time.Time, immediateInterval bool) (time.Duration, error) {
|
|
if c.UsesCron() {
|
|
schedule, err := parseCronSchedule(c.Cron)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return delayUntil(now, schedule.Next(now)), nil
|
|
}
|
|
if immediateInterval {
|
|
return 0, nil
|
|
}
|
|
return c.Interval.Duration, nil
|
|
}
|
|
|
|
func delayUntil(now, next time.Time) time.Duration {
|
|
if next.Before(now) {
|
|
return 0
|
|
}
|
|
return next.Sub(now)
|
|
}
|
|
|
|
func resetTimer(t *time.Timer, d time.Duration) {
|
|
if d < 0 {
|
|
d = 0
|
|
}
|
|
stopTimer(t)
|
|
t.Reset(d)
|
|
}
|
|
|
|
func stopTimer(t *time.Timer) {
|
|
if !t.Stop() {
|
|
select {
|
|
case <-t.C:
|
|
default:
|
|
}
|
|
}
|
|
}
|