<commit_before>package sched

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/tsaf/conf"
	"github.com/StackExchange/tsaf/expr"
)

type Schedule struct {
	sync.Mutex

	Conf          *conf.Conf
	Freq          time.Duration
	Status        map[AlertKey]*State
	Notifications map[AlertKey]map[string]time.Time

	cache     *opentsdb.Cache
	runStates map[AlertKey]Status
	nc        chan interface{}
}

func (s *Schedule) MarshalJSON() ([]byte, error) {
	t := struct {
		Alerts map[string]*conf.Alert
		Freq   time.Duration
		Status map[string]*State
	}{
		s.Conf.Alerts,
		s.Freq,
		make(map[string]*State),
	}
	for k, v := range s.Status {
		if v.Last().Status < ST_WARN {
			continue
		}
		t.Status[k.String()] = v
	}
	return json.Marshal(&t)
}

var DefaultSched = &Schedule{
	Freq: time.Minute * 5,
}

// Loads a configuration into the default schedule
func Load(c *conf.Conf) {
	DefaultSched.Load(c)
}

// Runs the default schedule.
func Run() error {
	return DefaultSched.Run()
}

func (s *Schedule) Load(c *conf.Conf) {
	s.Conf = c
	s.RestoreState()
}

// Restores notification and alert state from the file on disk.
func (s *Schedule) RestoreState() {
	s.Lock()
	defer s.Unlock()
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	s.Notifications = nil
	s.Status = make(map[AlertKey]*State)
	f, err := os.Open(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	dec := gob.NewDecoder(f)
	notifications := make(map[AlertKey]map[string]time.Time)
	if err := dec.Decode(&notifications); err != nil {
		log.Println(err)
		return
	}
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				log.Println("sched: notification not present during restore:", name)
				continue
			}
			_, present = s.Conf.Alerts[ak.Name]
			if !present {
				log.Println("sched: alert not present during restore:", ak.Name)
				continue
			}
			s.AddNotification(ak, n, t)
		}
	}
	for {
		var ak AlertKey
		var st State
		if err := dec.Decode(&ak); err == io.EOF {
			break
		} else if err != nil {
			log.Println(err)
			return
		}
		if err := dec.Decode(&st); err != nil {
			log.Println(err)
			return
		}
		if _, present := s.Conf.Alerts[ak.Name]; !present {
			log.Println("sched: alert no longer present, ignoring:", ak)
			continue
		}
		s.Status[ak] = &st
	}
}

func (s *Schedule) Save() {
	s.Lock()
	defer s.Unlock()
	f, err := os.Create(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	enc := gob.NewEncoder(f)
	if err := enc.Encode(s.Notifications); err != nil {
		log.Println(err)
		return
	}
	for k, v := range s.Status {
		enc.Encode(k)
		enc.Encode(v)
	}
	if err := f.Close(); err != nil {
		log.Println(err)
		return
	}
	log.Println("sched: wrote state to", s.Conf.StateFile)
}

func (s *Schedule) Run() error {
	go func() {
		for {
			time.Sleep(time.Minute)
			s.Save()
		}
	}()
	s.nc = make(chan interface{}, 1)
	go s.Poll()
	for {
		wait := time.After(s.Freq)
		if s.Freq < time.Second {
			return fmt.Errorf("sched: frequency must be > 1 second")
		}
		if s.Conf == nil {
			return fmt.Errorf("sched: nil configuration")
		}
		start := time.Now()
		s.Check()
		fmt.Printf("run at %v took %v\n", start, time.Since(start))
		<-wait
	}
}

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	var timeout time.Duration
	for {
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
		timeout = s.CheckNotifications()
	}
}

// CheckNotifications processes past notification events. It returns the
func (s *Schedule) CheckNotifications() time.Duration {
	s.Lock()
	defer s.Unlock()
	timeout := time.Hour
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				if remaining < timeout {
					timeout = remaining
				}
				s.AddNotification(ak, n, t)
				continue
			}
			st, present := s.Status[ak]
			if !present {
				continue
			}
			a, present := s.Conf.Alerts[ak.Name]
			if !present {
				continue
			}
			s.Notify(st, a, n)
			if n.Timeout < timeout {
				timeout = n.Timeout
			}
		}
	}
	return timeout
}

func (s *Schedule) Check() {
	s.Lock()
	defer s.Unlock()
	s.runStates = make(map[AlertKey]Status)
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	for _, a := range s.Conf.Alerts {
		s.CheckAlert(a)
	}
	changed := false
	for ak, status := range s.runStates {
		state := s.Status[ak]
		change := state.Append(status)
		if change {
			changed = true
		}
		if change && status > ST_NORM {
			var subject = new(bytes.Buffer)
			a := s.Conf.Alerts[ak.Name]
			if err := s.ExecuteSubject(subject, a, state); err != nil {
				log.Println(err)
			}
			state.Subject = subject.String()
			notify := func(notifications map[string]*conf.Notification) {
				for _, n := range notifications {
					s.Notify(state, a, n)
				}
			}
			switch status {
			case ST_CRIT:
				notify(a.CritNotification)
			case ST_WARN:
				notify(a.WarnNotification)
			}
		}
	}
	if changed {
		s.nc <- true
	}
}

func (s *Schedule) CheckAlert(a *conf.Alert) {
	crits := s.CheckExpr(a, a.Crit, true, nil)
	s.CheckExpr(a, a.Warn, false, crits)<commit_msg>Always set subject if displayed<commit_after>package sched

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/tsaf/conf"
	"github.com/StackExchange/tsaf/expr"
)

