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	ExitIdle              func(*BalancerData)
37}
38
39// BalancerData contains data relevant to a stub balancer.
40type BalancerData struct {
41	// ClientConn is set by the builder.
42	ClientConn balancer.ClientConn
43	// BuildOptions is set by the builder.
44	BuildOptions balancer.BuildOptions
45	// Data may be used to store arbitrary user data.
46	Data interface{}
47}
48
49type bal struct {
50	bf BalancerFuncs
51	bd *BalancerData
52}
53
54func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error {
55	if b.bf.UpdateClientConnState != nil {
56		return b.bf.UpdateClientConnState(b.bd, c)
57	}
58	return nil
59}
60
61func (b *bal) ResolverError(e error) {
62	if b.bf.ResolverError != nil {
63		b.bf.ResolverError(b.bd, e)
64	}
65}
66
67func (b *bal) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
68	if b.bf.UpdateSubConnState != nil {
69		b.bf.UpdateSubConnState(b.bd, sc, scs)
70	}
71}
72
73func (b *bal) Close() {
74	if b.bf.Close != nil {
75		b.bf.Close(b.bd)
76	}
77}
78
79func (b *bal) ExitIdle() {
80	if b.bf.ExitIdle != nil {
81		b.bf.ExitIdle(b.bd)
82	}
83}
84
85type bb struct {
86	name string
87	bf   BalancerFuncs
88}
89
90func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
91	b := &bal{bf: bb.bf, bd: &BalancerData{ClientConn: cc, BuildOptions: opts}}
92	if b.bf.Init != nil {
93		b.bf.Init(b.bd)
94	}
95	return b
96}
97
98func (bb bb) Name() string { return bb.name }
99
100// Register registers a stub balancer builder which will call the provided
101// functions.  The name used should be unique.
102func Register(name string, bf BalancerFuncs) {
103	balancer.Register(bb{name: name, bf: bf})
104}
105