xds_test: Wait for server to enter serving mode in RBAC test (#8287)

test/xds/xds_server_integration_test.go

@@ -83,7 +83,7 @@ func (l *acceptNotifyingListener) Accept() (net.Conn, error) {
 // Returns the following:
 // - local listener on which the xDS-enabled gRPC server is serving on
 // - cleanup function to be invoked by the tests when done
-func setupGRPCServer(t *testing.T, bootstrapContents []byte) (net.Listener, func()) {
+func setupGRPCServer(t *testing.T, bootstrapContents []byte, opts ...grpc.ServerOption) (net.Listener, func()) {
 	t.Helper()
 
 	// Configure xDS credentials to be used on the server-side.
@@ -113,12 +113,15 @@ func setupGRPCServer(t *testing.T, bootstrapContents []byte) (net.Listener, func
 		},
 	}
 
-	if stub.S, err = xds.NewGRPCServer(grpc.Creds(creds), testModeChangeServerOption(t), xds.BootstrapContentsForTesting(bootstrapContents)); err != nil {
+	opts = append([]grpc.ServerOption{
+		grpc.Creds(creds),
+		testModeChangeServerOption(t),
+		xds.BootstrapContentsForTesting(bootstrapContents),
+	}, opts...)
+	if stub.S, err = xds.NewGRPCServer(opts...); err != nil {
 		t.Fatalf("Failed to create an xDS enabled gRPC server: %v", err)
 	}
 
-	stubserver.StartTestService(t, stub)
-
 	// Create a local listener and pass it to Serve().
 	lis, err := testutils.LocalTCPListener()
 	if err != nil {
@@ -130,11 +133,8 @@ func setupGRPCServer(t *testing.T, bootstrapContents []byte) (net.Listener, func
 		serverReady: *grpcsync.NewEvent(),
 	}
 
-	go func() {
+	stub.Listener = readyLis
-		if err := stub.S.Serve(readyLis); err != nil {
+	stubserver.StartTestService(t, stub)
-			t.Errorf("Serve() failed: %v", err)
-		}
-	}()
 
 	// Wait for the server to start running.
 	select {

test/xds/xds_server_rbac_test.go

@@ -33,12 +33,14 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/authz/audit"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/internal/testutils/xds/e2e"
 	"google.golang.org/grpc/internal/testutils/xds/e2e/setup"
 	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/xds"
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
@@ -842,10 +844,23 @@ func serverListenerWithBadRouteConfiguration(t *testing.T, host string, port uin
 
 func (s) TestRBAC_WithBadRouteConfiguration(t *testing.T) {
 	managementServer, nodeID, bootstrapContents, xdsResolver := setup.ManagementServerAndResolver(t)
+	// We need to wait for the server to enter SERVING mode before making RPCs
+	// to avoid flakes due to the server closing connections.
+	servingCh := make(chan struct{})
 
-	lis, cleanup2 := setupGRPCServer(t, bootstrapContents)
+	// Initialize a test gRPC server, assign it to the stub server, and start
+	// the test service.
+	opt := xds.ServingModeCallback(func(addr net.Addr, args xds.ServingModeChangeArgs) {
+		if args.Mode == connectivity.ServingModeServing {
+			close(servingCh)
+		}
+	})
+	lis, cleanup2 := setupGRPCServer(t, bootstrapContents, opt)
 	defer cleanup2()
 
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
 	host, port, err := hostPortFromListener(lis)
 	if err != nil {
 		t.Fatalf("failed to retrieve host and port of server: %v", err)
@@ -867,13 +882,17 @@ func (s) TestRBAC_WithBadRouteConfiguration(t *testing.T) {
 	inboundLis := serverListenerWithBadRouteConfiguration(t, host, port)
 	resources.Listeners = append(resources.Listeners, inboundLis)
 
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
 	// Setup the management server with client and server-side resources.
 	if err := managementServer.Update(ctx, resources); err != nil {
 		t.Fatal(err)
 	}
 
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for the xDS-enabled gRPC server to go SERVING")
+	case <-servingCh:
+	}
+
 	cc, err := grpc.NewClient(fmt.Sprintf("xds:///%s", serviceName), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(xdsResolver))
 	if err != nil {
 		t.Fatalf("grpc.NewClient() failed: %v", err)