google.golang.org/grpc@v1.72.2/producer_ext_test.go (about)

     1  /*
     2   *
     3   * Copyright 2024 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package grpc_test
    20  
    21  import (
    22  	"context"
    23  	"strings"
    24  	"sync/atomic"
    25  	"testing"
    26  
    27  	"google.golang.org/grpc"
    28  	"google.golang.org/grpc/balancer"
    29  	"google.golang.org/grpc/connectivity"
    30  	"google.golang.org/grpc/credentials/insecure"
    31  	"google.golang.org/grpc/internal/balancer/stub"
    32  	"google.golang.org/grpc/internal/stubserver"
    33  	"google.golang.org/grpc/internal/testutils"
    34  	testgrpc "google.golang.org/grpc/interop/grpc_testing"
    35  )
    36  
    37  // TestProducerStopsBeforeStateChange confirms that producers are stopped before
    38  // any state change notification is delivered to the LB policy.
    39  func (s) TestProducerStopsBeforeStateChange(t *testing.T) {
    40  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
    41  	defer cancel()
    42  
    43  	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
    44  	var lastProducer *testProducer
    45  	bf := stub.BalancerFuncs{
    46  		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
    47  			var sc balancer.SubConn
    48  			sc, err := bd.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{
    49  				StateListener: func(scs balancer.SubConnState) {
    50  					bd.ClientConn.UpdateState(balancer.State{
    51  						ConnectivityState: scs.ConnectivityState,
    52  						// We do not pass a picker, but since we don't perform
    53  						// RPCs, that's okay.
    54  					})
    55  					if !lastProducer.stopped.Load() {
    56  						t.Errorf("lastProducer not stopped before state change notification")
    57  					}
    58  					t.Logf("State is now %v; recreating producer", scs.ConnectivityState)
    59  					p, _ := sc.GetOrBuildProducer(producerBuilderSingleton)
    60  					lastProducer = p.(*testProducer)
    61  				},
    62  			})
    63  			if err != nil {
    64  				return err
    65  			}
    66  			p, _ := sc.GetOrBuildProducer(producerBuilderSingleton)
    67  			lastProducer = p.(*testProducer)
    68  			sc.Connect()
    69  			return nil
    70  		},
    71  	}
    72  	stub.Register(name, bf)
    73  
    74  	ss := stubserver.StubServer{
    75  		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
    76  			return nil
    77  		},
    78  	}
    79  	if err := ss.StartServer(); err != nil {
    80  		t.Fatal("Error starting server:", err)
    81  	}
    82  	defer ss.Stop()
    83  
    84  	cc, err := grpc.NewClient("dns:///"+ss.Address,
    85  		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"`+name+`":{}}]}`),
    86  		grpc.WithTransportCredentials(insecure.NewCredentials()),
    87  	)
    88  	if err != nil {
    89  		t.Fatalf("Error creating client: %v", err)
    90  	}
    91  	defer cc.Close()
    92  
    93  	go cc.Connect()
    94  	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
    95  
    96  	cc.Close()
    97  	testutils.AwaitState(ctx, t, cc, connectivity.Shutdown)
    98  }
    99  
   100  type producerBuilder struct{}
   101  
   102  type testProducer struct {
   103  	// There should be no race accessing this field, but use an atomic since
   104  	// the race checker probably can't detect that.
   105  	stopped atomic.Bool
   106  }
   107  
   108  // Build constructs and returns a producer and its cleanup function
   109  func (*producerBuilder) Build(cci any) (balancer.Producer, func()) {
   110  	p := &testProducer{}
   111  	return p, func() {
   112  		p.stopped.Store(true)
   113  	}
   114  }
   115  
   116  var producerBuilderSingleton = &producerBuilder{}