1/* 2 * 3 * Copyright 2020 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 stub implements a balancer for testing purposes. 20package stub 21 22import "google.golang.org/grpc/balancer" 23 24// BalancerFuncs contains all balancer.Balancer functions with a preceding 25// *BalancerData parameter for passing additional instance information. Any 26// nil functions will never be called. 27type BalancerFuncs struct { 28 // Init is called after ClientConn and BuildOptions are set in 29 // BalancerData. It may be used to initialize BalancerData.Data. 30 Init func(*BalancerData) 31 32 UpdateClientConnState func(*BalancerData, balancer.ClientConnState) error 33 ResolverError func(*BalancerData, error) 34 UpdateSubConnState func(*BalancerData, balancer.SubConn, balancer.SubConnState) 35 Close func(*BalancerData) 36} 37 38// BalancerData contains data relevant to a stub balancer. 39type BalancerData struct { 40 // ClientConn is set by the builder. 41 ClientConn balancer.ClientConn 42 // BuildOptions is set by the builder. 43 BuildOptions balancer.BuildOptions 44 // Data may be used to store arbitrary user data. 45 Data interface{} 46} 47 48type bal struct { 49 bf BalancerFuncs 50 bd *BalancerData 51} 52 53func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error { 54 if b.bf.UpdateClientConnState != nil { 55 return b.bf.UpdateClientConnState(b.bd, c) 56 } 57 return nil 58} 59 60func (b *bal) ResolverError(e error) { 61 if b.bf.ResolverError != nil { 62 b.bf.ResolverError(b.bd, e) 63 } 64} 65 66func (b *bal) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) { 67 if b.bf.UpdateSubConnState != nil { 68 b.bf.UpdateSubConnState(b.bd, sc, scs) 69 } 70} 71 72func (b *bal) Close() { 73 if b.bf.Close != nil { 74 b.bf.Close(b.bd) 75 } 76} 77 78type bb struct { 79 name string 80 bf BalancerFuncs 81} 82 83func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer { 84 b := &bal{bf: bb.bf, bd: &BalancerData{ClientConn: cc, BuildOptions: opts}} 85 if b.bf.Init != nil { 86 b.bf.Init(b.bd) 87 } 88 return b 89} 90 91func (bb bb) Name() string { return bb.name } 92 93// Register registers a stub balancer builder which will call the provided 94// functions. The name used should be unique. 95func Register(name string, bf BalancerFuncs) { 96 balancer.Register(bb{name: name, bf: bf}) 97} 98