vendor/google.golang.org/grpc/picker_wrapper_test.go

   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         "fmt"
  23         "sync/atomic"
  24         "testing"
  25         "time"
  26
  27         "golang.org/x/net/context"
  28         "google.golang.org/grpc/balancer"
  29         "google.golang.org/grpc/connectivity"
  30         _ "google.golang.org/grpc/grpclog/glogger"
  31         "google.golang.org/grpc/test/leakcheck"
  32         "google.golang.org/grpc/transport"
  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(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
  59         if atomic.AddInt64(&p.maxCalled, -1) < 0 {
  60                 return nil, nil, fmt.Errorf("Pick called to many times (> goroutineCount)")
  61         }
  62         if p.err != nil {
  63                 return nil, nil, p.err
  64         }
  65         return p.sc, nil, nil
  66 }
  67
  68 func TestBlockingPickTimeout(t *testing.T) {
  69         defer leakcheck.Check(t)
  70         bp := newPickerWrapper()
  71         ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
  72         defer cancel()
  73         if _, _, err := bp.pick(ctx, true, balancer.PickOptions{}); err != context.DeadlineExceeded {
  74                 t.Errorf("bp.pick returned error %v, want DeadlineExceeded", err)
  75         }
  76 }
  77
  78 func TestBlockingPick(t *testing.T) {
  79         defer leakcheck.Check(t)
  80         bp := newPickerWrapper()
  81         // All goroutines should block because picker is nil in bp.
  82         var finishedCount uint64
  83         for i := goroutineCount; i > 0; i-- {
  84                 go func() {
  85                         if tr, _, err := bp.pick(context.Background(), true, balancer.PickOptions{}); err != nil || tr != testT {
  86                                 t.Errorf("bp.pick returned non-nil error: %v", err)
  87                         }
  88                         atomic.AddUint64(&finishedCount, 1)
  89                 }()
  90         }
  91         time.Sleep(50 * time.Millisecond)
  92         if c := atomic.LoadUint64(&finishedCount); c != 0 {
  93                 t.Errorf("finished goroutines count: %v, want 0", c)
  94         }
  95         bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
  96 }
  97
  98 func TestBlockingPickNoSubAvailable(t *testing.T) {
  99         defer leakcheck.Check(t)
 100         bp := newPickerWrapper()
 101         var finishedCount uint64
 102         bp.updatePicker(&testingPicker{err: balancer.ErrNoSubConnAvailable, maxCalled: goroutineCount})
 103         // All goroutines should block because picker returns no sc available.
 104         for i := goroutineCount; i > 0; i-- {
 105                 go func() {
 106                         if tr, _, err := bp.pick(context.Background(), true, balancer.PickOptions{}); err != nil || tr != testT {
 107                                 t.Errorf("bp.pick returned non-nil error: %v", err)
 108                         }
 109                         atomic.AddUint64(&finishedCount, 1)
 110                 }()
 111         }
 112         time.Sleep(50 * time.Millisecond)
 113         if c := atomic.LoadUint64(&finishedCount); c != 0 {
 114                 t.Errorf("finished goroutines count: %v, want 0", c)
 115         }
 116         bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
 117 }
 118
 119 func TestBlockingPickTransientWaitforready(t *testing.T) {
 120         defer leakcheck.Check(t)
 121         bp := newPickerWrapper()
 122         bp.updatePicker(&testingPicker{err: balancer.ErrTransientFailure, maxCalled: goroutineCount})
 123         var finishedCount uint64
 124         // All goroutines should block because picker returns transientFailure and
 125         // picks are not failfast.
 126         for i := goroutineCount; i > 0; i-- {
 127                 go func() {
 128                         if tr, _, err := bp.pick(context.Background(), false, balancer.PickOptions{}); err != nil || tr != testT {
 129                                 t.Errorf("bp.pick returned non-nil error: %v", err)
 130                         }
 131                         atomic.AddUint64(&finishedCount, 1)
 132                 }()
 133         }
 134         time.Sleep(time.Millisecond)
 135         if c := atomic.LoadUint64(&finishedCount); c != 0 {
 136                 t.Errorf("finished goroutines count: %v, want 0", c)
 137         }
 138         bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
 139 }
 140
 141 func TestBlockingPickSCNotReady(t *testing.T) {
 142         defer leakcheck.Check(t)
 143         bp := newPickerWrapper()
 144         bp.updatePicker(&testingPicker{sc: testSCNotReady, maxCalled: goroutineCount})
 145         var finishedCount uint64
 146         // All goroutines should block because sc is not ready.
 147         for i := goroutineCount; i > 0; i-- {
 148                 go func() {
 149                         if tr, _, err := bp.pick(context.Background(), true, balancer.PickOptions{}); err != nil || tr != testT {
 150                                 t.Errorf("bp.pick returned non-nil error: %v", err)
 151                         }
 152                         atomic.AddUint64(&finishedCount, 1)
 153                 }()
 154         }
 155         time.Sleep(time.Millisecond)
 156         if c := atomic.LoadUint64(&finishedCount); c != 0 {
 157                 t.Errorf("finished goroutines count: %v, want 0", c)
 158         }
 159         bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})
 160 }