You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

68 lines
1.5 KiB

2 years ago
package selector
import (
"context"
"sync/atomic"
)
type DefaultSelector struct {
nodes atomic.Value
Balancer Balancer
NodeBuilder WeightedNodeBuilder
}
func (d *DefaultSelector) Apply(nodes []Node) {
newNodes := make([]WeightedNode, 0, len(nodes))
for _, v := range nodes {
newNodes = append(newNodes, d.NodeBuilder.Build(v))
}
d.nodes.Store(newNodes)
}
func (d *DefaultSelector) Select(ctx context.Context, opts ...SelectOption) (selected WeightedNode, di DoneFunc, err error) {
var options SelectOptions
for _, o := range opts {
o(&options)
}
nodes, ok := d.nodes.Load().([]WeightedNode)
if !ok {
return nil, nil, ErrNoAvailable
}
conditions := make([]WeightedNode, 0)
if len(options.NodeFilters) > 0 {
newNodes := make([]Node, 0)
for _, v := range nodes {
newNodes = append(newNodes, v)
}
for _, filter := range options.NodeFilters {
newNodes = filter(ctx, newNodes)
}
conditions = make([]WeightedNode, 0, len(newNodes))
for _, v := range newNodes {
conditions = append(conditions, v.(WeightedNode))
}
} else {
conditions = nodes
}
pick, doneFunc, err := d.Balancer.Pick(ctx, conditions)
if err != nil {
return nil, nil, err
}
if t, ok := FromPeerContext(ctx); ok {
t.Node = pick
}
return pick, doneFunc, nil
}
type DefaultBuilder struct {
Node WeightedNodeBuilder
Balancer BalancerBuilder
}
func (d *DefaultBuilder) Build() Selector {
return &DefaultSelector{
NodeBuilder: d.Node,
Balancer: d.Balancer.Build(),
}
}