github.com/wangyougui/gf/v2@v2.6.5/net/gsel/gsel_selector_weight.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/wangyougui/gf.
     6  
     7  package gsel
     8  
     9  import (
    10  	"context"
    11  	"sync"
    12  
    13  	"github.com/wangyougui/gf/v2/internal/intlog"
    14  	"github.com/wangyougui/gf/v2/net/gsvc"
    15  	"github.com/wangyougui/gf/v2/util/grand"
    16  )
    17  
    18  type selectorWeight struct {
    19  	mu    sync.RWMutex
    20  	nodes Nodes
    21  }
    22  
    23  func NewSelectorWeight() Selector {
    24  	return &selectorWeight{
    25  		nodes: make(Nodes, 0),
    26  	}
    27  }
    28  
    29  func (s *selectorWeight) Update(ctx context.Context, nodes Nodes) error {
    30  	intlog.Printf(ctx, `Update nodes: %s`, nodes.String())
    31  	var newNodes []Node
    32  	for _, v := range nodes {
    33  		node := v
    34  		for i := 0; i < s.getWeight(node); i++ {
    35  			newNodes = append(newNodes, node)
    36  		}
    37  	}
    38  	s.mu.Lock()
    39  	s.nodes = newNodes
    40  	s.mu.Unlock()
    41  	return nil
    42  }
    43  
    44  func (s *selectorWeight) Pick(ctx context.Context) (node Node, done DoneFunc, err error) {
    45  	s.mu.RLock()
    46  	defer s.mu.RUnlock()
    47  	if len(s.nodes) == 0 {
    48  		return nil, nil, nil
    49  	}
    50  	node = s.nodes[grand.Intn(len(s.nodes))]
    51  	intlog.Printf(ctx, `Picked node: %s`, node.Address())
    52  	return node, nil, nil
    53  }
    54  
    55  func (s *selectorWeight) getWeight(node Node) int {
    56  	return node.Service().GetMetadata().Get(gsvc.MDWeight).Int()
    57  }