snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/xds/clients/lrsclient/load_store.go

@@ -0,0 +1,438 @@
+/*
+ *
+ * Copyright 2025 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package lrsclient
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/internal/xds/clients"
+	lrsclientinternal "google.golang.org/grpc/internal/xds/clients/lrsclient/internal"
+)
+
+// A LoadStore aggregates loads for multiple clusters and services that are
+// intended to be reported via LRS.
+//
+// LoadStore stores loads reported to a single LRS server. Use multiple stores
+// for multiple servers.
+//
+// It is safe for concurrent use.
+type LoadStore struct {
+	// stop is the function to call to Stop the LoadStore reporting.
+	stop func(ctx context.Context)
+
+	// mu only protects the map (2 layers). The read/write to
+	// *PerClusterReporter doesn't need to hold the mu.
+	mu sync.Mutex
+	// clusters is a map with cluster name as the key. The second layer is a
+	// map with service name as the key. Each value (PerClusterReporter)
+	// contains data for a (cluster, service) pair.
+	//
+	// Note that new entries are added to this map, but never removed. This is
+	// potentially a memory leak. But the memory is allocated for each new
+	// (cluster,service) pair, and the memory allocated is just pointers and
+	// maps. So this shouldn't get too bad.
+	clusters map[string]map[string]*PerClusterReporter
+}
+
+func init() {
+	lrsclientinternal.TimeNow = time.Now
+}
+
+// newLoadStore creates a LoadStore.
+func newLoadStore() *LoadStore {
+	return &LoadStore{
+		clusters: make(map[string]map[string]*PerClusterReporter),
+	}
+}
+
+// Stop signals the LoadStore to stop reporting.
+//
+// Before closing the underlying LRS stream, this method may block until a
+// final load report send attempt completes or the provided context `ctx`
+// expires.
+//
+// The provided context must have a deadline or timeout set to prevent Stop
+// from blocking indefinitely if the final send attempt fails to complete.
+//
+// Calling Stop on an already stopped LoadStore is a no-op.
+func (ls *LoadStore) Stop(ctx context.Context) {
+	ls.stop(ctx)
+}
+
+// ReporterForCluster returns the PerClusterReporter for the given cluster and
+// service.
+func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) *PerClusterReporter {
+	ls.mu.Lock()
+	defer ls.mu.Unlock()
+	c, ok := ls.clusters[clusterName]
+	if !ok {
+		c = make(map[string]*PerClusterReporter)
+		ls.clusters[clusterName] = c
+	}
+
+	if p, ok := c[serviceName]; ok {
+		return p
+	}
+	p := &PerClusterReporter{
+		cluster:          clusterName,
+		service:          serviceName,
+		lastLoadReportAt: lrsclientinternal.TimeNow(),
+	}
+	c[serviceName] = p
+	return p
+}
+
+// stats returns the load data for the given cluster names. Data is returned in
+// a slice with no specific order.
+//
+// If no clusterName is given (an empty slice), all data for all known clusters
+// is returned.
+//
+// If a cluster's loadData is empty (no load to report), it's not appended to
+// the returned slice.
+func (ls *LoadStore) stats(clusterNames []string) []*loadData {
+	ls.mu.Lock()
+	defer ls.mu.Unlock()
+
+	var ret []*loadData
+	if len(clusterNames) == 0 {
+		for _, c := range ls.clusters {
+			ret = appendClusterStats(ret, c)
+		}
+		return ret
+	}
+	for _, n := range clusterNames {
+		if c, ok := ls.clusters[n]; ok {
+			ret = appendClusterStats(ret, c)
+		}
+	}
+
+	return ret
+}
+
+// PerClusterReporter records load data pertaining to a single cluster. It
+// provides methods to record call starts, finishes, server-reported loads,
+// and dropped calls.
+//
+// It is safe for concurrent use.
+//
+// TODO(purnesh42h): Use regular maps with mutexes instead of sync.Map here.
+// The latter is optimized for two common use cases: (1) when the entry for a
+// given key is only ever written once but read many times, as in caches that
+// only grow, or (2) when multiple goroutines read, write, and overwrite
+// entries for disjoint sets of keys. In these two cases, use of a Map may
+// significantly reduce lock contention compared to a Go map paired with a
+// separate Mutex or RWMutex.
+// Neither of these conditions are met here, and we should transition to a
+// regular map with a mutex for better type safety.
+type PerClusterReporter struct {
+	cluster, service string
+	drops            sync.Map // map[string]*uint64
+	localityRPCCount sync.Map // map[clients.Locality]*rpcCountData
+
+	mu               sync.Mutex
+	lastLoadReportAt time.Time
+}
+
+// CallStarted records a call started in the LoadStore.
+func (p *PerClusterReporter) CallStarted(locality clients.Locality) {
+	s, ok := p.localityRPCCount.Load(locality)
+	if !ok {
+		tp := newRPCCountData()
+		s, _ = p.localityRPCCount.LoadOrStore(locality, tp)
+	}
+	s.(*rpcCountData).incrInProgress()
+	s.(*rpcCountData).incrIssued()
+}
+
+// CallFinished records a call finished in the LoadStore.
+func (p *PerClusterReporter) CallFinished(locality clients.Locality, err error) {
+	f, ok := p.localityRPCCount.Load(locality)
+	if !ok {
+		// The map is never cleared, only values in the map are reset. So the
+		// case where entry for call-finish is not found should never happen.
+		return
+	}
+	f.(*rpcCountData).decrInProgress()
+	if err == nil {
+		f.(*rpcCountData).incrSucceeded()
+	} else {
+		f.(*rpcCountData).incrErrored()
+	}
+}
+
+// CallServerLoad records the server load in the LoadStore.
+func (p *PerClusterReporter) CallServerLoad(locality clients.Locality, name string, val float64) {
+	s, ok := p.localityRPCCount.Load(locality)
+	if !ok {
+		// The map is never cleared, only values in the map are reset. So the
+		// case where entry for callServerLoad is not found should never happen.
+		return
+	}
+	s.(*rpcCountData).addServerLoad(name, val)
+}
+
+// CallDropped records a call dropped in the LoadStore.
+func (p *PerClusterReporter) CallDropped(category string) {
+	d, ok := p.drops.Load(category)
+	if !ok {
+		tp := new(uint64)
+		d, _ = p.drops.LoadOrStore(category, tp)
+	}
+	atomic.AddUint64(d.(*uint64), 1)
+}
+
+// stats returns and resets all loads reported to the store, except inProgress
+// rpc counts.
+//
+// It returns nil if the store doesn't contain any (new) data.
+func (p *PerClusterReporter) stats() *loadData {
+	sd := newLoadData(p.cluster, p.service)
+	p.drops.Range(func(key, val any) bool {
+		d := atomic.SwapUint64(val.(*uint64), 0)
+		if d == 0 {
+			return true
+		}
+		sd.totalDrops += d
+		keyStr := key.(string)
+		if keyStr != "" {
+			// Skip drops without category. They are counted in total_drops, but
+			// not in per category. One example is drops by circuit breaking.
+			sd.drops[keyStr] = d
+		}
+		return true
+	})
+	p.localityRPCCount.Range(func(key, val any) bool {
+		countData := val.(*rpcCountData)
+		succeeded := countData.loadAndClearSucceeded()
+		inProgress := countData.loadInProgress()
+		errored := countData.loadAndClearErrored()
+		issued := countData.loadAndClearIssued()
+		if succeeded == 0 && inProgress == 0 && errored == 0 && issued == 0 {
+			return true
+		}
+
+		ld := localityData{
+			requestStats: requestData{
+				succeeded:  succeeded,
+				errored:    errored,
+				inProgress: inProgress,
+				issued:     issued,
+			},
+			loadStats: make(map[string]serverLoadData),
+		}
+		countData.serverLoads.Range(func(key, val any) bool {
+			sum, count := val.(*rpcLoadData).loadAndClear()
+			if count == 0 {
+				return true
+			}
+			ld.loadStats[key.(string)] = serverLoadData{
+				count: count,
+				sum:   sum,
+			}
+			return true
+		})
+		sd.localityStats[key.(clients.Locality)] = ld
+		return true
+	})
+
+	p.mu.Lock()
+	sd.reportInterval = lrsclientinternal.TimeNow().Sub(p.lastLoadReportAt)
+	p.lastLoadReportAt = lrsclientinternal.TimeNow()
+	p.mu.Unlock()
+
+	if sd.totalDrops == 0 && len(sd.drops) == 0 && len(sd.localityStats) == 0 {
+		return nil
+	}
+	return sd
+}
+
+// loadData contains all load data reported to the LoadStore since the most recent
+// call to stats().
+type loadData struct {
+	// cluster is the name of the cluster this data is for.
+	cluster string
+	// service is the name of the EDS service this data is for.
+	service string
+	// totalDrops is the total number of dropped requests.
+	totalDrops uint64
+	// drops is the number of dropped requests per category.
+	drops map[string]uint64
+	// localityStats contains load reports per locality.
+	localityStats map[clients.Locality]localityData
+	// reportInternal is the duration since last time load was reported (stats()
+	// was called).
+	reportInterval time.Duration
+}
+
+// localityData contains load data for a single locality.
+type localityData struct {
+	// requestStats contains counts of requests made to the locality.
+	requestStats requestData
+	// loadStats contains server load data for requests made to the locality,
+	// indexed by the load type.
+	loadStats map[string]serverLoadData
+}
+
+// requestData contains request counts.
+type requestData struct {
+	// succeeded is the number of succeeded requests.
+	succeeded uint64
+	// errored is the number of requests which ran into errors.
+	errored uint64
+	// inProgress is the number of requests in flight.
+	inProgress uint64
+	// issued is the total number requests that were sent.
+	issued uint64
+}
+
+// serverLoadData contains server load data.
+type serverLoadData struct {
+	// count is the number of load reports.
+	count uint64
+	// sum is the total value of all load reports.
+	sum float64
+}
+
+// appendClusterStats gets the Data for all the given clusters, append to ret,
+// and return the new slice.
+//
+// Data is only appended to ret if it's not empty.
+func appendClusterStats(ret []*loadData, clusters map[string]*PerClusterReporter) []*loadData {
+	for _, d := range clusters {
+		data := d.stats()
+		if data == nil {
+			// Skip this data if it doesn't contain any information.
+			continue
+		}
+		ret = append(ret, data)
+	}
+	return ret
+}
+
+func newLoadData(cluster, service string) *loadData {
+	return &loadData{
+		cluster:       cluster,
+		service:       service,
+		drops:         make(map[string]uint64),
+		localityStats: make(map[clients.Locality]localityData),
+	}
+}
+
+type rpcCountData struct {
+	// Only atomic accesses are allowed for the fields.
+	succeeded  *uint64
+	errored    *uint64
+	inProgress *uint64
+	issued     *uint64
+
+	// Map from load desc to load data (sum+count). Loading data from map is
+	// atomic, but updating data takes a lock, which could cause contention when
+	// multiple RPCs try to report loads for the same desc.
+	//
+	// To fix the contention, shard this map.
+	serverLoads sync.Map // map[string]*rpcLoadData
+}
+
+func newRPCCountData() *rpcCountData {
+	return &rpcCountData{
+		succeeded:  new(uint64),
+		errored:    new(uint64),
+		inProgress: new(uint64),
+		issued:     new(uint64),
+	}
+}
+
+func (rcd *rpcCountData) incrSucceeded() {
+	atomic.AddUint64(rcd.succeeded, 1)
+}
+
+func (rcd *rpcCountData) loadAndClearSucceeded() uint64 {
+	return atomic.SwapUint64(rcd.succeeded, 0)
+}
+
+func (rcd *rpcCountData) incrErrored() {
+	atomic.AddUint64(rcd.errored, 1)
+}
+
+func (rcd *rpcCountData) loadAndClearErrored() uint64 {
+	return atomic.SwapUint64(rcd.errored, 0)
+}
+
+func (rcd *rpcCountData) incrInProgress() {
+	atomic.AddUint64(rcd.inProgress, 1)
+}
+
+func (rcd *rpcCountData) decrInProgress() {
+	atomic.AddUint64(rcd.inProgress, ^uint64(0)) // atomic.Add(x, -1)
+}
+
+func (rcd *rpcCountData) loadInProgress() uint64 {
+	return atomic.LoadUint64(rcd.inProgress) // InProgress count is not clear when reading.
+}
+
+func (rcd *rpcCountData) incrIssued() {
+	atomic.AddUint64(rcd.issued, 1)
+}
+
+func (rcd *rpcCountData) loadAndClearIssued() uint64 {
+	return atomic.SwapUint64(rcd.issued, 0)
+}
+
+func (rcd *rpcCountData) addServerLoad(name string, d float64) {
+	loads, ok := rcd.serverLoads.Load(name)
+	if !ok {
+		tl := newRPCLoadData()
+		loads, _ = rcd.serverLoads.LoadOrStore(name, tl)
+	}
+	loads.(*rpcLoadData).add(d)
+}
+
+// rpcLoadData is data for server loads (from trailers or oob). Fields in this
+// struct must be updated consistently.
+//
+// The current solution is to hold a lock, which could cause contention. To fix,
+// shard serverLoads map in rpcCountData.
+type rpcLoadData struct {
+	mu    sync.Mutex
+	sum   float64
+	count uint64
+}
+
+func newRPCLoadData() *rpcLoadData {
+	return &rpcLoadData{}
+}
+
+func (rld *rpcLoadData) add(v float64) {
+	rld.mu.Lock()
+	rld.sum += v
+	rld.count++
+	rld.mu.Unlock()
+}
+
+func (rld *rpcLoadData) loadAndClear() (s float64, c uint64) {
+	rld.mu.Lock()
+	s, rld.sum = rld.sum, 0
+	c, rld.count = rld.count, 0
+	rld.mu.Unlock()
+	return s, c
+}