smallwebwaf/internal/metrics/busiest.go

package metrics

import (
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

// other is the label value under which the countries or AS numbers
// outside the busiest are counted.
const other = "other"

// busiest are the metrics by one thing the lookup finds of the client,
// its country or its AS number, for requests whose client's is known.
// The topN busiest countries or AS numbers, by their requests since the
// start, have series of their own, and the others are counted under
// other, so that there are never more than topN + 1 series. One that
// drops out of the busiest loses its series, and its next requests are
// counted under other; one that becomes one of them gets a series that
// counts from then on. Each series therefore only ever goes up.
type busiest struct {
	topN int

	requests      *prometheus.CounterVec
	requestBytes  *prometheus.CounterVec
	responseBytes *prometheus.CounterVec
	// refused are, by country, the requests the country lists refused; nil
	// by AS number.
	refused *prometheus.CounterVec

	mu sync.Mutex
	// seen is each country's or AS number's requests since the start, by
	// which they are ranked.
	seen map[string]int64
	// top are those with series of their own.
	top map[string]bool
}

// newCountries returns the metrics by country, with series of their own
// for the topN busiest countries.
func newCountries(topN int) *busiest {
	countries := newBusiest(topN, "country", "the client's country")
	countries.refused = counterVec("smallwebwaf_country_list_refusals_total",
		"Requests the country lists refused, by the client's country.",
		[]string{"country"})

	return countries
}

// newASNs returns the metrics by AS number, with series of their own for
// the topN busiest AS numbers.
func newASNs(topN int) *busiest {
	return newBusiest(topN, "asn", "the client's AS number")
}

// newBusiest returns the metrics by label, which is described as
// description, with series of their own for the topN busiest values.
func newBusiest(topN int, label, description string) *busiest {
	by := []string{label}

	return &busiest{
		topN: topN,
		requests: counterVec("smallwebwaf_"+label+"_requests_total",
			"Requests, by "+description+".", by),
		requestBytes: counterVec("smallwebwaf_"+label+"_request_bytes_total",
			"Request body bytes, by "+description+".", by),
		responseBytes: counterVec("smallwebwaf_"+label+"_response_bytes_total",
			"Response body bytes, by "+description+".", by),
		seen: map[string]int64{},
		top:  map[string]bool{},
	}
}

// Describe and Collect make the metrics a prometheus.Collector, so that
// they are registered together.
func (b *busiest) Describe(ch chan<- *prometheus.Desc) {
	for _, vec := range b.vecs() {
		vec.Describe(ch)
	}
}

// Collect is the other half of prometheus.Collector, with Describe.
func (b *busiest) Collect(ch chan<- prometheus.Metric) {
	for _, vec := range b.vecs() {
		vec.Collect(ch)
	}
}

// vecs returns the metrics: by AS number, those of requests and bytes; by
// country, the refusals by the country lists as well.
func (b *busiest) vecs() []*prometheus.CounterVec {
	vecs := []*prometheus.CounterVec{b.requests, b.requestBytes, b.responseBytes}
	if b.refused != nil {
		vecs = append(vecs, b.refused)
	}

	return vecs
}

// add counts a request from its log line, whose client's country or AS
// number, value, is known.
func (b *busiest) add(value string, line *requestlog.Line) {
	b.mu.Lock()
	defer b.mu.Unlock()

	b.seen[value]++

	label := b.label(value)
	b.requests.WithLabelValues(label).Inc()
	b.requestBytes.WithLabelValues(label).Add(float64(line.RequestBytes))
	b.responseBytes.WithLabelValues(label).Add(float64(line.ResponseBytes))

	if b.refused != nil && line.Action == requestlog.ActionCountryDenied {
		b.refused.WithLabelValues(label).Inc()
	}
}

// label returns the label a request from value is counted under: value
// while it is one of the busiest, other while it is not. A value busier
// than the least busy of them takes its place, and that one's series are
// dropped.
func (b *busiest) label(value string) string {
	if b.top[value] {
		return value
	}

	if len(b.top) < b.topN {
		b.top[value] = true

		return value
	}

	least := ""

	for top := range b.top {
		if least == "" || b.seen[top] < b.seen[least] {
			least = top
		}
	}

	if b.seen[value] <= b.seen[least] {
		return other
	}

	delete(b.top, least)

	for _, vec := range b.vecs() {
		vec.DeleteLabelValues(least)
	}

	b.top[value] = true

	return value
}