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.
67 lines
1.5 KiB
67 lines
1.5 KiB
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(), |
|
} |
|
}
|
|
|