type Schedule struct {
	sync.Mutex

	Conf          *conf.Conf
	Freq          time.Duration
	Status        map[AlertKey]*State
	Notifications map[AlertKey]map[string]time.Time

	cache     *opentsdb.Cache
	runStates map[AlertKey]Status
	nc        chan interface{}
}

func (s *Schedule) MarshalJSON() ([]byte, error) {
	t := struct {
		Alerts map[string]*conf.Alert
		Freq   time.Duration
		Status map[string]*State
	}{
		s.Conf.Alerts,
		s.Freq,
		make(map[string]*State),
	}
	for k, v := range s.Status {
		if v.Last().Status < ST_WARN {
			continue
		}
		t.Status[k.String()] = v
	}
	return json.Marshal(&t)
}

var DefaultSched = &Schedule{
	Freq: time.Minute * 5,
}

// Loads a configuration into the default schedule
func Load(c *conf.Conf) {
	DefaultSched.Load(c)
}

// Runs the default schedule.
func Run() error {
	return DefaultSched.Run()
}

func (s *Schedule) Load(c *conf.Conf) {
	s.Conf = c
	s.RestoreState()
}

// Restores notification and alert state from the file on disk.
func (s *Schedule) RestoreState() {
	s.Lock()
	defer s.Unlock()
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	s.Notifications = nil
	s.Status = make(map[AlertKey]*State)
	f, err := os.Open(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	dec := gob.NewDecoder(f)
	notifications := make(map[AlertKey]map[string]time.Time)
	if err := dec.Decode(&notifications); err != nil {
		log.Println(err)
		return
	}
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				log.Println("sched: notification not present during restore:", name)
				continue
			}
			_, present = s.Conf.Alerts[ak.Name]
			if !present {
				log.Println("sched: alert not present during restore:", ak.Name)
				continue
			}
			s.AddNotification(ak, n, t)
		}
	}
	for {
		var ak AlertKey
		var st State
		if err := dec.Decode(&ak); err == io.EOF {
			break
		} else if err != nil {
			log.Println(err)
			return
		}
		if err := dec.Decode(&st); err != nil {
			log.Println(err)
			return
		}
		if _, present := s.Conf.Alerts[ak.Name]; !present {
			log.Println("sched: alert no longer present, ignoring:", ak)
			continue
		}
		s.Status[ak] = &st
	}
}

func (s *Schedule) Save() {
	s.Lock()
	defer s.Unlock()
	f, err := os.Create(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	enc := gob.NewEncoder(f)
	if err := enc.Encode(s.Notifications); err != nil {
		log.Println(err)
		return
	}
	for k, v := range s.Status {
		enc.Encode(k)
		enc.Encode(v)
	}
	if err := f.Close(); err != nil {
		log.Println(err)
		return
	}
	log.Println("sched: wrote state to", s.Conf.StateFile)
}

func (s *Schedule) Run() error {
	go func() {
		for {
			time.Sleep(time.Minute)
			s.Save()
		}
	}()
	s.nc = make(chan interface{}, 1)
	go s.Poll()
	for {
		wait := time.After(s.Freq)
		if s.Freq < time.Second {
			return fmt.Errorf("sched: frequency must be > 1 second")
		}
		if s.Conf == nil {
			return fmt.Errorf("sched: nil configuration")
		}
		start := time.Now()
		s.Check()
		fmt.Printf("run at %v took %v\n", start, time.Since(start))
		<-wait
	}
}

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	var timeout time.Duration
	for {
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
		timeout = s.CheckNotifications()
	}
}

// CheckNotifications processes past notification events. It returns the
func (s *Schedule) CheckNotifications() time.Duration {
	s.Lock()
	defer s.Unlock()
	timeout := time.Hour
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				if remaining < timeout {
					timeout = remaining
				}
				s.AddNotification(ak, n, t)
				continue
			}
			st, present := s.Status[ak]
			if !present {
				continue
			}
			a, present := s.Conf.Alerts[ak.Name]
			if !present {
				continue
			}
			s.Notify(st, a, n)
			if n.Timeout < timeout {
				timeout = n.Timeout
			}
		}
	}
	return timeout
}

func (s *Schedule) Check() {
	s.Lock()
	defer s.Unlock()
	s.runStates = make(map[AlertKey]Status)
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	for _, a := range s.Conf.Alerts {
		s.CheckAlert(a)
	}
	changed := false
	for ak, status := range s.runStates {
		state := s.Status[ak]
		change := state.Append(status)
		if change {
			changed = true
		}
		if status > ST_NORM {
			var subject = new(bytes.Buffer)
			a := s.Conf.Alerts[ak.Name]
			if err := s.ExecuteSubject(subject, a, state); err != nil {
				log.Println(err)
			}
			state.Subject = subject.String()
			if change {
				notify := func(notifications map[string]*conf.Notification) {
					for _, n := range notifications {
						s.Notify(state, a, n)
					}
				}
				switch status {
				case ST_CRIT:
					notify(a.CritNotification)
				case ST_WARN:
					notify(a.WarnNotification)
				}
			}
		}
	}
	if changed {
		s.nc <- true
	}
}

func (s *Schedule) CheckAlert(a *conf.Alert) {
	crits := s.CheckExpr(a, a.Crit, true, nil)
	s.CheckExpr(a, a.Warn, false, crits)
{'fraction_non_alphanumeric': 0.1480393057213702, 'fraction_numerical': 0.0007346863807512168, 'mean_word_length': 2.986090775988287, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4762393', 'n_tokens_mistral': 4551, 'n_tokens_neox': 4065, 'n_words': 994}