main.go

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/NectGmbH/health"
	"github.com/fsnotify/fsnotify"
	"github.com/sirupsen/logrus"
)

type HealthCheck struct {
	*health.HealthCheck
	Group        string
	LinesInInput []int
}

type LoggingHealthCheckProvider struct {
	wrapped health.HealthCheckProvider
}

func NewLoggingHealthCheckProvider(base health.HealthCheckProvider) health.HealthCheckProvider {
	return &LoggingHealthCheckProvider{wrapped: base}
}

func (p *LoggingHealthCheckProvider) CheckHealth(healthCheck *health.HealthCheck) (string, bool) {
	logger := logrus.WithFields(logrus.Fields{
		"endpoint": healthCheck.GetAddress(),
	})

	logger.Debug("Begin check health")
	s, b := p.wrapped.CheckHealth(healthCheck)
	logger.WithFields(logrus.Fields{"msg": s, "healthy": b}).Debug("Finished check health")

	if !b {
		logger.WithField("reason", s).Warn("endpoint down")
	}

	return s, b
}

func main() {
	var inputFile string
	var outputFile string
	var header string
	var magic string
	var command string
	var replaceWithTimestamp string
	var debug bool
	var json bool
	flag.StringVar(&inputFile, "input", "", "input file to healthcheck svcs from")
	flag.StringVar(&outputFile, "output", "", "output file where the modified file should be written to")
	flag.StringVar(&header, "header", "", "header to add to the file, e.g. '# AUTOGENERATED DO NOT TOUCH'")
	flag.StringVar(&command, "command", "", "the command to execute when output file has changed")
	flag.StringVar(&magic, "magic", "HealthCheck", "Magical token to look for in lines")
	flag.StringVar(&replaceWithTimestamp, "replace-with-timestamp", "", "Optional feature to replace the passed token with the current timestamp")
	flag.BoolVar(&debug, "debug", false, "enables debug options")
	flag.BoolVar(&json, "json", false, "enables json logging")
	flag.Parse()

	if json {
		logrus.SetFormatter(&logrus.JSONFormatter{})
	}

	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	if inputFile == "" {
		logrus.Fatal("Missing -input")
	}

	if outputFile == "" {
		logrus.Fatal("Missing -output")
	}

	if magic == "" {
		logrus.Fatal("Missing -magic")
	}

	inputInfo, err := os.Stat(inputFile)
	if err != nil {
		logrus.WithField("input", inputFile).Fatal("couldn't stat input file")
	}

	input, err := readLines(inputFile)
	if err != nil {
		logrus.WithField("input", inputFile).Fatal("couldn't read input file")
	}

	notificationChannel := make(chan health.HealthCheckStatus)
	healthChecks, stopChans, err := setupHealthChecks(input, magic, notificationChannel)
	if err != nil {
		logrus.WithFields(logrus.Fields{"input": inputFile, "reason": err}).Fatal("couldn't setup healthchecks")
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		logrus.WithField("reason", err).Fatal("coudln't set up fsnotify watcher")
	}
	defer watcher.Close()

	lock := &sync.Mutex{}

	go func() {
		for {
			select {
			case event, ok := <-watcher.Events:
				if !ok {
					logrus.Fatal("couldn't read from watcher.Events chan")
					return
				}

				if event.Op&fsnotify.Write == fsnotify.Write {
					logrus.WithField("input", inputFile).Info("input changed, reloading.")

					logrus.Info("stopping healthchecks")
					for _, s := range stopChans {
						s <- struct{}{}
					}

					lock.Lock()
					inputInfo, err = os.Stat(inputFile)
					if err != nil {
						logrus.WithField("input", inputFile).Fatal("couldn't stat input file")
					}

					input, err = readLines(inputFile)
					if err != nil {
						logrus.WithField("input", input).Fatal("couldn't read input file")
					}

					healthChecks, stopChans, err = setupHealthChecks(input, magic, notificationChannel)
					if err != nil {
						logrus.WithFields(logrus.Fields{"input": inputFile, "reason": err}).Fatal("couldn't setup healthchecks")
					}

					logrus.Info("restarted healthchecks")

					lock.Unlock()
				}
			case err, ok := <-watcher.Errors:
				if !ok {
					logrus.Fatal("couldn't read from watcher.Errors chan")
					return
				}
				logrus.WithField("reason", err).Error("received error from fsnotify")
			}
		}
	}()

	err = watcher.Add(inputFile)
	if err != nil {
		logrus.WithFields(logrus.Fields{"input": inputFile, "reason": err}).Fatal("couldn't add fsnotify watcher for input")
	}

	logrus.Info("Listening for notification...")

	for {
		notification := <-notificationChannel

		addr := notification.GetAddress()
		logger := logrus.WithFields(logrus.Fields{
			"endpoint": addr,
			"healthy":  notification.Healthy,
			"msg":      notification.Message,
		})

		logger.Debug("received nottification")

		if !notification.DidChange {
			logger.Debug("endpoint didnt change, ignoring event.")
			continue
		}

		lock.Lock()
		inputInfoCopy := inputInfo
		inputCopy := input
		healthChecksCopy := make(map[string]HealthCheck)
		for k, v := range healthChecks {
			healthChecksCopy[k] = v
		}
		lock.Unlock()

		if _, ok := healthChecksCopy[addr]; !ok {
			logger.Info("received event for no more monitored endpoint, ignoring.")
			continue
		}

		logger.Info("endpoint changed healthiness, updating output...")

		f, err := os.OpenFile(outputFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, inputInfoCopy.Mode().Perm())
		if err != nil {
			logger.WithFields(logrus.Fields{"output": outputFile, "reason": err}).Fatal("Couldn't open output file")
		}

		writer := bufio.NewWriter(f)
		if header != "" {
			_, err := writer.WriteString(header + "\n")
			if err != nil {
				logger.WithFields(logrus.Fields{"output": outputFile, "reason": err}).Fatal("Couldn't write output file")
			}
		}

		lineIndexMap := make(map[int]string)
		groupHealth := make(map[string]bool) // Used to group, so we can fallback to all if everything is reported as down

		for k, v := range healthChecksCopy {
			for _, line := range v.LinesInInput {
				lineIndexMap[line] = k

				if v.Group != "" {
					_, ok := groupHealth[v.Group]
					if !ok || v.Healthy {
						groupHealth[v.Group] = v.Healthy
					}
				}
			}
		}

		for i, s := range inputCopy {
			key, isHealthCheckLine := lineIndexMap[i]
			if !isHealthCheckLine || healthChecksCopy[key].Healthy || (healthChecksCopy[key].Group != "" && !groupHealth[healthChecksCopy[key].Group]) {
				line := s

				if replaceWithTimestamp != "" {
					line = strings.ReplaceAll(line, replaceWithTimestamp, fmt.Sprintf("%d", time.Now().Unix()))
				}

				_, err := writer.WriteString(line + "\n")
				if err != nil {
					logger.WithFields(logrus.Fields{"output": outputFile, "reason": err}).Fatal("Couldn't write output file")
				}
			}
		}

		writer.Flush()

		if err := f.Close(); err != nil {
			logger.WithFields(logrus.Fields{"output": outputFile, "reason": err}).Fatal("Couldn't close output file")
		}

		logger.Info("updated output file")

		if command != "" {
			logger = logger.WithField("command", command)
			logger.Info("Executing reload command...")
			cmd := exec.Command("/bin/sh", "-c", command)
			stdout, err := cmd.Output()
			if err != nil {
				logger.WithField("reason", err).Error("command failed")
				continue
			} else {
				logger.WithField("stdout", string(stdout)).Info("command succeeded")
			}
		}
	}

}

