xds: generic lrs client for load reporting #8250

purnesh42H · 2025-04-15T06:45:57Z

This is the change to make generic LRS client for load reporting to LRS server.

The PR copies the existing

xds/internal/xdsclient/load/store.go,
xds/internal/xdsclient/transport/lrs/lrs_stream.go,
xds/internal/xdsclient/load/store_test.go
xds/internal/xdsclient/tests/loadreport_test.go

from internal xdsclient code and then modify them to use the generic client types and interfaces. Each "copy" commit is followed by the "modify" commit for that file. Reviewers can start from reviewing the "modify" commit.

PS: Currently loadreport_test.go has compilation error as so its commented out as it is depends on some of the functions added in #8183

RELEASE NOTES: None

codecov · 2025-04-15T06:55:19Z

Codecov Report

Attention: Patch coverage is 78.68132% with 97 lines in your changes missing coverage. Please review.

Project coverage is 82.29%. Comparing base (82e25c7) to head (a263158).
Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
xds/internal/clients/lrsclient/lrs_stream.go	66.85%	43 Missing and 15 partials ⚠️
xds/internal/clients/lrsclient/lrsclient.go	74.15%	17 Missing and 6 partials ⚠️
xds/internal/clients/lrsclient/load_store.go	91.44%	12 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #8250      +/-   ##
==========================================
+ Coverage   82.15%   82.29%   +0.14%     
==========================================
  Files         412      419       +7     
  Lines       40562    41787    +1225     
==========================================
+ Hits        33322    34390    +1068     
- Misses       5875     5949      +74     
- Partials     1365     1448      +83

Files with missing lines	Coverage Δ
xds/internal/clients/lrsclient/logging.go	`100.00% <100.00%> (ø)`
xds/internal/clients/lrsclient/load_store.go	`91.62% <91.44%> (+91.62%)`	⬆️
xds/internal/clients/lrsclient/lrsclient.go	`74.15% <74.15%> (+74.15%)`	⬆️
xds/internal/clients/lrsclient/lrs_stream.go	`66.85% <66.85%> (ø)`

... and 30 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

purnesh42H · 2025-04-21T05:46:18Z

@dfawley assigning this for review since #8183 is close now. loadreport_test.go is commented due to testing helpers introduced in 8183 but i have tested in my fork. Will add Easwar once he finishes 8183.

dfawley

Looks fine overall. Just a few comments inline.

dfawley · 2025-04-21T20:03:22Z

xds/internal/clients/lrsclient/load_store.go

