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 (
	staticMetricsFile    string
	prometheusDocFile    string
	generatedMetricsFile string
	dryRun               bool

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

func main() {
	flag.StringVar(&staticMetricsFile, "static-metrics", "scripts/metricsdocgen/metrics", "Path to static metrics file (manually maintained)")
	flag.StringVar(&generatedMetricsFile, "generated-metrics", "scripts/metricsdocgen/generated_metrics", "Path to generated metrics file (from scanner)")
	flag.StringVar(&prometheusDocFile, "prometheus-doc-file", "docs/admin/integrations/prometheus.md", "Path to Prometheus doc file")
	flag.BoolVar(&dryRun, "dry-run", false, "Dry run")
	flag.Parse()

	metrics, err := readAndMergeMetrics()
	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)
	}
}

// readMetricsFromFile reads metrics from a single Prometheus text format file.
func readMetricsFromFile(path string) ([]*dto.MetricFamily, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, xerrors.Errorf("can't open metrics file %s: %w", path, err)
	}
	defer f.Close()

	var metrics []*dto.MetricFamily

	decoder := expfmt.NewDecoder(f, expfmt.NewFormat(expfmt.TypeTextPlain))
	for {
		var m dto.MetricFamily
		err = decoder.Decode(&m)
		if errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return nil, xerrors.Errorf("decoding metrics from %s: %w", path, err)
		}
		metrics = append(metrics, &m)
	}

	return metrics, nil
}

// readAndMergeMetrics reads metrics from both generated and static files,
// merges them, and returns a sorted list. Generated metrics are produced
// by the AST scanner that extracts metric definitions from the coder source
// code while static metrics are manually maintained (e.g., go_*, process_*,
// external dependencies).
// Note: Static metrics take priority over generated metrics, allowing manual
// overrides for metrics that can't be accurately extracted by the scanner.
func readAndMergeMetrics() ([]*dto.MetricFamily, error) {
	generatedMetrics, err := readMetricsFromFile(generatedMetricsFile)
	if err != nil {
		return nil, xerrors.Errorf("reading generated metrics: %w", err)
	}

	staticMetrics, err := readMetricsFromFile(staticMetricsFile)
	if err != nil {
		return nil, xerrors.Errorf("reading static metrics: %w", err)
	}

	// Merge metrics, using a map to deduplicate by name.
	metricsByName := make(map[string]*dto.MetricFamily)

	// Add generated metrics first.
	for _, m := range generatedMetrics {
		metricsByName[*m.Name] = m
	}

	// Static metrics overwrite generated metrics if they exist.
	for _, m := range staticMetrics {
		metricsByName[*m.Name] = m
	}

	// Convert back to slice and sort.
	var metrics []*dto.MetricFamily
	for _, m := range metricsByName {
		metrics = append(metrics, m)
	}

	sort.Slice(metrics, func(i, j int) bool {
		return *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, 0o644)
	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
}