easwars review 1

Author: purnesh42H

xds/internal/clients/lrsclient/load_store.go

@@ -33,7 +33,7 @@ import (
 //
 // It is safe for concurrent use.
 type LoadStore struct {
-	lrsStream *streamImpl
+	stop func(ctx context.Context) // Function to call to Stop the LoadStore
 
 	// mu only protects the map (2 layers). The read/write to
 	// *PerClusterReporter doesn't need to hold the mu.
@@ -49,31 +49,23 @@ type LoadStore struct {
 	clusters map[string]map[string]*PerClusterReporter
 }
 
-// newStore creates a LoadStore.
-func newLoadStore(lrsStream *streamImpl) *LoadStore {
+// newLoadStore creates a LoadStore.
+func newLoadStore() *LoadStore {
 	return &LoadStore{
-		clusters:  make(map[string]map[string]*PerClusterReporter),
-		lrsStream: lrsStream,
+		clusters: make(map[string]map[string]*PerClusterReporter),
 	}
 }
 
 // Stop stops the LRS stream associated with this LoadStore.
 //
-// If this LoadStore is the only one using the underlying LRS stream, the
-// stream will be closed. If other LoadStores are also using the same stream,
-// the reference count to the stream is decremented, and the stream remains
-// open until all LoadStores have called Stop().
+// If this is the last reference to the underlying LRS stream, the
+// stream will be closed.
 //
-// If this is the last LoadStore for the stream, this method makes a last
-// 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.
+// The provided context should have a deadline or timeout set. If this is the
+// last reference, Stop will block until the context is done, allowing any
+// pending load reports to be flushed before closing the stream.
 func (ls *LoadStore) Stop(ctx context.Context) {
-	// Wait for the provided context to be done (timeout or cancellation).
-	if ctx != nil {
-		<-ctx.Done()
-	}
-	ls.lrsStream.stop()
+	ls.stop(ctx)
 }
 
 // ReporterForCluster returns the PerClusterReporter for the given cluster and
@@ -111,22 +103,22 @@ func (ls *LoadStore) ReporterForCluster(clusterName, serviceName string) *PerClu
 // 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
 	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
 }
 
@@ -314,12 +306,12 @@ type serverLoadData struct {
 	sum float64
 }
 
-// appendClusterStats gets Data for the given cluster, append to ret, and return
-// the new slice.
+// 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, cluster map[string]*PerClusterReporter) []*loadData {
-	for _, d := range cluster {
+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.

xds/internal/clients/lrsclient/load_store_test.go

@@ -39,6 +39,13 @@ type rpcData struct {
 	serverData              map[string]float64 // Will be reported with successful RPCs.
 }
 
+func verifyLoadStoreData(wantStoreData, gotStoreData []*loadData) error {
+	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval"), sortDataSlice); diff != "" {
+		return fmt.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	}
+	return nil
+}
+
 // TestDrops spawns a bunch of goroutines which report drop data. After the
 // goroutines have exited, the test dumps the stats from the Store and makes
 // sure they are as expected.
@@ -72,8 +79,8 @@ func TestDrops(t *testing.T) {
 	wg.Wait()
 
 	gotStoreData := ls.stats()
-	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval")); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData([]*loadData{wantStoreData}, []*loadData{gotStoreData}); err != nil {
+		t.Error(err)
 	}
 }
 
@@ -165,8 +172,8 @@ func TestLocalityStats(t *testing.T) {
 	}
 
 	gotStoreData := ls.stats()
-	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval")); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData([]*loadData{wantStoreData}, []*loadData{gotStoreData}); err != nil {
+		t.Error(err)
 	}
 }
 
@@ -260,8 +267,8 @@ func TestResetAfterStats(t *testing.T) {
 	ls := PerClusterReporter{}
 	reportLoad(&ls)
 	gotStoreData := ls.stats()
-	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval")); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData([]*loadData{wantStoreData}, []*loadData{gotStoreData}); err != nil {
+		t.Error(err)
 	}
 
 	// The above call to stats() should have reset all load reports except the
@@ -274,8 +281,8 @@ func TestResetAfterStats(t *testing.T) {
 	}
 	reportLoad(&ls)
 	gotStoreData = ls.stats()
-	if diff := cmp.Diff(wantStoreData, gotStoreData, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval")); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData([]*loadData{wantStoreData}, []*loadData{gotStoreData}); err != nil {
+		t.Error(err)
 	}
 }
 
