backend: optimize updating command metrics (#467)

Co-authored-by: xhe <[email protected]>

Author: djshow832

pkg/proxy/backend/metrics.go

@@ -4,21 +4,55 @@
 package backend
 
 import (
+	"sync"
 	"time"
 
 	"github.com/pingcap/tiproxy/pkg/metrics"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/pingcap/tiproxy/pkg/util/monotime"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
-func addCmdMetrics(cmd pnet.Command, addr string, startTime monotime.Time) {
-	label := cmd.String()
-	metrics.QueryTotalCounter.WithLabelValues(addr, label).Inc()
+type mcPerCmd struct {
+	counter  prometheus.Counter
+	observer prometheus.Observer
+}
 
+type cmdMetricsCache struct {
+	sync.Mutex
 	// The duration labels are different with TiDB: Labels in TiDB are statement types.
 	// However, the proxy is not aware of the statement types, so we use command types instead.
+	metrics map[string]*[pnet.ComEnd]mcPerCmd
+}
+
+func newCmdMetricsCache() cmdMetricsCache {
+	return cmdMetricsCache{
+		metrics: make(map[string]*[pnet.ComEnd]mcPerCmd),
+	}
+}
+
+var cache = newCmdMetricsCache()
+
+func addCmdMetrics(cmd pnet.Command, addr string, startTime monotime.Time) {
+	cache.Lock()
+	defer cache.Unlock()
+
+	addrMc, ok := cache.metrics[addr]
+	if !ok {
+		addrMc = &[pnet.ComEnd]mcPerCmd{}
+		cache.metrics[addr] = addrMc
+	}
+	mc := &addrMc[cmd]
+	if mc.counter == nil {
+		label := cmd.String()
+		*mc = mcPerCmd{
+			counter:  metrics.QueryTotalCounter.WithLabelValues(addr, label),
+			observer: metrics.QueryDurationHistogram.WithLabelValues(addr, label),
+		}
+	}
+	mc.counter.Inc()
 	cost := monotime.Since(startTime)
-	metrics.QueryDurationHistogram.WithLabelValues(addr, label).Observe(cost.Seconds())
+	mc.observer.Observe(cost.Seconds())
 }
 
 func readCmdCounter(cmd pnet.Command, addr string) (int, error) {

pkg/proxy/backend/metrics_test.go

@@ -0,0 +1,101 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package backend
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/pingcap/tiproxy/pkg/metrics"
+	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
+	"github.com/pingcap/tiproxy/pkg/util/monotime"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func BenchmarkAddCmdMetrics(b *testing.B) {
+	cmd := pnet.ComQuery
+	addr := "127.0.0.1:4000"
+	startTime := monotime.Now()
+	for i := 0; i < b.N; i++ {
+		addCmdMetrics(cmd, addr, startTime)
+	}
+}
+
+func BenchmarkAddCmdMetricsArrayPointer(b *testing.B) {
+	cmd := pnet.ComQuery
+	addr := "127.0.0.1:4000"
+	var lock sync.Mutex
+	counter := make(map[string]*[pnet.ComEnd]prometheus.Counter)
+	for i := 0; i < b.N; i++ {
+		lock.Lock()
+
+		addrCounter, ok := counter[addr]
+		if !ok {
+			addrCounter = &[pnet.ComEnd]prometheus.Counter{}
+			counter[addr] = addrCounter
+		}
+		counter := addrCounter[cmd]
+		if counter == nil {
+			counter = metrics.QueryTotalCounter.WithLabelValues(addr, cmd.String())
+			addrCounter[cmd] = counter
+		}
+		counter.Inc()
+
+		lock.Unlock()
+	}
+}
+
+func BenchmarkAddCmdMetricsArray(b *testing.B) {
+	cmd := pnet.ComQuery
+	addr := "127.0.0.1:4000"
+	var lock sync.Mutex
+	counter := make(map[string][pnet.ComEnd]prometheus.Counter)
+	for i := 0; i < b.N; i++ {
+		lock.Lock()
+
+		addrCounter, ok := counter[addr]
+		if !ok {
+			addrCounter = [pnet.ComEnd]prometheus.Counter{}
+			counter[addr] = addrCounter
+		}
+		counter := addrCounter[cmd]
+		if counter == nil {
+			counter = metrics.QueryTotalCounter.WithLabelValues(addr, cmd.String())
+			addrCounter[cmd] = counter
+		}
+		counter.Inc()
+
+		lock.Unlock()
+	}
+}
+
+func BenchmarkAddCmdMetricsMap(b *testing.B) {
+	cmd := pnet.ComQuery
+	addr := "127.0.0.1:4000"
+	var lock sync.Mutex
+	counter := make(map[string]map[pnet.Command]prometheus.Counter)
+	for i := 0; i < b.N; i++ {
+		lock.Lock()
+
+		addrCounter, ok := counter[addr]
+		if !ok {
+			addrCounter = make(map[pnet.Command]prometheus.Counter)
+			counter[addr] = addrCounter
+		}
+		counter, ok := addrCounter[cmd]
+		if !ok {
+			counter = metrics.QueryTotalCounter.WithLabelValues(addr, cmd.String())
+			addrCounter[cmd] = counter
+		}
+		counter.Inc()
+
+		lock.Unlock()
+	}
+}
+
+func BenchmarkAddCmdMetricsSimple(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		metrics.QueryTotalCounter.WithLabelValues("127.0.0.1:4000", pnet.ComQuery.String()).Inc()
+	}
+}