kit/selector/node/node.go

package node

import (
	"context"
	"git.diulo.com/mogfee/kit/selector"
	"google.golang.org/grpc/balancer"
	"sync/atomic"
	"time"
)

const (
	defaultWeight = 100
)

var (
	_ selector.WeightedNode        = (*Node)(nil)
	_ selector.WeightedNodeBuilder = (*Builder)(nil)
)

type Node struct {
	selector.Node
	lastPick int64
}
type Builder struct {
}

func (b *Builder) Build(node selector.Node) selector.WeightedNode {
	return &Node{
		Node:     node,
		lastPick: 0,
	}
}

func (n *Node) Raw() selector.Node {
	return n.Node
}

func (n *Node) Weight() float64 {
	if n.InitialWeight() != nil {
		return float64(*n.InitialWeight())
	}
	return defaultWeight
}

func (n *Node) Pick() selector.DoneFunc {
	now := time.Now().UnixNano()
	atomic.StoreInt64(&n.lastPick, now)
	return func(ctx context.Context, di balancer.DoneInfo) {}
}

func (n *Node) PickElapsed() time.Duration {
	return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick))
}