@@ -304,7 +311,7 @@ func TestStoreStats(t *testing.T) {
 		testLocality = "test-locality"
 	)
 
-	store := newLoadStore(nil)
+	store := newLoadStore()
 	for _, c := range testClusters {
 		for _, s := range testServices {
 			store.ReporterForCluster(c, s).CallStarted(testLocality)
@@ -339,9 +346,7 @@ func TestStoreStats(t *testing.T) {
 	// Call Stats with just "c0", this should return data for "c0", and not
 	// touch data for other clusters.
 	gotC0 := store.stats([]string{"c0"})
-	if diff := cmp.Diff(wantC0, gotC0, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval"), sortDataSlice); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
-	}
+	verifyLoadStoreData(wantC0, gotC0)
 
 	wantOther := []*loadData{
 		{
@@ -388,8 +393,8 @@ func TestStoreStats(t *testing.T) {
 	// Call Stats with empty slice, this should return data for all the
 	// remaining clusters, and not include c0 (because c0 data was cleared).
 	gotOther := store.stats(nil)
-	if diff := cmp.Diff(wantOther, gotOther, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval"), sortDataSlice); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData(wantOther, gotOther); err != nil {
+		t.Error(err)
 	}
 }
 
@@ -401,7 +406,7 @@ func TestStoreStatsEmptyDataNotReported(t *testing.T) {
 		testLocality = "test-locality"
 	)
 
-	store := newLoadStore(nil)
+	store := newLoadStore()
 	// "c0"'s RPCs all finish with success.
 	for _, s := range testServices {
 		store.ReporterForCluster("c0", s).CallStarted(testLocality)
@@ -441,8 +446,8 @@ func TestStoreStatsEmptyDataNotReported(t *testing.T) {
 	// Call Stats with empty slice, this should return data for all the
 	// clusters.
 	got0 := store.stats(nil)
-	if diff := cmp.Diff(want0, got0, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval"), sortDataSlice); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData(want0, got0); err != nil {
+		t.Error(err)
 	}
 
 	want1 := []*loadData{
@@ -462,7 +467,7 @@ func TestStoreStatsEmptyDataNotReported(t *testing.T) {
 	// Call Stats with empty slice again, this should return data only for "c1",
 	// because "c0" data was cleared, but "c1" has in-progress RPCs.
 	got1 := store.stats(nil)
-	if diff := cmp.Diff(want1, got1, cmpopts.EquateEmpty(), cmp.AllowUnexported(loadData{}, localityData{}, requestData{}, serverLoadData{}), cmpopts.IgnoreFields(loadData{}, "reportInterval"), sortDataSlice); diff != "" {
-		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
+	if err := verifyLoadStoreData(want1, got1); err != nil {
+		t.Error(err)
 	}
 }

xds/internal/clients/lrsclient/lrs_stream.go

@@ -59,7 +59,6 @@ type streamImpl struct {
 
 	cancelStream context.CancelFunc // Cancel the stream. If nil, the stream is not active.
 	loadStore    *LoadStore         // LoadStore returned to user for pushing loads.
-	cleanup      func()             // Function to call after the stream is stopped.
 }
 
 // streamOpts holds the options for creating an lrsStream.
@@ -75,14 +74,20 @@ type streamOpts struct {
 // The actual streaming RPC call is initiated when the first call to ReportLoad
 // is made, and is terminated when the last call to ReportLoad is canceled.
 func newStreamImpl(opts streamOpts) *streamImpl {
+	ctx, cancel := context.WithCancel(context.Background())
+
 	lrs := &streamImpl{
-		transport: opts.transport,
-		backoff:   opts.backoff,
-		nodeProto: opts.nodeProto,
+		transport:    opts.transport,
+		backoff:      opts.backoff,
+		nodeProto:    opts.nodeProto,
+		cancelStream: cancel,
+		doneCh:       make(chan struct{}),
 	}
 
 	l := grpclog.Component("xds")
 	lrs.logger = igrpclog.NewPrefixLogger(l, opts.logPrefix+fmt.Sprintf("[lrs-stream %p] ", lrs))
+	lrs.loadStore = newLoadStore()
+	go lrs.runner(ctx)
 	return lrs
 }
 
@@ -183,7 +188,9 @@ func (lrs *streamImpl) recvFirstLoadStatsResponse(stream clients.Stream) ([]stri
 	}
 	var resp v3lrspb.LoadStatsResponse
 	if err := proto.Unmarshal(r, &resp); err != nil {
-		lrs.logger.Infof("Failed to unmarshal response to LoadStatsResponse: %v", err)
+		if lrs.logger.V(2) {
+			lrs.logger.Infof("Failed to unmarshal response to LoadStatsResponse: %v", err)
+		}
 		return nil, time.Duration(0), fmt.Errorf("lrs: unexpected message type %T", r)
 	}
 	if lrs.logger.V(perRPCVerbosityLevel) {
@@ -260,7 +267,9 @@ func (lrs *streamImpl) sendLoadStatsRequest(stream clients.Stream, loads []*load
 	}
 	msg, err := proto.Marshal(req)
 	if err != nil {
-		lrs.logger.Warningf("Failed to marshal LoadStatsRequest: %v", err)
+		if lrs.logger.V(2) {
+			lrs.logger.Infof("Failed to marshal LoadStatsRequest: %v", err)
+		}
 		return err
 	}
 	err = stream.Send(msg)
@@ -278,11 +287,6 @@ func getStreamError(stream clients.Stream) error {
 	}
 }
 
-// stop calls the registered cleanup function by LRS client to stop the stream.
-func (lrs *streamImpl) stop() {
-	lrs.cleanup()
-}
-
 // localityFromString converts a json representation of locality, into a
 // clients.Locality struct.
 func localityFromString(s string) (ret clients.Locality, _ error) {

xds/internal/clients/lrsclient/lrsclient.go

@@ -63,7 +63,7 @@ type LRSClient struct {
 func New(config Config) (*LRSClient, error) {
 	switch {
 	case config.Node.ID == "":
-		return nil, errors.New("lrsclient: node ID is empty")
+		return nil, errors.New("lrsclient: node ID in node is empty")
 	case config.TransportBuilder == nil:
 		return nil, errors.New("lrsclient: transport builder is nil")
 	}
@@ -81,6 +81,9 @@ func New(config Config) (*LRSClient, error) {
 
 // ReportLoad creates and returns a LoadStore for the caller to report loads
 // using a LoadReportingStream.
+//
+// Caller must call Stop on the returned LoadStore when they are done reporting
+// load to this server.
 func (c *LRSClient) ReportLoad(si clients.ServerIdentifier) (*LoadStore, error) {
 	lrs, err := c.getOrCreateLRSStream(si)
 	if err != nil {
@@ -133,17 +136,13 @@ func (c *LRSClient) getOrCreateLRSStream(serverIdentifier clients.ServerIdentifi
 		nodeProto: nodeProto,
 		logPrefix: logPrefix,
 	})
-	ctx, cancel := context.WithCancel(context.Background())
-	lrs.cancelStream = cancel
-	lrs.doneCh = make(chan struct{})
-	lrs.loadStore = newLoadStore(lrs)
-	go lrs.runner(ctx)
 
-	// Register a cleanup function that decrements the reference count, stops
+	// Register a stop function that decrements the reference count, stops
 	// the LRS stream when the last reference is removed and closes the
 	// transport and removes the lrs stream and its references from the
-	// respective maps.
-	cleanup := func() {
+	// respective maps. Before closing the stream, it waits for the provided
+	// context to be done (timeout or cancellation).
+	stop := func(ctx context.Context) {
 		c.mu.Lock()
 		defer c.mu.Unlock()
 
@@ -156,12 +155,12 @@ func (c *LRSClient) getOrCreateLRSStream(serverIdentifier clients.ServerIdentifi
 			return
 		}
 
-		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
+		// Wait for the provided context to be done (timeout or cancellation)
+		// before closing the stream.
+		if ctx != nil {
+			<-ctx.Done()
 		}
+
 		lrs.cancelStream()
 		lrs.cancelStream = nil
 		lrs.logger.Infof("Stopping LRS stream")
@@ -170,7 +169,7 @@ func (c *LRSClient) getOrCreateLRSStream(serverIdentifier clients.ServerIdentifi
 		delete(c.lrsStreams, serverIdentifier)
 		tr.Close()
 	}
-	lrs.cleanup = cleanup
+	lrs.loadStore.stop = stop
 
 	c.lrsStreams[serverIdentifier] = lrs
 	c.lrsRefs[serverIdentifier] = 1