Get all TiKV stores from PD, then get all metrics of TiKVs and return to Prometheus

tikv_metric_grasper.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"sort"

	"github.com/golang/protobuf/proto"
	"github.com/pingcap/kvproto/pkg/debugpb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"google.golang.org/grpc"
)

func perror(err error) {
	if err != nil {
		fmt.Printf("%s\n", err)
		os.Exit(1)
	}
}

func main() {
	pdConn, err := grpc.Dial("127.0.0.1:2379", grpc.WithInsecure())
	perror(err)
	defer pdConn.Close()

	pdClient := pdpb.NewPDClient(pdConn)
	ctx := context.Background()

	members, err := pdClient.GetMembers(ctx, &pdpb.GetMembersRequest{})
	perror(err)

	clusterID := members.GetHeader().GetClusterId()

	stores, err := pdClient.GetAllStores(ctx, &pdpb.GetAllStoresRequest{
		Header: &pdpb.RequestHeader{
			ClusterId: clusterID,
		},
	})
	perror(err)

	allMetrics := make([]*dto.MetricFamily, 0, 1024)
	for _, store := range stores.GetStores() {
		fmt.Printf("%d %s\n", store.GetId(), store.GetAddress())

		labels := map[string]string{
			"job":      fmt.Sprintf("tikv_%d", store.GetId()),
			"instance": store.GetAddress(),
		}

		tikvConn, err := grpc.Dial(store.GetAddress(), grpc.WithInsecure())
		perror(err)

		tikvClient := debugpb.NewDebugClient(tikvConn)
		metrics, err := tikvClient.GetMetrics(ctx, &debugpb.GetMetricsRequest{})
		perror(err)

		mData := metrics.GetPrometheus()

		var parser expfmt.TextParser
		metricFamilies, err := parser.TextToMetricFamilies(bytes.NewBufferString(mData))
		perror(err)

		sanitizeLabels(metricFamilies, labels)

		for _, m := range metricFamilies {
			allMetrics = append(allMetrics, m)
		}

		tikvConn.Close()
	}

	for _, m := range allMetrics {
		var w bytes.Buffer
		// TODO: we can return all to prometheus when GET /metrics
		expfmt.MetricFamilyToText(&w, m)
		fmt.Printf("%s", w.String())
	}
}

// Refer https://github.com/prometheus/pushgateway/blob/master/handler/push.go#L235
// But here just add job and instance to the labels simply
func sanitizeLabels(
	metricFamilies map[string]*dto.MetricFamily,
	groupingLabels map[string]string,
) {
	for _, mf := range metricFamilies {
		for _, m := range mf.GetMetric() {
			for key, value := range groupingLabels {
				l := &dto.LabelPair{
					Name:  proto.String(key),
					Value: proto.String(value),
				}
				m.Label = append(m.Label, l)
			}
			sort.Sort(prometheus.LabelPairSorter(m.Label))
		}
	}
}