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