coder/scripts/metricsdocgen/main.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"io"
	"log"
	"os"
	"sort"
	"strings"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"golang.org/x/xerrors"
)

var (
	metricsFile       string
	prometheusDocFile string
	dryRun            bool

	generatorPrefix = []byte("<!-- Code generated by 'make docs/admin/prometheus.md'. DO NOT EDIT -->")
	generatorSuffix = []byte("<!-- End generated by 'make docs/admin/prometheus.md'. -->")
)

func main() {
	flag.StringVar(&metricsFile, "metrics-file", "scripts/metricsdocgen/metrics", "Path to Prometheus metrics file")
	flag.StringVar(&prometheusDocFile, "prometheus-doc-file", "docs/admin/prometheus.md", "Path to Prometheus doc file")
	flag.BoolVar(&dryRun, "dry-run", false, "Dry run")
	flag.Parse()

	metrics, err := readMetrics()
	if err != nil {
		log.Fatal("can't read metrics: ", err)
	}

	doc, err := readPrometheusDoc()
	if err != nil {
		log.Fatal("can't read Prometheus doc: ", err)
	}

	doc, err = updatePrometheusDoc(doc, metrics)
	if err != nil {
		log.Fatal("can't update Prometheus doc: ", err)
	}

	if dryRun {
		log.Println(string(doc))
		return
	}

	err = writePrometheusDoc(doc)
	if err != nil {
		log.Fatal("can't write updated Prometheus doc: ", err)
	}
}

func readMetrics() ([]dto.MetricFamily, error) {
	f, err := os.Open(metricsFile)
	if err != nil {
		return nil, xerrors.New("can't open metrics file")
	}

	var metrics []dto.MetricFamily

	decoder := expfmt.NewDecoder(f, expfmt.FmtProtoText)
	for {
		var m dto.MetricFamily
		err = decoder.Decode(&m)
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return nil, err
		}
		metrics = append(metrics, m)
	}

	sort.Slice(metrics, func(i, j int) bool {
		return sort.StringsAreSorted([]string{*metrics[i].Name, *metrics[j].Name})
	})
	return metrics, nil
}

func readPrometheusDoc() ([]byte, error) {
	doc, err := os.ReadFile(prometheusDocFile)
	if err != nil {
		return nil, err
	}
	return doc, nil
}

func updatePrometheusDoc(doc []byte, metricFamilies []dto.MetricFamily) ([]byte, error) {
	i := bytes.Index(doc, generatorPrefix)
	if i < 0 {
		return nil, xerrors.New("generator prefix tag not found")
	}
	tableStartIndex := i + len(generatorPrefix) + 1

	j := bytes.Index(doc[tableStartIndex:], generatorSuffix)
	if j < 0 {
		return nil, xerrors.New("generator suffix tag not found")
	}
	tableEndIndex := tableStartIndex + j

	var buffer bytes.Buffer
	_, _ = buffer.Write(doc[:tableStartIndex])
	_ = buffer.WriteByte('\n')

	_, _ = buffer.WriteString("| Name | Type | Description | Labels |\n")
	_, _ = buffer.WriteString("| - | - | - | - |\n")
	for _, mf := range metricFamilies {
		_, _ = buffer.WriteString("| ")
		_, _ = buffer.Write([]byte("`" + *mf.Name + "`"))
		_, _ = buffer.WriteString(" | ")
		_, _ = buffer.Write([]byte(strings.ToLower(mf.Type.String())))
		_, _ = buffer.WriteString(" | ")
		if mf.Help != nil {
			_, _ = buffer.Write([]byte(*mf.Help))
		}
		_, _ = buffer.WriteString(" | ")

		labels := map[string]struct{}{}
		metrics := mf.GetMetric()
		for _, m := range metrics {
			for _, label := range m.Label {
				labels["`"+*label.Name+"`"] = struct{}{}
			}
		}

		if len(labels) > 0 {
			_, _ = buffer.WriteString(strings.Join(sortedKeys(labels), " "))
		}

		_, _ = buffer.WriteString(" |\n")
	}

	_ = buffer.WriteByte('\n')
	_, _ = buffer.Write(doc[tableEndIndex:])
	return buffer.Bytes(), nil
}

func writePrometheusDoc(doc []byte) error {
	// G306: Expect WriteFile permissions to be 0600 or less
	/* #nosec G306 */
	err := os.WriteFile(prometheusDocFile, doc, 0644)
	if err != nil {
		return err
	}
	return nil
}

func sortedKeys(m map[string]struct{}) []string {
	var keys []string
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}