96 lines
2.4 KiB
Go
96 lines
2.4 KiB
Go
|
// Copyright 2018 The etcd 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 picker
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"go.uber.org/zap"
|
||
|
|
"go.uber.org/zap/zapcore"
|
||
|
|
"google.golang.org/grpc/balancer"
|
||
|
|
"google.golang.org/grpc/resolver"
|
||
|
|
)
|
||
|
|
|
||
|
|
// newRoundrobinBalanced returns a new roundrobin balanced picker.
|
||
|
|
func newRoundrobinBalanced(cfg Config) Picker {
|
||
|
|
scs := make([]balancer.SubConn, 0, len(cfg.SubConnToResolverAddress))
|
||
|
|
for sc := range cfg.SubConnToResolverAddress {
|
||
|
|
scs = append(scs, sc)
|
||
|
|
}
|
||
|
|
return &rrBalanced{
|
||
|
|
p: RoundrobinBalanced,
|
||
|
|
lg: cfg.Logger,
|
||
|
|
scs: scs,
|
||
|
|
scToAddr: cfg.SubConnToResolverAddress,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type rrBalanced struct {
|
||
|
|
p Policy
|
||
|
|
|
||
|
|
lg *zap.Logger
|
||
|
|
|
||
|
|
mu sync.RWMutex
|
||
|
|
next int
|
||
|
|
scs []balancer.SubConn
|
||
|
|
scToAddr map[balancer.SubConn]resolver.Address
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rb *rrBalanced) String() string { return rb.p.String() }
|
||
|
|
|
||
|
|
// Pick is called for every client request.
|
||
|
|
func (rb *rrBalanced) Pick(ctx context.Context, opts balancer.PickOptions) (balancer.SubConn, func(balancer.DoneInfo), error) {
|
||
|
|
rb.mu.RLock()
|
||
|
|
n := len(rb.scs)
|
||
|
|
rb.mu.RUnlock()
|
||
|
|
if n == 0 {
|
||
|
|
return nil, nil, balancer.ErrNoSubConnAvailable
|
||
|
|
}
|
||
|
|
|
||
|
|
rb.mu.Lock()
|
||
|
|
cur := rb.next
|
||
|
|
sc := rb.scs[cur]
|
||
|
|
picked := rb.scToAddr[sc].Addr
|
||
|
|
rb.next = (rb.next + 1) % len(rb.scs)
|
||
|
|
rb.mu.Unlock()
|
||
|
|
|
||
|
|
rb.lg.Debug(
|
||
|
|
"picked",
|
||
|
|
zap.String("picker", rb.p.String()),
|
||
|
|
zap.String("address", picked),
|
||
|
|
zap.Int("subconn-index", cur),
|
||
|
|
zap.Int("subconn-size", n),
|
||
|
|
)
|
||
|
|
|
||
|
|
doneFunc := func(info balancer.DoneInfo) {
|
||
|
|
// TODO: error handling?
|
||
|
|
fss := []zapcore.Field{
|
||
|
|
zap.Error(info.Err),
|
||
|
|
zap.String("picker", rb.p.String()),
|
||
|
|
zap.String("address", picked),
|
||
|
|
zap.Bool("success", info.Err == nil),
|
||
|
|
zap.Bool("bytes-sent", info.BytesSent),
|
||
|
|
zap.Bool("bytes-received", info.BytesReceived),
|
||
|
|
}
|
||
|
|
if info.Err == nil {
|
||
|
|
rb.lg.Debug("balancer done", fss...)
|
||
|
|
} else {
|
||
|
|
rb.lg.Warn("balancer failed", fss...)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return sc, doneFunc, nil
|
||
|
|
}
|