kit/selector/random/random.go

package random

import (
	"context"
	"git.diulo.com/mogfee/kit/selector"
	"git.diulo.com/mogfee/kit/selector/node"
	"sync/atomic"
)

var _ selector.Balancer = (*Balancer)(nil)

func New() selector.Selector {
	return NewBuilder().Build()
}

type Balancer struct {
	index atomic.Int64
}

func (r *Balancer) Pick(ctx context.Context, nodes []selector.WeightedNode) (selector.WeightedNode, selector.DoneFunc, error) {
	if len(nodes) == 0 {
		return nil, nil, selector.ErrNoAvailable
	}
	//index := rand.Intn(len(nodes))
	index := r.index.Load()
	if int(index) > len(nodes) {
		index = 0
	}
	r.index.Add(1)
	selected := nodes[index]
	d := selected.Pick()
	return selected, d, nil
}

func NewBuilder() selector.Builder {
	return &selector.DefaultBuilder{
		Node:     &node.Builder{},
		Balancer: &Builder{},
	}
}

type Builder struct {
}

func (b *Builder) Build() selector.Balancer {
	return &Balancer{}
}