consul_collector.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"os/exec"
	"strconv"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	consulAgent = "consul agent"
)

func init() {
	prometheus.MustRegister(newConsulCollector("consul info"))
	prometheus.MustRegister(prometheus.NewProcessCollectorPIDFn(consulAgentPid, "consul"))
}

// consulCollector implements the prometheus.Collector interface.
type consulCollector struct {
	info    string
	metrics map[string]prometheus.Gauge
}

func newConsulCollector(info string) prometheus.Collector {
	return &consulCollector{
		info:    info,
		metrics: map[string]prometheus.Gauge{},
	}
}

func (c *consulCollector) Collect(metricc chan<- prometheus.Metric) {
	err := c.updateMetrics()
	if err != nil {
		return
	}

	for _, m := range c.metrics {
		m.Collect(metricc)
	}
}

func (c *consulCollector) Describe(descc chan<- *prometheus.Desc) {
	err := c.updateMetrics()
	if err != nil {
		return
	}

	for _, m := range c.metrics {
		descc <- m.Desc()
	}
}

func (c *consulCollector) updateMetrics() error {
	stats, err := getConsulStats(c.info)
	if err != nil {
		return fmt.Errorf("consul info failed: %s", err)
	}

	for name, value := range stats {
		if _, ok := c.metrics[name]; !ok {
			c.metrics[name] = prometheus.NewGauge(
				prometheus.GaugeOpts{
					Namespace: "consul",
					Name:      name,
					Help:      fmt.Sprintf("%s from consul info", name),
				},
			)
		}

		c.metrics[name].Set(float64(value))
	}

	return nil
}

type consulStats map[string]int64

func getConsulStats(info string) (consulStats, error) {
	cmd := strings.Split(info, " ")
	output, err := exec.Command(cmd[0], cmd[1:]...).Output()
	if err != nil {
		return nil, err
	}

	stats, err := parseConsulStats(output)
	if err != nil {
		return nil, err
	}

	return stats, nil
}

func parseConsulStats(b []byte) (consulStats, error) {
	var (
		s     = bufio.NewScanner(bytes.NewBuffer(b))
		stats = consulStats{}

		ignoredFields = map[string]struct{}{
			"arch":         struct{}{},
			"os":           struct{}{},
			"state":        struct{}{},
			"version":      struct{}{},
			"last_contact": struct{}{},
		}
		ignoredKeys = map[string]struct{}{
			"build": struct{}{},
		}

		key string
	)

	for s.Scan() {
		line := strings.TrimSpace(s.Text())

		if strings.Contains(line, ":") {
			key = strings.TrimSuffix(line, ":")
		}

		if _, ok := ignoredKeys[key]; ok {
			continue
		}

		if strings.Contains(line, "=") {
			var (
				parts        = strings.SplitN(line, "=", 2)
				field, value = strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
				name         = strings.Join([]string{key, field}, "_")
			)

			if _, ok := ignoredFields[field]; ok {
				continue
			}

			switch value {
			case "true":
				stats[name] = 1
			case "false":
				stats[name] = 0
			case "never":
				stats[name] = 0
			default:
				i, err := strconv.ParseInt(value, 10, 64)
				if err != nil {
					return nil, err
				}
				stats[name] = i
			}
		}
	}
	if err := s.Err(); err != nil {
		return nil, err
	}

	return stats, nil
}

func consulAgentPid() (int, error) {
	out, err := exec.Command("pgrep", "-f", consulAgent).Output()
	if err != nil {
		return 0, fmt.Errorf("could not get pid of %s: %s", consulAgent, err)
	}
	pid, err := strconv.Atoi(strings.TrimSpace(string(out)))
	if err != nil {
		return 0, fmt.Errorf("could not parse pid of %s: %s", consulAgent, err)
	}
	return pid, nil
}