1/*
 2 *
 3 * Copyright 2022 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
19package balancer
20
21import "google.golang.org/grpc/connectivity"
22
23// ConnectivityStateEvaluator takes the connectivity states of multiple SubConns
24// and returns one aggregated connectivity state.
25//
26// It's not thread safe.
27type ConnectivityStateEvaluator struct {
28	numReady            uint64 // Number of addrConns in ready state.
29	numConnecting       uint64 // Number of addrConns in connecting state.
30	numTransientFailure uint64 // Number of addrConns in transient failure state.
31	numIdle             uint64 // Number of addrConns in idle state.
32}
33
34// RecordTransition records state change happening in subConn and based on that
35// it evaluates what aggregated state should be.
36//
37//   - If at least one SubConn in Ready, the aggregated state is Ready;
38//   - Else if at least one SubConn in Connecting, the aggregated state is Connecting;
39//   - Else if at least one SubConn is Idle, the aggregated state is Idle;
40//   - Else if at least one SubConn is TransientFailure (or there are no SubConns), the aggregated state is Transient Failure.
41//
42// Shutdown is not considered.
43func (cse *ConnectivityStateEvaluator) RecordTransition(oldState, newState connectivity.State) connectivity.State {
44	// Update counters.
45	for idx, state := range []connectivity.State{oldState, newState} {
46		updateVal := 2*uint64(idx) - 1 // -1 for oldState and +1 for new.
47		switch state {
48		case connectivity.Ready:
49			cse.numReady += updateVal
50		case connectivity.Connecting:
51			cse.numConnecting += updateVal
52		case connectivity.TransientFailure:
53			cse.numTransientFailure += updateVal
54		case connectivity.Idle:
55			cse.numIdle += updateVal
56		}
57	}
58	return cse.CurrentState()
59}
60
61// CurrentState returns the current aggregate conn state by evaluating the counters
62func (cse *ConnectivityStateEvaluator) CurrentState() connectivity.State {
63	// Evaluate.
64	if cse.numReady > 0 {
65		return connectivity.Ready
66	}
67	if cse.numConnecting > 0 {
68		return connectivity.Connecting
69	}
70	if cse.numIdle > 0 {
71		return connectivity.Idle
72	}
73	return connectivity.TransientFailure
74}