converge/cmd/converge/prometheus.go

package main

import (
	"converge/pkg/models"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"log"
	"net/http"
)

const NAMESPACE = "converge"

var (
	// remember previous values of agent guids and clients so that we can increment
	// the cumulative counters.
	lastAgentGuids  map[string]bool = make(map[string]bool)
	lastClientGuids map[string]bool = make(map[string]bool)

	cumulativeAgentCount = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: NAMESPACE,
		Name:      "agent_count_total",
		Help:      "Total number of agents connected over time",
	})
	cumulativeClientCount = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: NAMESPACE,
		Name:      "client_count_total",
		Help:      "Total number of clients connected over time",
	})

	agentCount = promauto.NewGauge(prometheus.GaugeOpts{
		Namespace: NAMESPACE,
		Name:      "agent_count",
		Help:      "Current number of agents",
	})
	clientCount = promauto.NewGauge(prometheus.GaugeOpts{
		Namespace: NAMESPACE,
		Name:      "client_count",
		Help:      "Current number of clients",
	})

	agentStartTime = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Namespace: NAMESPACE,
		Name:      "agent_start_time_seconds",
		Help:      "Time the agent started",
	}, []string{"agent_guid"})
	clientStartTime = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Namespace: NAMESPACE,
		Name:      "client_start_time_seconds",
		Help:      "Time the client started",
	}, []string{"client_guid"})

	agentInfo = promauto.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: NAMESPACE,
			Name:      "agent_info",
			Help:      "A flexible gauge with dynamic labels, always set to 1",
		},
		[]string{
			"agent_guid",
			"agent_id",
			"agent_username",
			"agent_hostname",
			"agent_pwd",
			"agent_os",
			"agent_shell",
		})

	clientInfo = promauto.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: NAMESPACE,
			Name:      "client_info",
			Help:      "A flexible gauge with dynamic labels, always set to 1",
		},
		[]string{"client_guid",
			"client_id",
			"agent_id",
			"agent_guid",
			"client_sessiontype",
			"client_username",
			"client_hostname",
			"client_pwd",
			"client_os",
			"client_shell",
		}, // Label names
	)
)

func agentLabels(agent models.Agent) prometheus.Labels {
	return prometheus.Labels{
		"agent_guid":     agent.Guid,
		"agent_id":       agent.PublicId,
		"agent_username": agent.EnvironmentInfo.Username,
		"agent_hostname": agent.EnvironmentInfo.Hostname,
		"agent_pwd":      agent.EnvironmentInfo.Pwd,
		"agent_os":       agent.EnvironmentInfo.OS,
		"agent_shell":    agent.EnvironmentInfo.Shell,
	}
}

func clientLabels(client models.Client) prometheus.Labels {
	return prometheus.Labels{
		"client_guid":        client.Guid,
		"client_id":          client.ClientId,
		"agent_id":           client.PublicId,
		"agent_guid":         client.AgentGuid,
		"client_sessiontype": client.SessionType,
		"client_username":    client.EnvironmentInfo.Username,
		"client_hostname":    client.EnvironmentInfo.Hostname,
		"client_pwd":         client.EnvironmentInfo.Pwd,
		"client_os":          client.EnvironmentInfo.OS,
		"client_shell":       client.EnvironmentInfo.Shell,
	}
}

func agentActive(agent models.Agent) {
	agentInfo.With(agentLabels(agent)).Set(1)
}

func clientActive(client models.Client) {
	clientInfo.With(clientLabels(client)).Set(1)
}

func setupPrometheus(mux *http.ServeMux, notifications chan *models.State) {
	go func() {
		for {
			state := <-notifications
			updateMetrics(state)
		}
	}()
	mux.Handle("/metrics", promhttp.Handler())
}

func updateMetrics(state *models.State) {
	// This implemnetation has a small probability that the metric will be in a partially
	// initialized state. This is however unlikely. It would lead to in incorrect determination
	// that an agent or client is not available. However, each agent and client will have a UID
	// so that is still possible to identify the client or agent even though some values might
	// become 0.

	log.Printf("Got notification %v", *state)

	agentGuids := make(map[string]bool)
	clientGuids := make(map[string]bool)

	agentCount.Set(float64(len(state.Agents)))
	agentInfo.Reset()
	agentStartTime.Reset()
	for _, agent := range state.Agents {
		if !lastAgentGuids[agent.Guid] {
			agentStartTime.
				With(prometheus.Labels{"agent_guid": agent.Guid}).
				Set(float64(agent.StartTime.Unix()))
			cumulativeAgentCount.Inc()
		}
		agentGuids[agent.Guid] = true
		agentActive(agent)
	}
	lastAgentGuids = agentGuids

	clientCount.Set(float64(len(state.Clients)))
	clientInfo.Reset()
	clientStartTime.Reset()
	for _, client := range state.Clients {
		if !lastClientGuids[client.Guid] {
			clientStartTime.With(prometheus.Labels{"client_guid": client.Guid}).Set(float64(client.StartTime.Unix()))
			cumulativeClientCount.Inc()
		}
		clientGuids[client.Guid] = true
		clientActive(client)
	}
	lastClientGuids = clientGuids
}