kit/selector/default_selector.go

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(),
	}
}