revert back loadStore.Stop() to accept context

Author: purnesh42H

xds/internal/balancer/clusterimpl/clusterimpl.go

@@ -24,6 +24,7 @@
 package clusterimpl
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"sync"
@@ -99,7 +100,7 @@ type clusterImplBalancer struct {
 	// The following fields are only accessed from balancer API methods, which
 	// are guaranteed to be called serially by gRPC.
 	xdsClient        xdsclient.XDSClient     // Sent down in ResolverState attributes.
-	cancelLoadReport func(time.Duration)     // To stop reporting load through the above xDS client.
+	cancelLoadReport func(context.Context)   // To stop reporting load through the above xDS client.
 	edsServiceName   string                  // EDS service name to report load for.
 	lrsServer        *bootstrap.ServerConfig // Load reporting server configuration.
 	dropCategories   []DropConfig            // The categories for drops.
@@ -221,7 +222,9 @@ func (b *clusterImplBalancer) updateLoadStore(newConfig *LBConfig) error {
 
 	if stopOldLoadReport {
 		if b.cancelLoadReport != nil {
-			b.cancelLoadReport(loadStoreStopTimeout)
+			stopCtx, stopCancel := context.WithTimeout(context.Background(), loadStoreStopTimeout)
+			defer stopCancel()
+			b.cancelLoadReport(stopCtx)
 			b.cancelLoadReport = nil
 			if !startNewLoadReport {
 				// If a new LRS stream will be started later, no need to update
@@ -347,7 +350,9 @@ func (b *clusterImplBalancer) Close() {
 	b.childState = balancer.State{}
 
 	if b.cancelLoadReport != nil {
-		b.cancelLoadReport(loadStoreStopTimeout)
+		stopCtx, stopCancel := context.WithTimeout(context.Background(), loadStoreStopTimeout)
+		defer stopCancel()
+		b.cancelLoadReport(stopCtx)
 		b.cancelLoadReport = nil
 	}
 	b.logger.Infof("Shutdown")

xds/internal/clients/lrsclient/load_store.go

@@ -19,6 +19,7 @@
 package lrsclient
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -36,7 +37,7 @@ import (
 // It is safe for concurrent use.
 type LoadStore struct {
 	// stop is the function to call to Stop the LoadStore reporting.
-	stop func(timeout time.Duration)
+	stop func(ctx context.Context)
 
 	// mu only protects the map (2 layers). The read/write to
 	// *PerClusterReporter doesn't need to hold the mu.
@@ -66,13 +67,13 @@ func newLoadStore() *LoadStore {
 // 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 timeout duration
+// final load report send attempt completes or the provided context `ctx`
 // expires.
 //
-// The `timeout` duration should be set to prevent Stop from blocking
-// indefinitely if the final send attempt fails to complete.
-func (ls *LoadStore) Stop(timeout time.Duration) {
-	ls.stop(timeout)
+// The provided context must have a deadline or timeout set to prevent Stop
+// from blocking indefinitely if the final send attempt fails to complete.
+func (ls *LoadStore) Stop(ctx context.Context) {
+	ls.stop(ctx)
 }
 
 // ReporterForCluster returns the PerClusterReporter for the given cluster and

xds/internal/clients/lrsclient/loadreport_test.go

@@ -147,7 +147,9 @@ func (s) TestReportLoad_ConnectionCreation(t *testing.T) {
 	if err != nil {
 		t.Fatalf("client.ReportLoad() failed: %v", err)
 	}
-	defer loadStore1.Stop(defaultTestShortTimeout)
+	ssCtx, ssCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer ssCancel()
+	defer loadStore1.Stop(ssCtx)
 
 	// Call the load reporting API to report load to the first management
 	// server, and ensure that a connection to the server is created.
@@ -232,7 +234,9 @@ func (s) TestReportLoad_ConnectionCreation(t *testing.T) {
 	}
 
 	// Stop this load reporting stream, server should see error canceled.
-	loadStore2.Stop(defaultTestShortTimeout)
+	ssCtx, ssCancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer ssCancel()
+	loadStore2.Stop(ssCtx)
 
 	// Server should receive a stream canceled error. There may be additional
 	// load reports from the client in the channel.
@@ -419,15 +423,19 @@ func (s) TestReportLoad_StreamCreation(t *testing.T) {
 
 	// Cancel the first load reporting call, and ensure that the stream does not
 	// close (because we have another call open).
-	loadStore1.Stop(defaultTestShortTimeout)
+	ssCtx, ssCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer ssCancel()
+	loadStore1.Stop(ssCtx)
 	sCtx, sCancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
 	defer sCancel()
 	if _, err := lrsServer.LRSStreamCloseChan.Receive(sCtx); err != context.DeadlineExceeded {
 		t.Fatal("LRS stream closed when expected to stay open")
 	}
 
 	// Stop the second load reporting call, and ensure the stream is closed.
-	loadStore2.Stop(defaultTestShortTimeout)
+	ssCtx, ssCancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer ssCancel()
+	loadStore2.Stop(ssCtx)
 	if _, err := lrsServer.LRSStreamCloseChan.Receive(ctx); err != nil {
 		t.Fatal("Timeout waiting for LRS stream to close")
 	}
@@ -442,16 +450,18 @@ func (s) TestReportLoad_StreamCreation(t *testing.T) {
 	if _, err := lrsServer.LRSStreamOpenChan.Receive(ctx); err != nil {
 		t.Fatalf("Timeout when waiting for LRS stream to be created: %v", err)
 	}
-	loadStore3.Stop(defaultTestShortTimeout)
+	ssCtx, ssCancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer ssCancel()
+	loadStore3.Stop(ssCtx)
 }
 
-// TestReportLoad_StopWithTimeout tests the behavior of LoadStore.Stop() when
-// called with a timeout duration. It verifies that:
-//   - Stop() blocks until the timeout expires or final load send attempt is
+// TestReportLoad_StopWithContext tests the behavior of LoadStore.Stop() when
+// called with a context. It verifies that:
+//   - Stop() blocks until the context expires or final load send attempt is
 //     made.
 //   - Final load report is seen on the server after stop is called.
 //   - The stream is closed after Stop() returns.
-func (s) TestReportLoad_StopWithTimeout(t *testing.T) {
+func (s) TestReportLoad_StopWithContext(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 
@@ -536,11 +546,11 @@ func (s) TestReportLoad_StopWithTimeout(t *testing.T) {
 		t.Fatalf("Unexpected diff in LRS request (-got, +want):\n%s", diff)
 	}
 
-	// Create a timeout duration for Stop() that remains until the end of test
-	// to ensure that only possibility of Stop() to finish is if final load
-	// send attempt is made. If final load attempt is not made, test itself
-	// will timeout.
-	largeStopTimeout := 10 * defaultTestTimeout
+	// Create a context for Stop() that remains until the end of test to ensure
+	// that only possibility of Stop()s to finish is if final load send attempt
+	// is made. If final load attempt is not made, test will timeout.
+	stopCtx, stopCancel := context.WithCancel(ctx)
+	defer stopCancel()
 
 	// Push more loads.
 	loadStore.ReporterForCluster("cluster2", "eds2").CallDropped("test")
@@ -549,7 +559,7 @@ func (s) TestReportLoad_StopWithTimeout(t *testing.T) {
 	// Call Stop in a separate goroutine. It will block until
 	// final load send attempt is made.
 	go func() {
-		loadStore.Stop(largeStopTimeout)
+		loadStore.Stop(stopCtx)
 		close(stopCalled)
 	}()
 	<-stopCalled

xds/internal/clients/lrsclient/lrsclient.go

@@ -22,6 +22,7 @@
 package lrsclient
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sync"
@@ -140,8 +141,8 @@ func (c *LRSClient) getOrCreateLRSStream(serverIdentifier clients.ServerIdentifi
 	// 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. Before closing the stream, it waits for the provided
-	// timeout duration for the final load report attempt to complete.
-	stop := func(timeout time.Duration) {
+	// context to be done (timeout or cancellation).
+	stop := func(ctx context.Context) {
 		c.mu.Lock()
 		defer c.mu.Unlock()
 
@@ -156,16 +157,13 @@ func (c *LRSClient) getOrCreateLRSStream(serverIdentifier clients.ServerIdentifi
 
 		lrs.finalSendRequest <- struct{}{}
 
-		timer := time.NewTimer(timeout)
-		defer timer.Stop()
-
 		select {
 		case err := <-lrs.finalSendDone:
 			if err != nil {
 				c.logger.Warningf("Final send attempt failed: %v", err)
 			}
-		case <-timer.C:
-			c.logger.Warningf("Timed out before finishing the final send attempt: %v", err)
+		case <-ctx.Done():
+			c.logger.Warningf("Context canceled before finishing the final send attempt: %v", err)
 		}
 
 		lrs.cancelStream()

xds/internal/testutils/fakeclient/client.go

@@ -21,7 +21,6 @@ package fakeclient
 
 import (
 	"context"
-	"time"
 
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/internal/xds/bootstrap"
@@ -81,14 +80,14 @@ func (*stream) Recv() ([]byte, error) {
 }
 
 // ReportLoad starts reporting load about clusterName to server.
-func (xdsC *Client) ReportLoad(server *bootstrap.ServerConfig) (loadStore *lrsclient.LoadStore, cancel func(time.Duration)) {
+func (xdsC *Client) ReportLoad(server *bootstrap.ServerConfig) (loadStore *lrsclient.LoadStore, cancel func(context.Context)) {
 	lrsClient, _ := lrsclient.New(lrsclient.Config{Node: clients.Node{ID: "fake-node-id"}, TransportBuilder: &transportBuilder{}})
 	xdsC.loadStore, _ = lrsClient.ReportLoad(clients.ServerIdentifier{ServerURI: server.ServerURI()})
 
 	xdsC.loadReportCh.Send(ReportLoadArgs{Server: server})
 
-	return xdsC.loadStore, func(timeout time.Duration) {
-		xdsC.loadStore.Stop(timeout)
+	return xdsC.loadStore, func(ctx context.Context) {
+		xdsC.loadStore.Stop(ctx)
 		xdsC.lrsCancelCh.Send(nil)
 	}
 }

xds/internal/xdsclient/client.go

@@ -21,7 +21,7 @@
 package xdsclient
 
 import (
-	"time"
+	"context"
 
 	v3statuspb "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
 	"google.golang.org/grpc/internal/xds/bootstrap"
@@ -49,7 +49,7 @@ type XDSClient interface {
 	// the watcher is canceled. Callers need to handle this case.
 	WatchResource(rType xdsresource.Type, resourceName string, watcher xdsresource.ResourceWatcher) (cancel func())
 
-	ReportLoad(*bootstrap.ServerConfig) (*lrsclient.LoadStore, func(time.Duration))
+	ReportLoad(*bootstrap.ServerConfig) (*lrsclient.LoadStore, func(context.Context))
 
 	BootstrapConfig() *bootstrap.Config
 }

xds/internal/xdsclient/clientimpl_loadreport.go

@@ -18,8 +18,8 @@
 package xdsclient
 
 import (
+	"context"
 	"sync"
-	"time"
 
 	"google.golang.org/grpc/internal/xds/bootstrap"
 	"google.golang.org/grpc/xds/internal/clients"
@@ -31,26 +31,24 @@ import (
 // reports to the same server share the LRS stream.
 //
 // It returns a lrsclient.LoadStore for the user to report loads.
-func (c *clientImpl) ReportLoad(server *bootstrap.ServerConfig) (*lrsclient.LoadStore, func(time.Duration)) {
+func (c *clientImpl) ReportLoad(server *bootstrap.ServerConfig) (*lrsclient.LoadStore, func(context.Context)) {
 	if c.lrsClient == nil {
 		lrsConfig := lrsclient.Config{Node: c.gConfig.Node, TransportBuilder: c.gConfig.TransportBuilder}
 		lrsC, err := lrsclient.New(lrsConfig)
 		if err != nil {
 			c.logger.Warningf("Failed to create an lrs client to the management server to report load: %v", server, err)
-			return nil, func(time.Duration) {}
+			return nil, func(context.Context) {}
 		}
 		c.lrsClient = lrsC
 	}
 
 	load, err := c.lrsClient.ReportLoad(clients.ServerIdentifier{ServerURI: server.ServerURI(), Extensions: grpctransport.ServerIdentifierExtension{ConfigName: server.SelectedCreds().Type}})
 	if err != nil {
 		c.logger.Warningf("Failed to create a load store to the management server to report load: %v", server, err)
-		return nil, func(time.Duration) {}
+		return nil, func(context.Context) {}
 	}
 	var loadStop sync.Once
-	return load, func(timeout time.Duration) {
-		loadStop.Do(func() {
-			load.Stop(timeout)
-		})
+	return load, func(ctx context.Context) {
+		loadStop.Do(func() { load.Stop(ctx) })
 	}
 }

xds/internal/xdsclient/tests/loadreport_test.go

@@ -144,7 +144,9 @@ func (s) TestReportLoad_ConnectionCreation(t *testing.T) {
 	// Call the load reporting API to report load to the first management
 	// server, and ensure that a connection to the server is created.
 	store1, lrsCancel1 := client.ReportLoad(serverCfg1)
-	defer lrsCancel1(defaultTestShortTimeout)
+	sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel()
+	defer lrsCancel1(sCtx)
 	if _, err := newConnChan1.Receive(ctx); err != nil {
 		t.Fatal("Timeout when waiting for a connection to the first management server, after starting load reporting")
 	}
@@ -159,7 +161,9 @@ func (s) TestReportLoad_ConnectionCreation(t *testing.T) {
 	// Call the load reporting API to report load to the second management
 	// server, and ensure that a connection to the server is created.
 	store2, lrsCancel2 := client.ReportLoad(serverCfg2)
-	defer lrsCancel2(defaultTestShortTimeout)
+	sCtx2, sCancel2 := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel2()
+	defer lrsCancel2(sCtx2)
 	if _, err := newConnChan2.Receive(ctx); err != nil {
 		t.Fatal("Timeout when waiting for a connection to the second management server, after starting load reporting")
 	}
@@ -227,7 +231,9 @@ func (s) TestReportLoad_ConnectionCreation(t *testing.T) {
 	}
 
 	// Cancel this load reporting stream, server should see error canceled.
-	lrsCancel2(defaultTestShortTimeout)
+	sCtx2, sCancel2 = context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel2()
+	lrsCancel2(sCtx2)
 
 	// Server should receive a stream canceled error. There may be additional
 	// load reports from the client in the channel.
@@ -403,15 +409,19 @@ func (s) TestReportLoad_StreamCreation(t *testing.T) {
 
 	// Cancel the first load reporting call, and ensure that the stream does not
 	// close (because we have another call open).
-	cancel1(defaultTestShortTimeout)
+	sCtx1, sCancel1 := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel1()
+	cancel1(sCtx1)
 	sCtx, sCancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
 	defer sCancel()
 	if _, err := lrsServer.LRSStreamCloseChan.Receive(sCtx); err != context.DeadlineExceeded {
 		t.Fatal("LRS stream closed when expected to stay open")
 	}
 
 	// Cancel the second load reporting call, and ensure the stream is closed.
-	cancel2(defaultTestShortTimeout)
+	sCtx2, sCancel2 := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel2()
+	cancel2(sCtx2)
 	if _, err := lrsServer.LRSStreamCloseChan.Receive(ctx); err != nil {
 		t.Fatal("Timeout waiting for LRS stream to close")
 	}
@@ -423,5 +433,7 @@ func (s) TestReportLoad_StreamCreation(t *testing.T) {
 	if _, err := lrsServer.LRSStreamOpenChan.Receive(ctx); err != nil {
 		t.Fatalf("Timeout when waiting for LRS stream to be created: %v", err)
 	}
-	cancel3(defaultTestShortTimeout)
+	sCtx3, sCancel3 := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel3()
+	cancel3(sCtx3)
 }