Files
gf/contrib/rpc/grpcx/internal/balancer/balancer_picker.go
2023-03-08 14:12:51 +08:00

65 lines
2.1 KiB
Go

// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.
package balancer
import (
"google.golang.org/grpc/balancer"
"github.com/gogf/gf/v2/net/gsel"
)
// Picker implements grpc balancer.Picker,
// which is used by gRPC to pick a SubConn to send an RPC.
// Balancer is expected to generate a new picker from its snapshot every time its
// internal state has changed.
//
// The pickers used by gRPC can be updated by ClientConn.UpdateState().
type Picker struct {
selector gsel.Selector
}
// Pick returns the connection to use for this RPC and related information.
//
// Pick should not block. If the balancer needs to do I/O or any blocking
// or time-consuming work to service this call, it should return
// ErrNoSubConnAvailable, and the Pick call will be repeated by gRPC when
// the Picker is updated (using ClientConn.UpdateState).
//
// If an error is returned:
//
// - If the error is ErrNoSubConnAvailable, gRPC will block until a new
// Picker is provided by the balancer (using ClientConn.UpdateState).
//
// - If the error is a status error (implemented by the grpc/status
// package), gRPC will terminate the RPC with the code and message
// provided.
//
// - For all other errors, wait for ready RPCs will wait, but non-wait for
// ready RPCs will be terminated with this error's Error() string and
// status code Unavailable.
func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
node, done, err := p.selector.Pick(info.Ctx)
if err != nil {
return balancer.PickResult{}, err
}
return balancer.PickResult{
SubConn: node.(*Node).conn,
Done: func(di balancer.DoneInfo) {
if done == nil {
return
}
done(info.Ctx, gsel.DoneInfo{
Err: di.Err,
Trailer: di.Trailer,
BytesSent: di.BytesSent,
BytesReceived: di.BytesReceived,
ServerLoad: di.ServerLoad,
})
},
}, nil
}