meter, config: meter public and private response traffic independently (#993)

djshow832 · web-flow · commit 704568c68648 · 2025-11-25T08:52:32.000Z
diff --git a/lib/config/proxy.go b/lib/config/proxy.go
@@ -61,6 +61,8 @@ type ProxyServerOnline struct {
 	// In k8s, the pod terminationGracePeriodSeconds can be set to very long so that these configs can be updated online.
 	GracefulWaitBeforeShutdown int `yaml:"graceful-wait-before-shutdown,omitempty" toml:"graceful-wait-before-shutdown,omitempty" json:"graceful-wait-before-shutdown,omitempty" reloadable:"true"`
 	GracefulCloseConnTimeout   int `yaml:"graceful-close-conn-timeout,omitempty" toml:"graceful-close-conn-timeout,omitempty" json:"graceful-close-conn-timeout,omitempty" reloadable:"true"`
+	// Public and private traffic are metered separately.
+	PublicEndpoints []string `yaml:"public-endpoints,omitempty" toml:"public-endpoints,omitempty" json:"public-endpoints,omitempty" reloadable:"true"`
 }
 
 type ProxyServer struct {
diff --git a/pkg/balance/router/group.go b/pkg/balance/router/group.go
@@ -17,6 +17,7 @@ import (
 	"github.com/pingcap/tiproxy/pkg/balance/observer"
 	"github.com/pingcap/tiproxy/pkg/balance/policy"
 	"github.com/pingcap/tiproxy/pkg/metrics"
+	"github.com/pingcap/tiproxy/pkg/util/netutil"
 	"go.uber.org/zap"
 )
 
@@ -76,21 +77,15 @@ func NewGroup(values []string, bpCreator func(lg *zap.Logger) policy.BalancePoli
 }
 
 func (g *Group) parseValues() error {
-	var parseErr error
 	switch g.matchType {
 	case MatchClientCIDR, MatchProxyCIDR:
-		cidrList := make([]*net.IPNet, 0, len(g.values))
-		for _, v := range g.values {
-			_, cidr, err := net.ParseCIDR(v)
-			if err == nil {
-				cidrList = append(cidrList, cidr)
-			} else {
-				parseErr = err
-			}
+		cidrList, parseErr := netutil.ParseCIDRList(g.values)
+		if parseErr != nil {
+			return parseErr
 		}
 		g.cidrList = cidrList
 	}
-	return parseErr
+	return nil
 }
 
 func (g *Group) Match(clientInfo ClientInfo) bool {
@@ -100,26 +95,11 @@ func (g *Group) Match(clientInfo ClientInfo) bool {
 		if g.matchType == MatchClientCIDR {
 			addr = clientInfo.ClientAddr
 		}
-		if addr == nil || reflect.ValueOf(addr).IsNil() {
-			return false
-		}
-		value := addr.String()
-		ipStr, _, err := net.SplitHostPort(value)
+		contains, err := netutil.CIDRContainsIP(g.cidrList, addr)
 		if err != nil {
-			g.lg.Error("parsing address failed", zap.String("addr", value), zap.Error(err))
-			return false
+			g.lg.Error("checking CIDR failed", zap.String("addr", addr.String()), zap.Error(err))
 		}
-		ip := net.ParseIP(ipStr)
-		if ip == nil {
-			g.lg.Error("parsing IP failed", zap.String("ip", value))
-			return false
-		}
-		for _, cidr := range g.cidrList {
-			if cidr.Contains(ip) {
-				return true
-			}
-		}
-		return false
+		return contains
 	}
 	return true
 }
diff --git a/pkg/balance/router/group_test.go b/pkg/balance/router/group_test.go
@@ -26,7 +26,7 @@ func TestParseCIDR(t *testing.T) {
 	}{
 		{
 			cidrs:   []string{"1.1.1.1"},
-			success: false,
+			success: true,
 		},
 		{
 			cidrs:   []string{"1.1.1.1/32"},
diff --git a/pkg/manager/meter/meter.go b/pkg/manager/meter/meter.go
@@ -25,11 +25,16 @@ const (
 	// The timeout can not be too long because the pod grace termination period is fixed.
 	writeTimeout = 10 * time.Second
 	category     = "proxy"
+
+	crossAZKey         = "crossZone_bytes"
+	publicEndpointKey  = "public_outBound_bytes"
+	privateEndpointKey = "private_outBound_bytes"
 )
 
 type MeterData struct {
-	respBytes    int64
-	crossAZBytes int64
+	publicRespBytes  int64
+	privateRespBytes int64
+	crossAZBytes     int64
 }
 
 type Meter struct {
@@ -62,11 +67,15 @@ func NewMeter(cfg *config.Config, lg *zap.Logger) (*Meter, error) {
 	}, nil
 }
 
-func (m *Meter) IncTraffic(clusterID string, respBytes, crossAZBytes int64) {
+func (m *Meter) IncTraffic(clusterID string, respBytes, crossAZBytes int64, fromPublicEndpoint bool) {
 	m.Lock()
 	defer m.Unlock()
 	orig := m.data[clusterID]
-	orig.respBytes += respBytes
+	if fromPublicEndpoint {
+		orig.publicRespBytes += respBytes
+	} else {
+		orig.privateRespBytes += respBytes
+	}
 	orig.crossAZBytes += crossAZBytes
 	m.data[clusterID] = orig
 }
@@ -109,11 +118,12 @@ func (m *Meter) flush(ts int64, timeout time.Duration) {
 	array := make([]map[string]any, 0, len(data))
 	for clusterID, d := range data {
 		array = append(array, map[string]any{
-			"version":         "1",
-			"cluster_id":      clusterID,
-			"source_name":     category,
-			"crossZone_bytes": &common.MeteringValue{Value: uint64(d.crossAZBytes), Unit: "bytes"},
-			"outBound_bytes":  &common.MeteringValue{Value: uint64(d.respBytes), Unit: "bytes"},
+			"version":          "1",
+			"cluster_id":       clusterID,
+			"source_name":      category,
+			crossAZKey:         &common.MeteringValue{Value: uint64(d.crossAZBytes), Unit: "bytes"},
+			privateEndpointKey: &common.MeteringValue{Value: uint64(d.privateRespBytes), Unit: "bytes"},
+			publicEndpointKey:  &common.MeteringValue{Value: uint64(d.publicRespBytes), Unit: "bytes"},
 		})
 	}
 
diff --git a/pkg/manager/meter/meter_test.go b/pkg/manager/meter/meter_test.go
@@ -72,19 +72,23 @@ func TestNewMeter(t *testing.T) {
 func TestWrite(t *testing.T) {
 	m, reader := createLocalMeter(t, t.TempDir())
 	ts := time.Now().Unix() / 60 * 60
-	m.IncTraffic("cluster-1", 100, 200)
-	m.IncTraffic("cluster-2", 200, 300)
+	m.IncTraffic("cluster-1", 100, 200, false)
+	m.IncTraffic("cluster-2", 200, 300, false)
+	m.IncTraffic("cluster-1", 500, 200, true)
+	m.IncTraffic("cluster-2", 600, 300, true)
 	m.flush(ts, time.Second)
 
 	data := readMeteringData(t, reader, ts)
 	require.Len(t, data, 2)
-	resp, crossAZ := getValuesFromData(t, data, "cluster-1")
-	require.Equal(t, int64(100), resp)
-	require.Equal(t, int64(200), crossAZ)
-
-	resp, crossAZ = getValuesFromData(t, data, "cluster-2")
-	require.Equal(t, int64(200), resp)
-	require.Equal(t, int64(300), crossAZ)
+	publicResp, privateResp, crossAZ := getValuesFromData(t, data, "cluster-1")
+	require.Equal(t, int64(500), publicResp)
+	require.Equal(t, int64(100), privateResp)
+	require.Equal(t, int64(400), crossAZ)
+
+	publicResp, privateResp, crossAZ = getValuesFromData(t, data, "cluster-2")
+	require.Equal(t, int64(600), publicResp)
+	require.Equal(t, int64(200), privateResp)
+	require.Equal(t, int64(600), crossAZ)
 }
 
 func TestLoop(t *testing.T) {
@@ -97,34 +101,40 @@ func TestLoop(t *testing.T) {
 	for range 10 {
 		wg.Run(func() {
 			for range 100 {
-				m.IncTraffic("cluster-1", 1, 2)
-				m.IncTraffic("cluster-2", 1, 2)
+				m.IncTraffic("cluster-1", 1, 2, false)
+				m.IncTraffic("cluster-2", 1, 2, false)
+				m.IncTraffic("cluster-1", 10, 2, true)
+				m.IncTraffic("cluster-2", 10, 2, true)
 				time.Sleep(time.Millisecond)
 			}
 		}, nil)
 	}
 	wg.Wait()
 	require.NoError(t, m.Close())
 
-	totalResp, totalCrossAZ := make(map[string]int64), make(map[string]int64)
+	totalPublicResp, totalPrivateResp, totalCrossAZ := make(map[string]int64), make(map[string]int64), make(map[string]int64)
 	for ts := startTime / 60 * 60; ts <= startTime/60*60+60; ts += 60 {
 		data := readMeteringData(t, reader, ts)
 		if len(data) == 0 {
 			continue
 		}
 		require.Len(t, data, 2)
-		resp, crossAZ := getValuesFromData(t, data, "cluster-1")
-		totalResp["cluster-1"] += resp
+		publicResp, privateResp, crossAZ := getValuesFromData(t, data, "cluster-1")
+		totalPublicResp["cluster-1"] += publicResp
+		totalPrivateResp["cluster-1"] += privateResp
 		totalCrossAZ["cluster-1"] += crossAZ
-		resp, crossAZ = getValuesFromData(t, data, "cluster-2")
-		totalResp["cluster-2"] += resp
+		publicResp, privateResp, crossAZ = getValuesFromData(t, data, "cluster-2")
+		totalPublicResp["cluster-2"] += publicResp
+		totalPrivateResp["cluster-2"] += privateResp
 		totalCrossAZ["cluster-2"] += crossAZ
 	}
 
-	require.Equal(t, int64(1000), totalResp["cluster-1"])
-	require.Equal(t, int64(2000), totalCrossAZ["cluster-1"])
-	require.Equal(t, int64(1000), totalResp["cluster-2"])
-	require.Equal(t, int64(2000), totalCrossAZ["cluster-2"])
+	require.Equal(t, int64(10000), totalPublicResp["cluster-1"])
+	require.Equal(t, int64(1000), totalPrivateResp["cluster-1"])
+	require.Equal(t, int64(4000), totalCrossAZ["cluster-1"])
+	require.Equal(t, int64(10000), totalPublicResp["cluster-2"])
+	require.Equal(t, int64(1000), totalPrivateResp["cluster-2"])
+	require.Equal(t, int64(4000), totalCrossAZ["cluster-2"])
 }
 
 func createLocalMeter(t *testing.T, dir string) (*Meter, *meteringreader.MeteringReader) {
@@ -183,17 +193,20 @@ func readMeteringData(t *testing.T, reader *meteringreader.MeteringReader, ts in
 	return meteringData.Data
 }
 
-func getValuesFromData(t *testing.T, data []map[string]any, clusterID string) (int64, int64) {
+func getValuesFromData(t *testing.T, data []map[string]any, clusterID string) (int64, int64, int64) {
 	for i := range data {
 		if data[i]["cluster_id"] == clusterID {
-			outBound, ok := data[i]["outBound_bytes"].(map[string]any)
+			publicOutBound, ok := data[i][publicEndpointKey].(map[string]any)
+			require.True(t, ok)
+			privateOutBound, ok := data[i][privateEndpointKey].(map[string]any)
 			require.True(t, ok)
-			crossZone, ok := data[i]["crossZone_bytes"].(map[string]any)
+			crossZone, ok := data[i][crossAZKey].(map[string]any)
 			require.True(t, ok)
-			require.Equal(t, "bytes", outBound["unit"])
+			require.Equal(t, "bytes", publicOutBound["unit"])
+			require.Equal(t, "bytes", privateOutBound["unit"])
 			require.Equal(t, "bytes", crossZone["unit"])
-			return int64(outBound["value"].(float64)), int64(crossZone["value"].(float64))
+			return int64(publicOutBound["value"].(float64)), int64(privateOutBound["value"].(float64)), int64(crossZone["value"].(float64))
 		}
 	}
-	return 0, 0
+	return 0, 0, 0
 }
diff --git a/pkg/proxy/backend/backend_conn_mgr.go b/pkg/proxy/backend/backend_conn_mgr.go
@@ -90,6 +90,7 @@ const (
 type BCConfig struct {
 	HealthyKeepAlive     config.KeepAlive
 	UnhealthyKeepAlive   config.KeepAlive
+	FromPublicEndpoints  func(addr net.Addr) bool
 	TickerInterval       time.Duration
 	CheckBackendInterval time.Duration
 	DialTimeout          time.Duration
@@ -115,7 +116,7 @@ func (cfg *BCConfig) check() {
 }
 
 type Meter interface {
-	IncTraffic(clusterID string, respBytes, crossAZBytes int64)
+	IncTraffic(clusterID string, respBytes, crossAZBytes int64, fromPublicEndpoint bool)
 }
 
 // BackendConnManager migrates a session from one BackendConnection to another.
@@ -162,10 +163,11 @@ type BackendConnManager struct {
 		sync.Mutex
 		m map[any]any
 	}
-	connectionID uint64
-	quitSource   ErrorSource
-	cpt          capture.Capture
-	meter        Meter
+	connectionID       uint64
+	quitSource         ErrorSource
+	cpt                capture.Capture
+	meter              Meter
+	fromPublicEndpoint bool
 }
 
 // NewBackendConnManager creates a BackendConnManager.
@@ -235,6 +237,9 @@ func (mgr *BackendConnManager) Connect(ctx context.Context, clientIO pnet.Packet
 	mgr.handshakeHandler.OnHandshake(mgr, mgr.ServerAddr(), nil, SrcNone)
 	endTime := time.Now()
 	addHandshakeMetrics(mgr.ServerAddr(), endTime.Sub(startTime))
+	if mgr.config.FromPublicEndpoints != nil {
+		mgr.fromPublicEndpoint = mgr.config.FromPublicEndpoints(clientIO.ProxyAddr())
+	}
 	mgr.updateTraffic(*mgr.backendIO.Load())
 
 	mgr.cmdProcessor.capability = mgr.authenticator.capability
@@ -451,7 +456,7 @@ func (mgr *BackendConnManager) updateTraffic(backendIO pnet.PacketIO) {
 			if !mgr.curBackend.Local() {
 				crossAZBytes = int64(outBytes - mgr.outBytes + inBytes - mgr.inBytes)
 			}
-			mgr.meter.IncTraffic(keyspace, int64(inBytes-mgr.inBytes), crossAZBytes)
+			mgr.meter.IncTraffic(keyspace, int64(inBytes-mgr.inBytes), crossAZBytes, mgr.fromPublicEndpoint)
 		}
 	}
 	mgr.inBytes, mgr.inPackets, mgr.outBytes, mgr.outPackets = inBytes, inPackets, outBytes, outPackets
diff --git a/pkg/proxy/backend/backend_conn_mgr_test.go b/pkg/proxy/backend/backend_conn_mgr_test.go
@@ -1369,6 +1369,9 @@ func TestTrafficMetrics(t *testing.T) {
 	meter := newMeter()
 	ts := newBackendMgrTester(t, func(cfg *testConfig) {
 		cfg.proxyConfig.meter = meter
+		cfg.proxyConfig.bcConfig.FromPublicEndpoints = func(net.Addr) bool {
+			return true
+		}
 	})
 	var inBytes, inPackets, outBytes, outPackets int
 	runners := []runner{
@@ -1398,7 +1401,8 @@ func TestTrafficMetrics(t *testing.T) {
 				require.NoError(t, err)
 				require.True(t, inBytes2 > inBytes && inPackets2 > inPackets && outBytes2 > outBytes && outPackets2 > outPackets)
 				require.True(t, inBytes2 > 4096 && inPackets2 > 1000)
-				require.True(t, meter.respBytes["key1"] > 4096)
+				require.True(t, meter.publicRespBytes["key1"] > 4096)
+				require.EqualValues(t, 0, meter.privateRespBytes["key1"])
 				inBytes, inPackets, outBytes, outPackets = inBytes2, inPackets2, outBytes2, outPackets2
 				// The first backend is local, so no cross-az traffic.
 				crossLocationBytes2, err := metrics.ReadCounter(metrics.CrossLocationBytesCounter)
@@ -1444,7 +1448,8 @@ func TestTrafficMetrics(t *testing.T) {
 				inBytes2, inPackets2, outBytes2, outPackets2, err := readTraffic(addr)
 				require.NoError(t, err)
 				require.True(t, inBytes2 > inBytes1 && inPackets2 > inPackets1 && outBytes2 > outBytes1 && outPackets2 > outPackets1)
-				require.True(t, meter.respBytes["key2"] > 0)
+				require.True(t, meter.publicRespBytes["key2"] > 0)
+				require.EqualValues(t, 0, meter.privateRespBytes["key2"])
 				// The second backend is remote, so exists cross-az traffic.
 				crossLocationBytes2, err := metrics.ReadCounter(metrics.CrossLocationBytesCounter)
 				require.NoError(t, err)
diff --git a/pkg/proxy/backend/mock_proxy_test.go b/pkg/proxy/backend/mock_proxy_test.go
@@ -160,18 +160,24 @@ func (mc *mockCapture) Close() {
 var _ Meter = (*mockMeter)(nil)
 
 type mockMeter struct {
-	crossAZBytes map[string]int64
-	respBytes    map[string]int64
+	crossAZBytes     map[string]int64
+	privateRespBytes map[string]int64
+	publicRespBytes  map[string]int64
 }
 
 func newMeter() *mockMeter {
 	return &mockMeter{
-		crossAZBytes: make(map[string]int64),
-		respBytes:    make(map[string]int64),
+		crossAZBytes:     make(map[string]int64),
+		publicRespBytes:  make(map[string]int64),
+		privateRespBytes: make(map[string]int64),
 	}
 }
 
-func (m *mockMeter) IncTraffic(clusterID string, respBytes, crossAZBytes int64) {
+func (m *mockMeter) IncTraffic(clusterID string, respBytes, crossAZBytes int64, fromPublicEndpoint bool) {
 	m.crossAZBytes[clusterID] += crossAZBytes
-	m.respBytes[clusterID] += respBytes
+	if fromPublicEndpoint {
+		m.publicRespBytes[clusterID] += respBytes
+	} else {
+		m.privateRespBytes[clusterID] += respBytes
+	}
 }
diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go
@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/tiproxy/pkg/proxy/keepalive"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/capture"
+	"github.com/pingcap/tiproxy/pkg/util/netutil"
 	"github.com/pingcap/tiproxy/pkg/util/waitgroup"
 	"go.uber.org/zap"
 )
@@ -29,6 +30,7 @@ type serverState struct {
 	healthyKeepAlive   config.KeepAlive
 	unhealthyKeepAlive config.KeepAlive
 	clients            map[uint64]*client.ClientConnection
+	publicEndpoints    []*net.IPNet
 	maxConnections     uint64
 	connBufferSize     int
 	requireBackendTLS  bool
@@ -84,6 +86,10 @@ func NewSQLServer(logger *zap.Logger, cfg *config.Config, certMgr *cert.CertMana
 }
 
 func (s *SQLServer) reset(cfg *config.Config) {
+	cidrList, parseErr := netutil.ParseCIDRList(cfg.Proxy.PublicEndpoints)
+	if parseErr != nil {
+		s.logger.Warn("failed to parse public endpoints", zap.Error(parseErr))
+	}
 	s.mu.Lock()
 	s.mu.tcpKeepAlive = cfg.Proxy.FrontendKeepalive.Enabled
 	s.mu.maxConnections = cfg.Proxy.MaxConnections
@@ -94,6 +100,7 @@ func (s *SQLServer) reset(cfg *config.Config) {
 	s.mu.healthyKeepAlive = cfg.Proxy.BackendHealthyKeepalive
 	s.mu.unhealthyKeepAlive = cfg.Proxy.BackendUnhealthyKeepalive
 	s.mu.connBufferSize = cfg.Proxy.ConnBufferSize
+	s.mu.publicEndpoints = cidrList
 	s.mu.Unlock()
 }
 
diff --git a/pkg/util/netutil/netutil.go b/pkg/util/netutil/netutil.go
diff --git a/pkg/util/netutil/netutil_test.go b/pkg/util/netutil/netutil_test.go

Original file line number	Diff line number	Diff line change
`@@ -61,6 +61,8 @@ type ProxyServerOnline struct {`
`61`	`61`	`// In k8s, the pod terminationGracePeriodSeconds can be set to very long so that these configs can be updated online.`
`62`	`62`	GracefulWaitBeforeShutdown int `yaml:"graceful-wait-before-shutdown,omitempty" toml:"graceful-wait-before-shutdown,omitempty" json:"graceful-wait-before-shutdown,omitempty" reloadable:"true"`
`63`	`63`	GracefulCloseConnTimeout int `yaml:"graceful-close-conn-timeout,omitempty" toml:"graceful-close-conn-timeout,omitempty" json:"graceful-close-conn-timeout,omitempty" reloadable:"true"`
	`64`	`+ // Public and private traffic are metered separately.`
	`65`	+ PublicEndpoints []string `yaml:"public-endpoints,omitempty" toml:"public-endpoints,omitempty" json:"public-endpoints,omitempty" reloadable:"true"`
`64`	`66`	`}`
`65`	`67`
`66`	`68`	`type ProxyServer struct {`
Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@ func TestParseCIDR(t *testing.T) {`
`26`	`26`	`}{`
`27`	`27`	`{`
`28`	`28`	`cidrs: []string{"1.1.1.1"},`
`29`		`- success: false,`
	`29`	`+ success: true,`
`30`	`30`	`},`
`31`	`31`	`{`
`32`	`32`	`cidrs: []string{"1.1.1.1/32"},`