mirror of https://github.com/grpc/grpc-go.git
141 lines
3.9 KiB
Go
141 lines
3.9 KiB
Go
/*
|
|
*
|
|
* Copyright 2023 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.
|
|
*
|
|
*/
|
|
|
|
// Package xds contains various xds interop helpers for usage in interop tests.
|
|
package xds
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"google.golang.org/grpc/balancer"
|
|
"google.golang.org/grpc/balancer/roundrobin"
|
|
"google.golang.org/grpc/internal/pretty"
|
|
"google.golang.org/grpc/metadata"
|
|
"google.golang.org/grpc/serviceconfig"
|
|
)
|
|
|
|
func init() {
|
|
balancer.Register(rpcBehaviorBB{})
|
|
}
|
|
|
|
const name = "test.RpcBehaviorLoadBalancer"
|
|
|
|
type rpcBehaviorBB struct{}
|
|
|
|
func (rpcBehaviorBB) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
|
|
b := &rpcBehaviorLB{
|
|
ClientConn: cc,
|
|
}
|
|
// round_robin child to complete balancer tree with a usable leaf policy and
|
|
// have RPCs actually work.
|
|
builder := balancer.Get(roundrobin.Name)
|
|
if builder == nil {
|
|
// Shouldn't happen, defensive programming. Registered from import of
|
|
// roundrobin package.
|
|
return nil
|
|
}
|
|
rr := builder.Build(b, bOpts)
|
|
if rr == nil {
|
|
// Shouldn't happen, defensive programming.
|
|
return nil
|
|
}
|
|
b.Balancer = rr
|
|
return b
|
|
}
|
|
|
|
func (rpcBehaviorBB) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
|
|
lbCfg := &lbConfig{}
|
|
if err := json.Unmarshal(s, lbCfg); err != nil {
|
|
return nil, fmt.Errorf("rpc-behavior-lb: unable to marshal lbConfig: %s, error: %v", string(s), err)
|
|
}
|
|
return lbCfg, nil
|
|
|
|
}
|
|
|
|
func (rpcBehaviorBB) Name() string {
|
|
return name
|
|
}
|
|
|
|
type lbConfig struct {
|
|
serviceconfig.LoadBalancingConfig `json:"-"`
|
|
RPCBehavior string `json:"rpcBehavior,omitempty"`
|
|
}
|
|
|
|
// rpcBehaviorLB is a load balancer that wraps a round robin balancer and
|
|
// appends the rpc-behavior metadata field to any metadata in pick results based
|
|
// on what is specified in configuration.
|
|
type rpcBehaviorLB struct {
|
|
// embed a ClientConn to wrap only UpdateState() operation
|
|
balancer.ClientConn
|
|
// embed a Balancer to wrap only UpdateClientConnState() operation
|
|
balancer.Balancer
|
|
|
|
mu sync.Mutex
|
|
cfg *lbConfig
|
|
}
|
|
|
|
func (b *rpcBehaviorLB) UpdateClientConnState(s balancer.ClientConnState) error {
|
|
lbCfg, ok := s.BalancerConfig.(*lbConfig)
|
|
if !ok {
|
|
return fmt.Errorf("test.RpcBehaviorLoadBalancer:received config with unexpected type %T: %s", s.BalancerConfig, pretty.ToJSON(s.BalancerConfig))
|
|
}
|
|
b.mu.Lock()
|
|
b.cfg = lbCfg
|
|
b.mu.Unlock()
|
|
return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
|
|
ResolverState: s.ResolverState,
|
|
})
|
|
}
|
|
|
|
func (b *rpcBehaviorLB) UpdateState(state balancer.State) {
|
|
b.mu.Lock()
|
|
rpcBehavior := b.cfg.RPCBehavior
|
|
b.mu.Unlock()
|
|
|
|
b.ClientConn.UpdateState(balancer.State{
|
|
ConnectivityState: state.ConnectivityState,
|
|
Picker: newRPCBehaviorPicker(state.Picker, rpcBehavior),
|
|
})
|
|
}
|
|
|
|
// rpcBehaviorPicker wraps a picker and adds the rpc-behavior metadata field
|
|
// into the child pick result's metadata.
|
|
type rpcBehaviorPicker struct {
|
|
childPicker balancer.Picker
|
|
rpcBehavior string
|
|
}
|
|
|
|
// Pick appends the rpc-behavior metadata entry to the pick result of the child.
|
|
func (p *rpcBehaviorPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
|
|
pr, err := p.childPicker.Pick(info)
|
|
if err != nil {
|
|
return balancer.PickResult{}, err
|
|
}
|
|
pr.Metadata = metadata.Join(pr.Metadata, metadata.Pairs("rpc-behavior", p.rpcBehavior))
|
|
return pr, nil
|
|
}
|
|
|
|
func newRPCBehaviorPicker(childPicker balancer.Picker, rpcBehavior string) *rpcBehaviorPicker {
|
|
return &rpcBehaviorPicker{
|
|
childPicker: childPicker,
|
|
rpcBehavior: rpcBehavior,
|
|
}
|
|
}
|