github.com/bazelbuild/remote-apis-sdks@v0.0.0-20240425170053-8a36686a6350/go/pkg/balancer/gcp_balancer_test.go (about)

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/google/uuid"
     9  
    10  	grpcbalancer "google.golang.org/grpc/balancer"
    11  	"google.golang.org/grpc/connectivity"
    12  	"google.golang.org/grpc/resolver"
    13  )
    14  
    15  func TestGCPBalancer_UpdatingConnectionStateIsMutuallyExclusive(t *testing.T) {
    16  	builder := newBuilder()
    17  	var subconns []*fakeSubConn
    18  	for i := 0; i < MinConnections; i++ {
    19  		subconns = append(subconns, &fakeSubConn{id: uuid.New().String()})
    20  	}
    21  
    22  	cc := fakeClientConn{subconns: subconns}
    23  	balancer := builder.Build(cc, grpcbalancer.BuildOptions{})
    24  	var wg sync.WaitGroup
    25  	wg.Add(2)
    26  
    27  	// Invoke both UpdateClientConnState() and UpdatesubConnState() simultaneously
    28  	// and make sure that it doesn't result in a race condition. The test would fail
    29  	// if there's a race condition.
    30  	go func() {
    31  		balancer.UpdateClientConnState(grpcbalancer.ClientConnState{})
    32  		wg.Done()
    33  	}()
    34  	go func() {
    35  		for i := 0; i < MinConnections; i++ {
    36  			// Sending a connectivity.Ready state will actually trigger regenerating the picker
    37  			// which is where we expect the race condition to occur.
    38  			balancer.UpdateSubConnState(subconns[i], grpcbalancer.SubConnState{ConnectivityState: connectivity.Ready})
    39  		}
    40  		wg.Done()
    41  	}()
    42  	wg.Wait()
    43  }
    44  
    45  type fakeClientConn struct {
    46  	grpcbalancer.ClientConn
    47  
    48  	subconns []*fakeSubConn
    49  }
    50  
    51  var (
    52  	// Note: These cannot be moved to fakeClientConn since gRPC Builder() interface
    53  	// does not accept a pointer to fakeClientConn and so the mutex will be copied if
    54  	// passed by value.
    55  	// https://godoc.org/google.golang.org/grpc/balancer#Builder
    56  	idx int
    57  	mu  sync.Mutex
    58  )
    59  
    60  func (f fakeClientConn) NewSubConn([]resolver.Address, grpcbalancer.NewSubConnOptions) (grpcbalancer.SubConn, error) {
    61  	mu.Lock()
    62  	defer mu.Unlock()
    63  	idx++
    64  	idx = idx % MinConnections
    65  	return f.subconns[idx], nil
    66  }
    67  
    68  func (f fakeClientConn) RemoveSubConn(grpcbalancer.SubConn) {}
    69  
    70  func (f fakeClientConn) UpdateState(grpcbalancer.State) {}
    71  
    72  func (f fakeClientConn) ResolveNow(resolver.ResolveNowOptions) {}
    73  
    74  func (f fakeClientConn) Target() string {
    75  	return ""
    76  }
    77  
    78  type fakeSubConn struct {
    79  	id string
    80  }
    81  
    82  func (fakeSubConn) UpdateAddresses([]resolver.Address) {}
    83  
    84  func (fakeSubConn) Connect() {
    85  	// Sleep to simulate connecting to an actual server.
    86  	time.Sleep(100 * time.Millisecond)
    87  }
    88  
    89  func (fakeSubConn) GetOrBuildProducer(grpcbalancer.ProducerBuilder) (grpcbalancer.Producer, func()) {
    90  	return nil, nil
    91  }
    92  
    93  func (fakeSubConn) Shutdown() {}