prometheus-latency-exporter/main.go

package main

import (
	"bufio"
	"database/sql"
	"flag"
	"fmt"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type Service struct {
	Name     string `yaml:"name"`
	Host     string `yaml:"host"`
	Port     string `yaml:"port"`
	Type     string `yaml:"type"`
	Username string `yaml:"username"`
	Password string `yaml:"password"`
}

type Config struct {
	Services []Service `yaml:"services"`
}

var (
	latency = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "service_latency_seconds",
			Help: "Latency of services.",
		},
		[]string{"service"},
	)
	addr  = flag.String("a", "0.0.0.0", "address to use")
	port  = flag.String("p", "9062", "port to run on")
	conf  = flag.String("c", "config.yml", "configuration file to use")
	inter = flag.Int("i", 30, "time interval in seconds between two checks")
	debug = flag.Bool("d", false, "enable debug logs")
)

func init() {
	prometheus.MustRegister(latency)
}

func main() {
	flag.Parse()
	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		fmt.Fprintln(w, "Prometheus Latency Exporter: <a href='/metrics'>Metrics</a>")
	})
	go func() {
		for {
			checkLatency()
			time.Sleep(time.Duration(*inter) * time.Second)
		}
	}()
	log.Println("INFO: Starting server on address:", *addr, "port:", *port)
	err := http.ListenAndServe(net.JoinHostPort(*addr, *port), nil)
	if err != nil {
		log.Fatal("ERROR: Failed to start server, error: ", err)
	}
}

func checkLatency() {
	var config Config
	source, err := ioutil.ReadFile(*conf)
	if err != nil {
		log.Fatalf("ERROR: Failed to read configuration file, error: %v", err)
	}

	err = yaml.Unmarshal(source, &config)
	if err != nil {
		log.Fatalf("ERROR: Failed to unmarshal configuration, error: %v", err)
	}

	for _, service := range config.Services {
		if service.Type == "" {
			service.Type = "tcp"
		}
		if *debug == true {
			log.Printf("DEBUG: Dialing %s: host: %s, port: %s, type: %s",
				service.Name, service.Host, service.Port, service.Type)
		}

		start := time.Now()

		if strings.ToLower(service.Type) == "mysql" {
			dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/", service.Username, service.Password, service.Host, service.Port)
			db, err := sql.Open("mysql", dsn)
			if err != nil {
				log.Printf("ERROR: Failed to check service: %s, error: %s", service.Name, err)
				continue
			}
			err = db.Ping()
			if err != nil {
				log.Printf("ERROR: Failed to perform MySQL check on: %s, error: %s", service.Name, err)
				continue
			}
			defer db.Close()

		} else {
			conn, err := net.Dial("tcp", net.JoinHostPort(service.Host, service.Port))
			if err != nil {
				log.Printf("ERROR: Failed to check service: %s, error: %s", service.Name, err)
				continue
			}

			// Performs a Redis check if needed
			if strings.ToLower(service.Type) == "redis" {
				_, err := conn.Write([]byte("PING\r\n"))
				if err != nil {
					log.Printf("ERROR: Failed to check service: %s, error: %s", service.Name, err)
					continue
				}
				reply, _ := bufio.NewReader(conn).ReadString('\n')
				if reply != "+PONG\r\n" {
					log.Printf("ERROR: Bad answer from %s (type redis): %s", service.Name, reply)
					continue
				}
			}
			conn.Close()
		}

		duration := time.Since(start)
		latency.With(prometheus.Labels{"service": service.Name}).Set(duration.Seconds())
		if *debug == true {
			log.Printf("DEBUG: Latency %s: %v seconds", service.Name, duration.Seconds())
		}
	}
}