func setupHealthChecks(input []string, magic string, channel chan health.HealthCheckStatus) (map[string]HealthCheck, []chan struct{}, error) {
	logrus.Info("setting up health checks")

	healthChecks := make(map[string]HealthCheck)
	stopChans := make([]chan struct{}, 0)

	activeChans := int32(0)

	for i, line := range input {
		logger := logrus.WithFields(logrus.Fields{"line": i})

		idx := strings.Index(line, magic+"{")

		if idx < 0 {
			continue
		}

		logger.WithField("lineContent", line).Debug("found healthcheck line")

		reader := bytes.NewReader([]byte(line[idx+len(magic):]))
		dec := json.NewDecoder(reader)

		data := struct {
			Group              string `json:"group"`
			Host               string `json:"host"`
			Port               int    `json:"port"`
			Proto              string `json:"proto"`
			InsecureSkipVerify bool   `json:"insecureSkipVerify"`
			Timeout            string `json:"timeout"`
			Interval           string `json:"interval"`
			Path               string `json:"path"`
		}{}

		if err := dec.Decode(&data); err != nil {
			return nil, nil, fmt.Errorf("couldn't decode health check config, see: %v", err)
		}

		if data.Proto == "" {
			data.Proto = "tcp"
		}

		if data.Timeout == "" {
			data.Timeout = "5s"
		}

		if data.Interval == "" {
			data.Interval = "3s"
		}

		timeout, err := time.ParseDuration(data.Timeout)
		if err != nil {
			return nil, nil, fmt.Errorf("couldn't parse timeout duration `%s`, see: %v", data.Timeout, err)
		}

		interval, err := time.ParseDuration(data.Interval)
		if err != nil {
			return nil, nil, fmt.Errorf("couldn't parse interval duration `%s`, see: %v", data.Interval, err)
		}

		provider, err := health.GetHealthCheckProvider(data.Proto)
		if err != nil {
			return nil, nil, fmt.Errorf("couldn't get healthcheck provider, see: %v", err)
		}

		if data.Proto == "https" || data.Proto == "http" {
			provider = health.NewHTTPHealthCheckProvider(data.Proto == "https", data.InsecureSkipVerify, data.Path)
		}

		provider = NewLoggingHealthCheckProvider(provider)

		// FIXME: currently host+port is the unique key. that might suck.
		healthCheck := health.NewHealthCheck(
			net.ParseIP(data.Host),
			data.Port,
			provider,
			interval,
			interval*10,
			timeout,
		)

		logger.WithFields(logrus.Fields{
			"endpoint":           healthCheck.GetAddress(),
			"proto":              data.Proto,
			"insecureSkipVerify": data.InsecureSkipVerify,
			"timeout":            data.Timeout,
			"interval":           data.Interval,
		}).Info("set up health check")

		stopChan := make(chan struct{})
		stopChans = append(stopChans, stopChan)

		notificationCh := healthCheck.Monitor(stopChan)
		go (func() {
			for n := range notificationCh {
				channel <- n
			}

			if atomic.AddInt32(&activeChans, -1) == 0 {
				close(channel)
			}
		})()

		healthChecks[healthCheck.GetAddress()] = HealthCheck{HealthCheck: healthCheck, Group: data.Group, LinesInInput: []int{i}}
	}

	return healthChecks, stopChans, nil
}

func readLines(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	return lines, scanner.Err()
}