-func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) PerClusterReporter {
-	panic("unimplemented")
+func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) *PerClusterReporter {
+	if ls == nil {


I think it's fine to panic if a nil LoadStore is used. Why not? It seems like a pretty severe programming error.

Removed nil check. @easwars any reason why this check is there in existing code?

Probably for tests to use a nil load store. If that is not required anymore and tests are happy, we should be good to remove the nil check.

dfawley · 2025-04-21T20:06:00Z

xds/internal/clients/lrsclient/load_store.go

 }

 // CallStarted records a call started in the LoadStore.
 func (p *PerClusterReporter) CallStarted(locality string) {
-	panic("unimplemented")
+	if p == nil {


As above. And below.

dfawley · 2025-04-21T20:08:39Z

xds/internal/clients/lrsclient/load_store.go

+}
+
+func (rcd *rpcCountData) decrInProgress() {
+	atomic.AddUint64(rcd.inProgress, negativeOneUInt64) // atomic.Add(x, -1)


Nit: the const doesn't seem to buy us anything, since we're already needing to comment what this means. IMO delete the constant and inline it.

Made it inline

dfawley · 2025-04-21T20:10:34Z

xds/internal/clients/lrsclient/load_store.go

+	s = rld.sum
+	rld.sum = 0
+	c = rld.count
+	rld.count = 0


You could do something like this, which might(?) be more quickly understood:

Suggested change

s = rld.sum

rld.sum = 0

c = rld.count

rld.count = 0

s, rld.sum = rld.sum, 0

c, rld.count = rld.count, 0

Or,

Suggested change

s = rld.sum

rld.sum = 0

c = rld.count

rld.count = 0

s, c = rld.sum, rld.count

rld.sum, rld.count = 0, 0

Did the first one

dfawley · 2025-04-21T20:11:05Z

xds/internal/clients/lrsclient/load_store.go

+	c = rld.count
+	rld.count = 0
+	rld.mu.Unlock()
+	return


Please don't use bare returns that return values. It can be hard to understand what's going on, mainly for longer functions.

Suggested change

return

return s, c

Or pair with the second option above:

func (rld *rpcLoadData) loadAndClear() (float64, int64) { rld.mu.Lock() defer rld.mu.Unlock() s, c := rld.sum, rld.count rld.sum, rld.count = 0, 0 return s, c }

yeah added them to return

dfawley · 2025-04-21T20:39:43Z

xds/internal/clients/lrsclient/loadreport_test.go

+	return c, err
+}
+
+/*


Why is this commented out?

I have mentioned in the PR description. It is because the tests had compilation errors because some things were added in 8183 PR. Now its merged so i have rebased on latest and uncommented it.

dfawley · 2025-04-21T20:41:58Z

I mostly skimmed the changes - @easwars may also want to take a quick pass.

The commits here aren't quite as easy to review as the last change, since they go file-by-file. It would have been easier if one commit copied all the files, so that we could just skip that one commit when reviewing.

easwars · 2025-04-23T23:56:35Z

xds/internal/clients/lrsclient/load_store.go

 	// 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.


Could you please add this to the list of things that are being tracked? Thanks.

yes, i had already added. Initially, I was thinking to first change in internal code but i think its fine to fix here after migration.

easwars · 2025-04-23T23:58:25Z

xds/internal/clients/lrsclient/load_store.go

+	clusters map[string]map[string]*PerClusterReporter
+}
+
+// newStore creates a LoadStore.


This comments needs to be updated to match the function name.

easwars · 2025-04-23T23:59:42Z

xds/internal/clients/lrsclient/load_store.go

+	// Wait for the provided context to be done (timeout or cancellation).
+	if ctx != nil {
+		<-ctx.Done()
+	}


How and why would a user use this feature?

The idea is that if its last reference to underlying stream, stop should wait for context to be done which will allow any pending load to be flushed during the wait. If not, then just go ahead and close the straem.

I updated the logic to wait only if its the last reference just before closing the straem.

easwars · 2025-04-24T00:01:42Z

xds/internal/clients/lrsclient/load_store.go

-func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) PerClusterReporter {
-	panic("unimplemented")
+func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) *PerClusterReporter {
+	if ls == nil {


Probably for tests to use a nil load store. If that is not required anymore and tests are happy, we should be good to remove the nil check.

easwars · 2025-04-24T00:03:56Z

xds/internal/clients/lrsclient/load_store.go

+// 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 {
+	var ret []*loadData


Nit: Can we please move this variable declaration to be on the line above the if len(clusterNames) == 0 { ... }. The code just reads better when it starts with a call to lock and a deferred call to unlock and everything else follows after that.

easwars · 2025-04-24T00:07:43Z

xds/internal/clients/lrsclient/load_store.go

+// appendClusterStats gets Data for the given cluster, append to ret, and return
+// the new slice.


This comment needs some updating since there is no single "given cluster", but a collection of clusters.

easwars · 2025-04-24T00:11:33Z

xds/internal/clients/lrsclient/load_store_test.go

+	wg.Wait()
+
+	gotStoreData := ls.stats()
+	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval")); diff != "" {


This diff command is repeated multiple times with the same (or maybe almost same) set of options. Maybe a helper function can be written to do the comparison and the test can simply call if err := verifyStoreData(got, want); err != nil { ... }?

Moved to helper

easwars · 2025-04-24T00:16:22Z

xds/internal/clients/lrsclient/lrs_stream.go

+	}
+	var resp v3lrspb.LoadStatsResponse
+	if err := proto.Unmarshal(r, &resp); err != nil {
+		lrs.logger.Infof("Failed to unmarshal response to LoadStatsResponse: %v", err)


There is a warning log for the marshal error, but an info log for the unmarshal error. Please use a consistent log level, and if you end up choosing info, please add a verbosity check too.

Made both Infof guarded by verbosity check

easwars · 2025-04-24T00:22:49Z

xds/internal/clients/lrsclient/lrsclient.go

+	// The LRSClient owns a bunch of streams to individual LRS servers.
+	//
+	// Once all references to a stream are dropped, the stream is closed.


Once all references to a channel are dropped

What channel are we referring to here?

the streams which are running for each server. are you suggesting to mention them as channel?

easwars · 2025-04-24T00:27:19Z

xds/internal/clients/lrsclient/lrsclient.go

+		if lrs.cancelStream == nil {
+			// It is possible that Stop() is called before the cleanup function
+			// is called, thereby setting cancelStream to nil. Hence we need a
+			// nil check here bofore invoking the cancel function.
+			return
+		}


I don't quite follow this. The only place that sets lrs.cancelStream to nil is right below in this same closure.

I removed it. May be it had a reason in internal code because it was within xds client? Here its not possible

easwars · 2025-04-24T13:50:03Z

xds/internal/clients/lrsclient/lrsclient.go

+func New(config Config) (*LRSClient, error) {
+	switch {
+	case config.Node.ID == "":
+		return nil, errors.New("lrsclient: node ID is empty")


Nit: the error message says "node ID", but this is the complete node configuration. So, please update the error message to reflect that.

the case is checking for node ID. We shouldn't have node without ID. Updated to mention that.

easwars · 2025-04-24T13:57:45Z

xds/internal/clients/lrsclient/load_store.go

@@ -30,6 +33,28 @@ import "context"
 //
 // It is safe for concurrent use.
 type LoadStore struct {
+	lrsStream *streamImpl


The only reason we need the streamImpl here is to call stop on it. Why can't we simply have a function pointer here instead?

yeah i moved the function to LoadStore instead and setting to load store directly from lrs client. Also, renamed to Stop

easwars · 2025-04-24T13:58:55Z

xds/internal/clients/lrsclient/lrsclient.go

+	ctx, cancel := context.WithCancel(context.Background())
+	lrs.cancelStream = cancel
+	lrs.doneCh = make(chan struct{})
+	lrs.loadStore = newLoadStore(lrs)
+	go lrs.runner(ctx)


Can all of this be done in newStreamImpl?

easwars · 2025-04-24T14:03:18Z

xds/internal/clients/lrsclient/lrsclient.go

+		delete(c.lrsStreams, serverIdentifier)
+		tr.Close()
+	}
+	lrs.cleanup = cleanup


This is how I see the API being used:

user creates a new LRS client with lrsclient.New()

when they want to report load to a server, they would call ReportLoad on the returned LRS client from above

this would create a new streamImpl if required or reuse an existing one

the user is returned a reference to the associated LoadStore

going forward, all user interactions are with the LoadStore and the PerClusterReporter returned from the store

Eventually, they would call Stop() on the store to indicate that they no longer wish to report loads

this calls streamImpl.stop(), which calls this cleanup function

we need a once func somewhere in this call path to ensure that multiple calls to LoadStore.Stop() does not decrement the ref count multiple times

yeah that's the flow.

we need a once func somewhere in this call path to ensure that multiple calls to LoadStore.Stop() does not decrement the ref count multiple times

once func is not possible because LoadStore is shared? The stop function though has a check if lrs reference is already 0 and stop is called, it returns early but log an error.

easwars · 2025-04-24T14:05:02Z

xds/internal/clients/lrsclient/load_store.go

@@ -43,38 +68,372 @@ type LoadStore struct {
 // attempt to flush any unreported load data to the LRS server. It will either
 // wait for this attempt to complete, or for the provided context to be done
 // before canceling the LRS stream.
-func (ls *LoadStore) Stop(ctx context.Context) error {
-	panic("unimplemented")
+func (ls *LoadStore) Stop(ctx context.Context) {


I think we can remove this paragraph about reference counting from this docstring.

Also, where do we make the last attempt to flush any unsent loads?

i modified it to just callout the last reference to stream and if this is the last reference to stream, only then wait for context to be done.

Also, where do we make the last attempt to flush any unsent loads?

I mentioned in the other comment. Basically, we wait for context to complete to allow some time for load to be flushed based on load reporting interval because i think we don't want to make multiple attempts in less than reporting interval? wdyt?

easwars · 2025-04-24T14:05:53Z

xds/internal/clients/lrsclient/lrsclient.go

-func (*LRSClient) ReportLoad(_ clients.ServerIdentifier) *LoadStore {
-	panic("unimplemented")
+// ReportLoad creates and returns a LoadStore for the caller to report loads
+// using a LoadReportingStream.


I think we should mention here that the caller should call Stop on the returned LoadStore when they are done reporting load to this server.

purnesh42H added Type: Feature New features or improvements in behavior Area: xDS Includes everything xDS related, including LB policies used with xDS. labels Apr 15, 2025

purnesh42H added this to the 1.73 Release milestone Apr 15, 2025

purnesh42H force-pushed the generic-xds-client-lrs-client-e2e branch from 78ab34a to ce9ba3d Compare April 19, 2025 18:36

purnesh42H requested a review from dfawley April 21, 2025 05:39

purnesh42H assigned dfawley Apr 21, 2025

purnesh42H force-pushed the generic-xds-client-lrs-client-e2e branch 2 times, most recently from e140792 to 2e2674f Compare April 21, 2025 18:22

dfawley approved these changes Apr 21, 2025

View reviewed changes

dfawley assigned purnesh42H and unassigned dfawley Apr 21, 2025

easwars self-assigned this Apr 22, 2025

purnesh42H added 15 commits April 23, 2025 21:41

copy store.go from internal xdsclient/load

3c7a0de

delete load_store.go

32f4707

rename store.go to load_store.go

0f4358a

modify load_store.go for generic lrs client

997cebd

copy store_test.go from internal xdsclient/load

d85d63e

rename store_test.go to load_store_test.go

b69c7b4

modify load_store_test.go

f5afaa5

make vet happy

6e95f47

copy lrs_stream.go from internal xdsclient/transport/lrs

05ce351

modify lrs_stream.go and load_store.go for lrs stream reporting

3cbb974

copy loadreport_test.go from internal xdsclient/test

2f056b1

modify loadreport_test.go for generic lrs client

b048ef8

modify lrsclient.go

8afee6a

comment loadreport_test.go until 8183 is done

6ae274b

make vet happy

871aa44

purnesh42H added 3 commits April 23, 2025 21:41

explicitly unlock lrs stream mu because cleanup has client level lock

9b0c28f

keep refcounting in LRS client

6ababcd

dfawley review 1

aaa667d

purnesh42H force-pushed the generic-xds-client-lrs-client-e2e branch from 0a8d352 to aaa667d Compare April 23, 2025 16:22

purnesh42H removed their assignment Apr 23, 2025

easwars reviewed Apr 24, 2025

View reviewed changes

easwars assigned purnesh42H and unassigned easwars Apr 24, 2025

easwars reviewed Apr 24, 2025

View reviewed changes

easwars review 1

a263158

purnesh42H force-pushed the generic-xds-client-lrs-client-e2e branch from 9b88d68 to a263158 Compare April 25, 2025 06:15

purnesh42H requested a review from easwars April 25, 2025 06:17

purnesh42H assigned easwars and unassigned purnesh42H Apr 25, 2025

		// appendClusterStats gets Data for the given cluster, append to ret, and return
		// the new slice.

xds: generic lrs client for load reporting #8250

Are you sure you want to change the base?

xds: generic lrs client for load reporting #8250

Conversation

purnesh42H commented Apr 15, 2025 • edited Loading

codecov bot commented Apr 15, 2025 • edited Loading

Codecov Report

purnesh42H commented Apr 21, 2025 • edited Loading

dfawley left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purnesh42H Apr 23, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dfawley commented Apr 21, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purnesh42H Apr 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purnesh42H Apr 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purnesh42H Apr 25, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purnesh42H commented Apr 15, 2025 •

edited

Loading

codecov bot commented Apr 15, 2025 •

edited

Loading

purnesh42H commented Apr 21, 2025 •

edited

Loading

purnesh42H Apr 23, 2025 •

edited

Loading

purnesh42H Apr 25, 2025 •

edited

Loading

purnesh42H Apr 25, 2025 •

edited

Loading

purnesh42H Apr 25, 2025 •

edited

Loading