picker_wrapper_test.go

Documentation: google.golang.org/grpc

     1  /*
     2   *
     3   * Copyright 2017 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
    20  
    21  import (
    22  	"context"
    23  	"fmt"
    24  	"sync/atomic"
    25  	"testing"
    26  	"time"
    27  
    28  	"google.golang.org/grpc/balancer"
    29  	"google.golang.org/grpc/codes"
    30  	"google.golang.org/grpc/connectivity"
    31  	"google.golang.org/grpc/internal/transport"
    32  	"google.golang.org/grpc/status"
    33  )
    34  
    35  const goroutineCount = 5
    36  
    37  var (
    38  	testT  = &testTransport{}
    39  	testSC = &acBalancerWrapper{ac: &addrConn{
    40  		state:     connectivity.Ready,
    41  		transport: testT,
    42  	}}
    43  	testSCNotReady = &acBalancerWrapper{ac: &addrConn{
    44  		state: connectivity.TransientFailure,
    45  	}}
    46  )
    47  
    48  type testTransport struct {
    49  	transport.ClientTransport
    50  }
    51  
    52  type testingPicker struct {
    53  	err       error
    54  	sc        balancer.SubConn
    55  	maxCalled int64
    56  }
    57  
    58  func (p *testingPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
    59  	if atomic.AddInt64(&p.maxCalled, -1) < 0 {
    60  		return balancer.PickResult{}, fmt.Errorf("pick called to many times (> goroutineCount)")
    61  	}
    62  	if p.err != nil {
    63  		return balancer.PickResult{}, p.err
    64  	}
    65  	return balancer.PickResult{SubConn: p.sc}, nil
    66  }
    67  
    68  func (s) TestBlockingPickTimeout(t *testing.T) {
    69  	bp := newPickerWrapper(nil)
    70  	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
    71  	defer cancel()
    72  	if _, _, err := bp.pick(ctx, true, balancer.PickInfo{}); status.Code(err) != codes.DeadlineExceeded {
    73  		t.Errorf("bp.pick returned error %v, want DeadlineExceeded", err)
    74  	}
    75  }
    76  
    77  func (s) TestBlockingPick(t *testing.T) {
    78  	bp := newPickerWrapper(nil)
    79  	// All goroutines should block because picker is nil in bp.
    80  	var finishedCount uint64
    81  	for i := goroutineCount; i > 0; i-- {
    82  		go func() {
    83  			if tr, _, err := bp.pick(context.Background(), true, balancer.PickInfo{}); err != nil || tr != testT {
    84  				t.Errorf("bp.pick returned non-nil error: %v", err)
    85  			}
    86  			atomic.AddUint64(&finishedCount, 1)
    87  		}()
    88  	}
    89  	time.Sleep(50 * time.Millisecond)
    90  	if c := atomic.LoadUint64(&finishedCount); c != 0 {
    91  		t.Errorf("finished goroutines count: %v, want 0", c)
    92  	}
    93  	bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
    94  }
    95  
    96  func (s) TestBlockingPickNoSubAvailable(t *testing.T) {
    97  	bp := newPickerWrapper(nil)
    98  	var finishedCount uint64
    99  	bp.updatePicker(&testingPicker{err: balancer.ErrNoSubConnAvailable, maxCalled: goroutineCount})
   100  	// All goroutines should block because picker returns no subConn available.
   101  	for i := goroutineCount; i > 0; i-- {
   102  		go func() {
   103  			if tr, _, err := bp.pick(context.Background(), true, balancer.PickInfo{}); err != nil || tr != testT {
   104  				t.Errorf("bp.pick returned non-nil error: %v", err)
   105  			}
   106  			atomic.AddUint64(&finishedCount, 1)
   107  		}()
   108  	}
   109  	time.Sleep(50 * time.Millisecond)
   110  	if c := atomic.LoadUint64(&finishedCount); c != 0 {
   111  		t.Errorf("finished goroutines count: %v, want 0", c)
   112  	}
   113  	bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
   114  }
   115  
   116  func (s) TestBlockingPickTransientWaitforready(t *testing.T) {
   117  	bp := newPickerWrapper(nil)
   118  	bp.updatePicker(&testingPicker{err: balancer.ErrTransientFailure, maxCalled: goroutineCount})
   119  	var finishedCount uint64
   120  	// All goroutines should block because picker returns transientFailure and
   121  	// picks are not failfast.
   122  	for i := goroutineCount; i > 0; i-- {
   123  		go func() {
   124  			if tr, _, err := bp.pick(context.Background(), false, balancer.PickInfo{}); err != nil || tr != testT {
   125  				t.Errorf("bp.pick returned non-nil error: %v", err)
   126  			}
   127  			atomic.AddUint64(&finishedCount, 1)
   128  		}()
   129  	}
   130  	time.Sleep(time.Millisecond)
   131  	if c := atomic.LoadUint64(&finishedCount); c != 0 {
   132  		t.Errorf("finished goroutines count: %v, want 0", c)
   133  	}
   134  	bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
   135  }
   136  
   137  func (s) TestBlockingPickSCNotReady(t *testing.T) {
   138  	bp := newPickerWrapper(nil)
   139  	bp.updatePicker(&testingPicker{sc: testSCNotReady, maxCalled: goroutineCount})
   140  	var finishedCount uint64
   141  	// All goroutines should block because subConn is not ready.
   142  	for i := goroutineCount; i > 0; i-- {
   143  		go func() {
   144  			if tr, _, err := bp.pick(context.Background(), true, balancer.PickInfo{}); err != nil || tr != testT {
   145  				t.Errorf("bp.pick returned non-nil error: %v", err)
   146  			}
   147  			atomic.AddUint64(&finishedCount, 1)
   148  		}()
   149  	}
   150  	time.Sleep(time.Millisecond)
   151  	if c := atomic.LoadUint64(&finishedCount); c != 0 {
   152  		t.Errorf("finished goroutines count: %v, want 0", c)
   153  	}
   154  	bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
   155  }
   156
View as plain text