dn42promsrv/dn42promsrv.go

//////////////////////////////////////////////////////////////////////////
// DN42 Prometheus Stats Server
//////////////////////////////////////////////////////////////////////////

package main

//////////////////////////////////////////////////////////////////////////

import (
	"context"
	"github.com/gorilla/mux"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	log "github.com/sirupsen/logrus"
	flag "github.com/spf13/pflag"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"
)

//////////////////////////////////////////////////////////////////////////

type Metric interface {
	Register()
	Collect()
}

//////////////////////////////////////////////////////////////////////////
// Set the log level

func setLogLevel(levelStr string) {

	if level, err := log.ParseLevel(levelStr); err != nil {
		// failed to set the level

		// set a sensible default and, of course, log the error
		log.SetLevel(log.InfoLevel)
		log.WithFields(log.Fields{
			"loglevel": levelStr,
			"error":    err,
		}).Error("Failed to set requested log level")

	} else {

		// set the requested level
		log.SetLevel(level)

	}
}

//////////////////////////////////////////////////////////////////////////
// collect metrics

func collectMetrics(stop chan bool,
	notify *sync.WaitGroup, interval time.Duration,
	metrics []Metric) {

	notify.Add(1)
	defer notify.Done()

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	log.WithFields(log.Fields{
		"Interval": interval,
	}).Info("Starting data collection")

	for {
		select {
		case <-stop:
			// stop updating
			return
		case <-ticker.C:
			// timer expired, perform an update

			for _, metric := range metrics {
				metric.Collect()
			}

		}
	}

}

//////////////////////////////////////////////////////////////////////////
// initialise metrics

func initMetrics() []Metric {

	metrics := make([]Metric, 1)

	metrics[0] = &DNSMetrics{}

	return metrics
}

//////////////////////////////////////////////////////////////////////////

func main() {

	// set a default log level, so that logging can be used immediately
	// the level will be overidden later on once the command line
	// options are loaded
	log.SetLevel(log.InfoLevel)
	log.Info("DN42 Stats Server Starting")

	// declare cmd line options
	var (
		logLevel        = flag.StringP("LogLevel", "l", "Info", "Log level")
		bindAddress     = flag.StringP("BindAddress", "b", ":8001", "Server bind address")
		refreshInterval = flag.StringP("Refresh", "i", "1m", "Refresh interval")
	)
	flag.Parse()

	// now initialise logging properly based on the cmd line options
	setLogLevel(*logLevel)

	// parse the refreshInterval and start data collection
	interval, err := time.ParseDuration(*refreshInterval)
	if err != nil {
		log.WithFields(log.Fields{
			"error":    err,
			"interval": *refreshInterval,
		}).Fatal("Unable to parse refresh interval")
	}

	// initialise and register metrics
	metrics := initMetrics()

	for _, metric := range metrics {
		metric.Register()
	}

	// start metric collection
	notify_complete := &sync.WaitGroup{}
	stop_collection := make(chan bool)
	go collectMetrics(stop_collection, notify_complete, interval, metrics)

	// initialise router and install prom handler
	router := mux.NewRouter()
	router.Handle("/metrics", promhttp.Handler())

	// initialise http server
	server := &http.Server{
		Addr:         *bindAddress,
		WriteTimeout: time.Second * 15,
		ReadTimeout:  time.Second * 15,
		IdleTimeout:  time.Second * 60,
		Handler:      router,
	}

	// run the server in a non-blocking goroutine

	log.WithFields(log.Fields{
		"BindAddress": *bindAddress,
	}).Info("Starting server")

	go func() {
		if err := server.ListenAndServe(); err != nil {
			log.WithFields(log.Fields{
				"error":       err,
				"BindAddress": *bindAddress,
			}).Fatal("Unable to start server")
		}
	}()

	// graceful shutdown via SIGINT (^C)
	csig := make(chan os.Signal, 1)
	signal.Notify(csig, os.Interrupt)

	// and block
	<-csig

	log.Info("Server shutting down")

	// deadline for server to shutdown
	ctx, cancel := context.WithTimeout(context.Background(), 10)
	defer cancel()

	// shutdown stats collection and the server
	close(stop_collection)
	notify_complete.Wait()
	server.Shutdown(ctx)

	// nothing left to do
	log.Info("Shutdown complete, all done")
	os.Exit(0)
}

//////////////////////////////////////////////////////////////////////////
// end of code