vendor/google.golang.org/grpc/balancer_switching_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         "math"
  24         "testing"
  25         "time"
  26
  27         "golang.org/x/net/context"
  28         _ "google.golang.org/grpc/grpclog/glogger"
  29         "google.golang.org/grpc/resolver"
  30         "google.golang.org/grpc/resolver/manual"
  31         "google.golang.org/grpc/test/leakcheck"
  32 )
  33
  34 func checkPickFirst(cc *ClientConn, servers []*server) error {
  35         var (
  36                 req   = "port"
  37                 reply string
  38                 err   error
  39         )
  40         connected := false
  41         for i := 0; i < 1000; i++ {
  42                 if err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc); ErrorDesc(err) == servers[0].port {
  43                         if connected {
  44                                 // connected is set to false if peer is not server[0]. So if
  45                                 // connected is true here, this is the second time we saw
  46                                 // server[0] in a row. Break because pickfirst is in effect.
  47                                 break
  48                         }
  49                         connected = true
  50                 } else {
  51                         connected = false
  52                 }
  53                 time.Sleep(time.Millisecond)
  54         }
  55         if !connected {
  56                 return fmt.Errorf("pickfirst is not in effect after 1 second, EmptyCall() = _, %v, want _, %v", err, servers[0].port)
  57         }
  58         // The following RPCs should all succeed with the first server.
  59         for i := 0; i < 3; i++ {
  60                 err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc)
  61                 if ErrorDesc(err) != servers[0].port {
  62                         return fmt.Errorf("Index %d: want peer %v, got peer %v", i, servers[0].port, err)
  63                 }
  64         }
  65         return nil
  66 }
  67
  68 func checkRoundRobin(cc *ClientConn, servers []*server) error {
  69         var (
  70                 req   = "port"
  71                 reply string
  72                 err   error
  73         )
  74
  75         // Make sure connections to all servers are up.
  76         for i := 0; i < 2; i++ {
  77                 // Do this check twice, otherwise the first RPC's transport may still be
  78                 // picked by the closing pickfirst balancer, and the test becomes flaky.
  79                 for _, s := range servers {
  80                         var up bool
  81                         for i := 0; i < 1000; i++ {
  82                                 if err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc); ErrorDesc(err) == s.port {
  83                                         up = true
  84                                         break
  85                                 }
  86                                 time.Sleep(time.Millisecond)
  87                         }
  88                         if !up {
  89                                 return fmt.Errorf("server %v is not up within 1 second", s.port)
  90                         }
  91                 }
  92         }
  93
  94         serverCount := len(servers)
  95         for i := 0; i < 3*serverCount; i++ {
  96                 err = Invoke(context.Background(), "/foo/bar", &req, &reply, cc)
  97                 if ErrorDesc(err) != servers[i%serverCount].port {
  98                         return fmt.Errorf("Index %d: want peer %v, got peer %v", i, servers[i%serverCount].port, err)
  99                 }
 100         }
 101         return nil
 102 }
 103
 104 func TestSwitchBalancer(t *testing.T) {
 105         defer leakcheck.Check(t)
 106         r, rcleanup := manual.GenerateAndRegisterManualResolver()
 107         defer rcleanup()
 108
 109         numServers := 2
 110         servers, _, scleanup := startServers(t, numServers, math.MaxInt32)
 111         defer scleanup()
 112
 113         cc, err := Dial(r.Scheme()+":///test.server", WithInsecure(), WithCodec(testCodec{}))
 114         if err != nil {
 115                 t.Fatalf("failed to dial: %v", err)
 116         }
 117         defer cc.Close()
 118         r.NewAddress([]resolver.Address{{Addr: servers[0].addr}, {Addr: servers[1].addr}})
 119         // The default balancer is pickfirst.
 120         if err := checkPickFirst(cc, servers); err != nil {
 121                 t.Fatalf("check pickfirst returned non-nil error: %v", err)
 122         }
 123         // Switch to roundrobin.
 124         cc.switchBalancer("roundrobin")
 125         if err := checkRoundRobin(cc, servers); err != nil {
 126                 t.Fatalf("check roundrobin returned non-nil error: %v", err)
 127         }
 128         // Switch to pickfirst.
 129         cc.switchBalancer("pickfirst")
 130         if err := checkPickFirst(cc, servers); err != nil {
 131                 t.Fatalf("check pickfirst returned non-nil error: %v", err)
 132         }
 133 }