route/vendor/google.golang.org/grpc/grpclb/grpclb_test.go

976 lines
27 KiB
Go
Raw Normal View History

2017-10-06 15:29:20 +00:00
/*
*
* Copyright 2016 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
2018-01-03 19:19:49 +00:00
//go:generate protoc --go_out=plugins=:$GOPATH grpc_lb_v1/messages/messages.proto
//go:generate protoc --go_out=plugins=grpc:$GOPATH grpc_lb_v1/service/service.proto
2017-10-06 15:29:20 +00:00
// Package grpclb_test is currently used only for grpclb testing.
package grpclb_test
import (
"errors"
"fmt"
"io"
"net"
2018-01-03 19:19:49 +00:00
"strconv"
2017-10-06 15:29:20 +00:00
"strings"
"sync"
"testing"
"time"
"github.com/golang/protobuf/proto"
"golang.org/x/net/context"
"google.golang.org/grpc"
2018-01-03 19:19:49 +00:00
"google.golang.org/grpc/balancer"
2017-10-06 15:29:20 +00:00
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials"
lbmpb "google.golang.org/grpc/grpclb/grpc_lb_v1/messages"
lbspb "google.golang.org/grpc/grpclb/grpc_lb_v1/service"
2018-01-03 19:19:49 +00:00
_ "google.golang.org/grpc/grpclog/glogger"
2017-10-06 15:29:20 +00:00
"google.golang.org/grpc/metadata"
2018-01-03 19:19:49 +00:00
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/manual"
"google.golang.org/grpc/status"
2017-10-06 15:29:20 +00:00
testpb "google.golang.org/grpc/test/grpc_testing"
2018-01-03 19:19:49 +00:00
"google.golang.org/grpc/test/leakcheck"
_ "google.golang.org/grpc/grpclog/glogger"
2017-10-06 15:29:20 +00:00
)
var (
2018-01-03 19:19:49 +00:00
lbServerName = "bar.com"
beServerName = "foo.com"
lbToken = "iamatoken"
2017-10-06 15:29:20 +00:00
// Resolver replaces localhost with fakeName in Next().
// Dialer replaces fakeName with localhost when dialing.
// This will test that custom dialer is passed from Dial to grpclb.
fakeName = "fake.Name"
)
type serverNameCheckCreds struct {
mu sync.Mutex
sn string
expected string
}
func (c *serverNameCheckCreds) ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
if _, err := io.WriteString(rawConn, c.sn); err != nil {
fmt.Printf("Failed to write the server name %s to the client %v", c.sn, err)
return nil, nil, err
}
return rawConn, nil, nil
}
func (c *serverNameCheckCreds) ClientHandshake(ctx context.Context, addr string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
c.mu.Lock()
defer c.mu.Unlock()
b := make([]byte, len(c.expected))
2018-01-03 19:19:49 +00:00
errCh := make(chan error, 1)
go func() {
_, err := rawConn.Read(b)
errCh <- err
}()
select {
case err := <-errCh:
if err != nil {
fmt.Printf("Failed to read the server name from the server %v", err)
return nil, nil, err
}
case <-ctx.Done():
return nil, nil, ctx.Err()
2017-10-06 15:29:20 +00:00
}
if c.expected != string(b) {
fmt.Printf("Read the server name %s want %s", string(b), c.expected)
return nil, nil, errors.New("received unexpected server name")
}
return rawConn, nil, nil
}
func (c *serverNameCheckCreds) Info() credentials.ProtocolInfo {
c.mu.Lock()
defer c.mu.Unlock()
return credentials.ProtocolInfo{}
}
func (c *serverNameCheckCreds) Clone() credentials.TransportCredentials {
c.mu.Lock()
defer c.mu.Unlock()
return &serverNameCheckCreds{
expected: c.expected,
}
}
func (c *serverNameCheckCreds) OverrideServerName(s string) error {
c.mu.Lock()
defer c.mu.Unlock()
c.expected = s
return nil
}
// fakeNameDialer replaces fakeName with localhost when dialing.
// This will test that custom dialer is passed from Dial to grpclb.
func fakeNameDialer(addr string, timeout time.Duration) (net.Conn, error) {
addr = strings.Replace(addr, fakeName, "localhost", 1)
return net.DialTimeout("tcp", addr, timeout)
}
type remoteBalancer struct {
2018-01-03 19:19:49 +00:00
sls chan *lbmpb.ServerList
2017-10-06 15:29:20 +00:00
statsDura time.Duration
done chan struct{}
mu sync.Mutex
stats lbmpb.ClientStats
}
2018-01-03 19:19:49 +00:00
func newRemoteBalancer(intervals []time.Duration) *remoteBalancer {
2017-10-06 15:29:20 +00:00
return &remoteBalancer{
2018-01-03 19:19:49 +00:00
sls: make(chan *lbmpb.ServerList, 1),
done: make(chan struct{}),
2017-10-06 15:29:20 +00:00
}
}
func (b *remoteBalancer) stop() {
2018-01-03 19:19:49 +00:00
close(b.sls)
2017-10-06 15:29:20 +00:00
close(b.done)
}
func (b *remoteBalancer) BalanceLoad(stream lbspb.LoadBalancer_BalanceLoadServer) error {
req, err := stream.Recv()
if err != nil {
return err
}
initReq := req.GetInitialRequest()
2018-01-03 19:19:49 +00:00
if initReq.Name != beServerName {
return status.Errorf(codes.InvalidArgument, "invalid service name: %v", initReq.Name)
2017-10-06 15:29:20 +00:00
}
resp := &lbmpb.LoadBalanceResponse{
LoadBalanceResponseType: &lbmpb.LoadBalanceResponse_InitialResponse{
InitialResponse: &lbmpb.InitialLoadBalanceResponse{
ClientStatsReportInterval: &lbmpb.Duration{
Seconds: int64(b.statsDura.Seconds()),
Nanos: int32(b.statsDura.Nanoseconds() - int64(b.statsDura.Seconds())*1e9),
},
},
},
}
if err := stream.Send(resp); err != nil {
return err
}
go func() {
for {
var (
req *lbmpb.LoadBalanceRequest
err error
)
if req, err = stream.Recv(); err != nil {
return
}
b.mu.Lock()
b.stats.NumCallsStarted += req.GetClientStats().NumCallsStarted
b.stats.NumCallsFinished += req.GetClientStats().NumCallsFinished
b.stats.NumCallsFinishedWithDropForRateLimiting += req.GetClientStats().NumCallsFinishedWithDropForRateLimiting
b.stats.NumCallsFinishedWithDropForLoadBalancing += req.GetClientStats().NumCallsFinishedWithDropForLoadBalancing
b.stats.NumCallsFinishedWithClientFailedToSend += req.GetClientStats().NumCallsFinishedWithClientFailedToSend
b.stats.NumCallsFinishedKnownReceived += req.GetClientStats().NumCallsFinishedKnownReceived
b.mu.Unlock()
}
}()
2018-01-03 19:19:49 +00:00
for v := range b.sls {
2017-10-06 15:29:20 +00:00
resp = &lbmpb.LoadBalanceResponse{
LoadBalanceResponseType: &lbmpb.LoadBalanceResponse_ServerList{
ServerList: v,
},
}
if err := stream.Send(resp); err != nil {
return err
}
}
<-b.done
return nil
}
type testServer struct {
testpb.TestServiceServer
2018-01-03 19:19:49 +00:00
addr string
fallback bool
2017-10-06 15:29:20 +00:00
}
const testmdkey = "testmd"
func (s *testServer) EmptyCall(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
2018-01-03 19:19:49 +00:00
return nil, status.Error(codes.Internal, "failed to receive metadata")
2017-10-06 15:29:20 +00:00
}
2018-01-03 19:19:49 +00:00
if !s.fallback && (md == nil || md["lb-token"][0] != lbToken) {
return nil, status.Errorf(codes.Internal, "received unexpected metadata: %v", md)
2017-10-06 15:29:20 +00:00
}
grpc.SetTrailer(ctx, metadata.Pairs(testmdkey, s.addr))
return &testpb.Empty{}, nil
}
func (s *testServer) FullDuplexCall(stream testpb.TestService_FullDuplexCallServer) error {
return nil
}
2018-01-03 19:19:49 +00:00
func startBackends(sn string, fallback bool, lis ...net.Listener) (servers []*grpc.Server) {
2017-10-06 15:29:20 +00:00
for _, l := range lis {
creds := &serverNameCheckCreds{
sn: sn,
}
s := grpc.NewServer(grpc.Creds(creds))
2018-01-03 19:19:49 +00:00
testpb.RegisterTestServiceServer(s, &testServer{addr: l.Addr().String(), fallback: fallback})
2017-10-06 15:29:20 +00:00
servers = append(servers, s)
go func(s *grpc.Server, l net.Listener) {
s.Serve(l)
}(s, l)
}
return
}
func stopBackends(servers []*grpc.Server) {
for _, s := range servers {
s.Stop()
}
}
type testServers struct {
lbAddr string
ls *remoteBalancer
lb *grpc.Server
beIPs []net.IP
bePorts []int
}
func newLoadBalancer(numberOfBackends int) (tss *testServers, cleanup func(), err error) {
var (
beListeners []net.Listener
ls *remoteBalancer
lb *grpc.Server
beIPs []net.IP
bePorts []int
)
for i := 0; i < numberOfBackends; i++ {
// Start a backend.
beLis, e := net.Listen("tcp", "localhost:0")
if e != nil {
err = fmt.Errorf("Failed to listen %v", err)
return
}
beIPs = append(beIPs, beLis.Addr().(*net.TCPAddr).IP)
bePorts = append(bePorts, beLis.Addr().(*net.TCPAddr).Port)
beListeners = append(beListeners, beLis)
}
2018-01-03 19:19:49 +00:00
backends := startBackends(beServerName, false, beListeners...)
2017-10-06 15:29:20 +00:00
// Start a load balancer.
lbLis, err := net.Listen("tcp", "localhost:0")
if err != nil {
err = fmt.Errorf("Failed to create the listener for the load balancer %v", err)
return
}
lbCreds := &serverNameCheckCreds{
2018-01-03 19:19:49 +00:00
sn: lbServerName,
2017-10-06 15:29:20 +00:00
}
lb = grpc.NewServer(grpc.Creds(lbCreds))
if err != nil {
err = fmt.Errorf("Failed to generate the port number %v", err)
return
}
2018-01-03 19:19:49 +00:00
ls = newRemoteBalancer(nil)
2017-10-06 15:29:20 +00:00
lbspb.RegisterLoadBalancerServer(lb, ls)
go func() {
lb.Serve(lbLis)
}()
tss = &testServers{
2018-01-03 19:19:49 +00:00
lbAddr: fakeName + ":" + strconv.Itoa(lbLis.Addr().(*net.TCPAddr).Port),
2017-10-06 15:29:20 +00:00
ls: ls,
lb: lb,
beIPs: beIPs,
bePorts: bePorts,
}
cleanup = func() {
defer stopBackends(backends)
defer func() {
ls.stop()
lb.Stop()
}()
}
return
}
func TestGRPCLB(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
2017-10-06 15:29:20 +00:00
tss, cleanup, err := newLoadBalancer(1)
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
be := &lbmpb.Server{
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
LoadBalanceToken: lbToken,
}
var bes []*lbmpb.Server
bes = append(bes, be)
sl := &lbmpb.ServerList{
Servers: bes,
}
2018-01-03 19:19:49 +00:00
tss.ls.sls <- sl
2017-10-06 15:29:20 +00:00
creds := serverNameCheckCreds{
2018-01-03 19:19:49 +00:00
expected: beServerName,
2017-10-06 15:29:20 +00:00
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
2018-01-03 19:19:49 +00:00
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
2018-01-03 19:19:49 +00:00
defer cc.Close()
2017-10-06 15:29:20 +00:00
testC := testpb.NewTestServiceClient(cc)
2018-01-03 19:19:49 +00:00
r.NewAddress([]resolver.Address{{
Addr: tss.lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}})
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false)); err != nil {
2017-10-06 15:29:20 +00:00
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
}
2018-01-03 19:19:49 +00:00
// The remote balancer sends response with duplicates to grpclb client.
func TestGRPCLBWeighted(t *testing.T) {
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
2017-10-06 15:29:20 +00:00
tss, cleanup, err := newLoadBalancer(2)
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
2018-01-03 19:19:49 +00:00
beServers := []*lbmpb.Server{{
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
LoadBalanceToken: lbToken,
}, {
IpAddress: tss.beIPs[1],
Port: int32(tss.bePorts[1]),
LoadBalanceToken: lbToken,
2017-10-06 15:29:20 +00:00
}}
2018-01-03 19:19:49 +00:00
portsToIndex := make(map[int]int)
for i := range beServers {
portsToIndex[tss.bePorts[i]] = i
}
2017-10-06 15:29:20 +00:00
creds := serverNameCheckCreds{
2018-01-03 19:19:49 +00:00
expected: beServerName,
2017-10-06 15:29:20 +00:00
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
2018-01-03 19:19:49 +00:00
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
2018-01-03 19:19:49 +00:00
defer cc.Close()
2017-10-06 15:29:20 +00:00
testC := testpb.NewTestServiceClient(cc)
2018-01-03 19:19:49 +00:00
r.NewAddress([]resolver.Address{{
Addr: tss.lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}})
sequences := []string{"00101", "00011"}
for _, seq := range sequences {
var (
bes []*lbmpb.Server
p peer.Peer
result string
)
for _, s := range seq {
bes = append(bes, beServers[s-'0'])
}
tss.ls.sls <- &lbmpb.ServerList{Servers: bes}
for i := 0; i < 1000; i++ {
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false), grpc.Peer(&p)); err != nil {
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
result += strconv.Itoa(portsToIndex[p.Addr.(*net.TCPAddr).Port])
2017-10-06 15:29:20 +00:00
}
2018-01-03 19:19:49 +00:00
// The generated result will be in format of "0010100101".
if !strings.Contains(result, strings.Repeat(seq, 2)) {
t.Errorf("got result sequence %q, want patten %q", result, seq)
2017-10-06 15:29:20 +00:00
}
}
}
2018-01-03 19:19:49 +00:00
func TestDropRequest(t *testing.T) {
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
2017-10-06 15:29:20 +00:00
tss, cleanup, err := newLoadBalancer(1)
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
2018-01-03 19:19:49 +00:00
tss.ls.sls <- &lbmpb.ServerList{
Servers: []*lbmpb.Server{{
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
LoadBalanceToken: lbToken,
DropForLoadBalancing: false,
}, {
DropForLoadBalancing: true,
}},
2017-10-06 15:29:20 +00:00
}
creds := serverNameCheckCreds{
2018-01-03 19:19:49 +00:00
expected: beServerName,
2017-10-06 15:29:20 +00:00
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
2018-01-03 19:19:49 +00:00
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
2018-01-03 19:19:49 +00:00
defer cc.Close()
2017-10-06 15:29:20 +00:00
testC := testpb.NewTestServiceClient(cc)
2018-01-03 19:19:49 +00:00
r.NewAddress([]resolver.Address{{
Addr: tss.lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}})
// The 1st, non-fail-fast RPC should succeed. This ensures both server
// connections are made, because the first one has DropForLoadBalancing set to true.
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false)); err != nil {
t.Fatalf("%v.SayHello(_, _) = _, %v, want _, <nil>", testC, err)
}
for _, failfast := range []bool{true, false} {
for i := 0; i < 3; i++ {
// Even RPCs should fail, because the 2st backend has
// DropForLoadBalancing set to true.
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(failfast)); status.Code(err) != codes.Unavailable {
t.Errorf("%v.EmptyCall(_, _) = _, %v, want _, %s", testC, err, codes.Unavailable)
}
// Odd RPCs should succeed since they choose the non-drop-request
// backend according to the round robin policy.
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(failfast)); err != nil {
t.Errorf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
}
2017-10-06 15:29:20 +00:00
}
}
// When the balancer in use disconnects, grpclb should connect to the next address from resolved balancer address list.
func TestBalancerDisconnects(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
2017-10-06 15:29:20 +00:00
var (
2018-01-03 19:19:49 +00:00
tests []*testServers
lbs []*grpc.Server
2017-10-06 15:29:20 +00:00
)
2018-01-03 19:19:49 +00:00
for i := 0; i < 2; i++ {
2017-10-06 15:29:20 +00:00
tss, cleanup, err := newLoadBalancer(1)
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
be := &lbmpb.Server{
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
LoadBalanceToken: lbToken,
}
var bes []*lbmpb.Server
bes = append(bes, be)
sl := &lbmpb.ServerList{
Servers: bes,
}
2018-01-03 19:19:49 +00:00
tss.ls.sls <- sl
2017-10-06 15:29:20 +00:00
2018-01-03 19:19:49 +00:00
tests = append(tests, tss)
2017-10-06 15:29:20 +00:00
lbs = append(lbs, tss.lb)
}
creds := serverNameCheckCreds{
2018-01-03 19:19:49 +00:00
expected: beServerName,
2017-10-06 15:29:20 +00:00
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
2018-01-03 19:19:49 +00:00
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
2018-01-03 19:19:49 +00:00
defer cc.Close()
2017-10-06 15:29:20 +00:00
testC := testpb.NewTestServiceClient(cc)
2018-01-03 19:19:49 +00:00
r.NewAddress([]resolver.Address{{
Addr: tests[0].lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}, {
Addr: tests[1].lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}})
var p peer.Peer
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false), grpc.Peer(&p)); err != nil {
2017-10-06 15:29:20 +00:00
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
2018-01-03 19:19:49 +00:00
if p.Addr.(*net.TCPAddr).Port != tests[0].bePorts[0] {
t.Fatalf("got peer: %v, want peer port: %v", p.Addr, tests[0].bePorts[0])
}
2017-10-06 15:29:20 +00:00
lbs[0].Stop()
2018-01-03 19:19:49 +00:00
// Stop balancer[0], balancer[1] should be used by grpclb.
// Check peer address to see if that happened.
for i := 0; i < 1000; i++ {
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false), grpc.Peer(&p)); err != nil {
2017-10-06 15:29:20 +00:00
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
2018-01-03 19:19:49 +00:00
}
if p.Addr.(*net.TCPAddr).Port == tests[1].bePorts[0] {
return
}
time.Sleep(time.Millisecond)
}
t.Fatalf("No RPC sent to second backend after 1 second")
}
type customGRPCLBBuilder struct {
balancer.Builder
name string
}
func (b *customGRPCLBBuilder) Name() string {
return b.name
}
const grpclbCustomFallbackName = "grpclb_with_custom_fallback_timeout"
func init() {
balancer.Register(&customGRPCLBBuilder{
Builder: grpc.NewLBBuilderWithFallbackTimeout(100 * time.Millisecond),
name: grpclbCustomFallbackName,
2017-10-06 15:29:20 +00:00
})
2018-01-03 19:19:49 +00:00
}
func TestFallback(t *testing.T) {
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
tss, cleanup, err := newLoadBalancer(1)
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
// Start a standalone backend.
beLis, err := net.Listen("tcp", "localhost:0")
if err != nil {
t.Fatalf("Failed to listen %v", err)
}
defer beLis.Close()
standaloneBEs := startBackends(beServerName, true, beLis)
defer stopBackends(standaloneBEs)
be := &lbmpb.Server{
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
LoadBalanceToken: lbToken,
}
var bes []*lbmpb.Server
bes = append(bes, be)
sl := &lbmpb.ServerList{
Servers: bes,
}
tss.ls.sls <- sl
creds := serverNameCheckCreds{
expected: beServerName,
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithBalancerName(grpclbCustomFallbackName),
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
defer cc.Close()
testC := testpb.NewTestServiceClient(cc)
r.NewAddress([]resolver.Address{{
Addr: "",
Type: resolver.GRPCLB,
ServerName: lbServerName,
}, {
Addr: beLis.Addr().String(),
Type: resolver.Backend,
ServerName: beServerName,
}})
var p peer.Peer
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false), grpc.Peer(&p)); err != nil {
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
if p.Addr.String() != beLis.Addr().String() {
t.Fatalf("got peer: %v, want peer: %v", p.Addr, beLis.Addr())
}
r.NewAddress([]resolver.Address{{
Addr: tss.lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}, {
Addr: beLis.Addr().String(),
Type: resolver.Backend,
ServerName: beServerName,
}})
for i := 0; i < 1000; i++ {
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false), grpc.Peer(&p)); err != nil {
2017-10-06 15:29:20 +00:00
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
2018-01-03 19:19:49 +00:00
if p.Addr.(*net.TCPAddr).Port == tss.bePorts[0] {
return
}
time.Sleep(time.Millisecond)
2017-10-06 15:29:20 +00:00
}
2018-01-03 19:19:49 +00:00
t.Fatalf("No RPC sent to backend behind remote balancer after 1 second")
2017-10-06 15:29:20 +00:00
}
type failPreRPCCred struct{}
func (failPreRPCCred) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
2018-01-03 19:19:49 +00:00
if strings.Contains(uri[0], failtosendURI) {
2017-10-06 15:29:20 +00:00
return nil, fmt.Errorf("rpc should fail to send")
}
return nil, nil
}
func (failPreRPCCred) RequireTransportSecurity() bool {
return false
}
func checkStats(stats *lbmpb.ClientStats, expected *lbmpb.ClientStats) error {
if !proto.Equal(stats, expected) {
return fmt.Errorf("stats not equal: got %+v, want %+v", stats, expected)
}
return nil
}
func runAndGetStats(t *testing.T, dropForLoadBalancing, dropForRateLimiting bool, runRPCs func(*grpc.ClientConn)) lbmpb.ClientStats {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
r, cleanup := manual.GenerateAndRegisterManualResolver()
defer cleanup()
tss, cleanup, err := newLoadBalancer(1)
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("failed to create new load balancer: %v", err)
}
defer cleanup()
2018-01-03 19:19:49 +00:00
tss.ls.sls <- &lbmpb.ServerList{
2017-10-06 15:29:20 +00:00
Servers: []*lbmpb.Server{{
2018-01-03 19:19:49 +00:00
IpAddress: tss.beIPs[0],
Port: int32(tss.bePorts[0]),
2017-10-06 15:29:20 +00:00
LoadBalanceToken: lbToken,
DropForLoadBalancing: dropForLoadBalancing,
DropForRateLimiting: dropForRateLimiting,
}},
2018-01-03 19:19:49 +00:00
}
2017-10-06 15:29:20 +00:00
tss.ls.statsDura = 100 * time.Millisecond
2018-01-03 19:19:49 +00:00
creds := serverNameCheckCreds{expected: beServerName}
2017-10-06 15:29:20 +00:00
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
2018-01-03 19:19:49 +00:00
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
grpc.WithTransportCredentials(&creds),
grpc.WithPerRPCCredentials(failPreRPCCred{}),
grpc.WithDialer(fakeNameDialer))
2017-10-06 15:29:20 +00:00
if err != nil {
t.Fatalf("Failed to dial to the backend %v", err)
}
defer cc.Close()
2018-01-03 19:19:49 +00:00
r.NewAddress([]resolver.Address{{
Addr: tss.lbAddr,
Type: resolver.GRPCLB,
ServerName: lbServerName,
}})
2017-10-06 15:29:20 +00:00
runRPCs(cc)
time.Sleep(1 * time.Second)
tss.ls.mu.Lock()
stats := tss.ls.stats
tss.ls.mu.Unlock()
return stats
}
2018-01-03 19:19:49 +00:00
const (
countRPC = 40
failtosendURI = "failtosend"
dropErrDesc = "request dropped by grpclb"
)
2017-10-06 15:29:20 +00:00
func TestGRPCLBStatsUnarySuccess(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
stats := runAndGetStats(t, false, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
// The first non-failfast RPC succeeds, all connections are up.
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false)); err != nil {
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
for i := 0; i < countRPC-1; i++ {
testC.EmptyCall(context.Background(), &testpb.Empty{})
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC),
NumCallsFinished: int64(countRPC),
NumCallsFinishedKnownReceived: int64(countRPC),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsUnaryDropLoadBalancing(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
c := 0
stats := runAndGetStats(t, true, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
for {
c++
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}); err != nil {
2018-01-03 19:19:49 +00:00
if strings.Contains(err.Error(), dropErrDesc) {
2017-10-06 15:29:20 +00:00
break
}
}
}
for i := 0; i < countRPC; i++ {
testC.EmptyCall(context.Background(), &testpb.Empty{})
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC + c),
NumCallsFinished: int64(countRPC + c),
NumCallsFinishedWithDropForLoadBalancing: int64(countRPC + 1),
NumCallsFinishedWithClientFailedToSend: int64(c - 1),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsUnaryDropRateLimiting(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
c := 0
stats := runAndGetStats(t, false, true, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
for {
c++
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}); err != nil {
2018-01-03 19:19:49 +00:00
if strings.Contains(err.Error(), dropErrDesc) {
2017-10-06 15:29:20 +00:00
break
}
}
}
for i := 0; i < countRPC; i++ {
testC.EmptyCall(context.Background(), &testpb.Empty{})
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC + c),
NumCallsFinished: int64(countRPC + c),
NumCallsFinishedWithDropForRateLimiting: int64(countRPC + 1),
NumCallsFinishedWithClientFailedToSend: int64(c - 1),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsUnaryFailedToSend(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
stats := runAndGetStats(t, false, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
// The first non-failfast RPC succeeds, all connections are up.
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false)); err != nil {
t.Fatalf("%v.EmptyCall(_, _) = _, %v, want _, <nil>", testC, err)
}
for i := 0; i < countRPC-1; i++ {
2018-01-03 19:19:49 +00:00
grpc.Invoke(context.Background(), failtosendURI, &testpb.Empty{}, nil, cc)
2017-10-06 15:29:20 +00:00
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC),
NumCallsFinished: int64(countRPC),
NumCallsFinishedWithClientFailedToSend: int64(countRPC - 1),
NumCallsFinishedKnownReceived: 1,
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsStreamingSuccess(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
stats := runAndGetStats(t, false, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
// The first non-failfast RPC succeeds, all connections are up.
stream, err := testC.FullDuplexCall(context.Background(), grpc.FailFast(false))
if err != nil {
t.Fatalf("%v.FullDuplexCall(_, _) = _, %v, want _, <nil>", testC, err)
}
for {
if _, err = stream.Recv(); err == io.EOF {
break
}
}
for i := 0; i < countRPC-1; i++ {
stream, err = testC.FullDuplexCall(context.Background())
if err == nil {
// Wait for stream to end if err is nil.
for {
if _, err = stream.Recv(); err == io.EOF {
break
}
}
}
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC),
NumCallsFinished: int64(countRPC),
NumCallsFinishedKnownReceived: int64(countRPC),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsStreamingDropLoadBalancing(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
c := 0
stats := runAndGetStats(t, true, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
for {
c++
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}); err != nil {
2018-01-03 19:19:49 +00:00
if strings.Contains(err.Error(), dropErrDesc) {
2017-10-06 15:29:20 +00:00
break
}
}
}
for i := 0; i < countRPC; i++ {
testC.FullDuplexCall(context.Background())
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC + c),
NumCallsFinished: int64(countRPC + c),
NumCallsFinishedWithDropForLoadBalancing: int64(countRPC + 1),
NumCallsFinishedWithClientFailedToSend: int64(c - 1),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsStreamingDropRateLimiting(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
c := 0
stats := runAndGetStats(t, false, true, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
for {
c++
if _, err := testC.EmptyCall(context.Background(), &testpb.Empty{}); err != nil {
2018-01-03 19:19:49 +00:00
if strings.Contains(err.Error(), dropErrDesc) {
2017-10-06 15:29:20 +00:00
break
}
}
}
for i := 0; i < countRPC; i++ {
testC.FullDuplexCall(context.Background())
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC + c),
NumCallsFinished: int64(countRPC + c),
NumCallsFinishedWithDropForRateLimiting: int64(countRPC + 1),
NumCallsFinishedWithClientFailedToSend: int64(c - 1),
}); err != nil {
t.Fatal(err)
}
}
func TestGRPCLBStatsStreamingFailedToSend(t *testing.T) {
2018-01-03 19:19:49 +00:00
defer leakcheck.Check(t)
2017-10-06 15:29:20 +00:00
stats := runAndGetStats(t, false, false, func(cc *grpc.ClientConn) {
testC := testpb.NewTestServiceClient(cc)
// The first non-failfast RPC succeeds, all connections are up.
stream, err := testC.FullDuplexCall(context.Background(), grpc.FailFast(false))
if err != nil {
t.Fatalf("%v.FullDuplexCall(_, _) = _, %v, want _, <nil>", testC, err)
}
for {
if _, err = stream.Recv(); err == io.EOF {
break
}
}
for i := 0; i < countRPC-1; i++ {
2018-01-03 19:19:49 +00:00
grpc.NewClientStream(context.Background(), &grpc.StreamDesc{}, cc, failtosendURI)
2017-10-06 15:29:20 +00:00
}
})
if err := checkStats(&stats, &lbmpb.ClientStats{
NumCallsStarted: int64(countRPC),
NumCallsFinished: int64(countRPC),
NumCallsFinishedWithClientFailedToSend: int64(countRPC - 1),
NumCallsFinishedKnownReceived: 1,
}); err != nil {
t.Fatal(err)
}
}