github.com/gogf/gf/v2@v2.7.4/net/gsel/gsel_selector_least_connection.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/gogf/gf.
     6  
     7  package gsel
     8  
     9  import (
    10  	"context"
    11  	"sync"
    12  
    13  	"github.com/gogf/gf/v2/container/gtype"
    14  	"github.com/gogf/gf/v2/internal/intlog"
    15  )
    16  
    17  type selectorLeastConnection struct {
    18  	mu    sync.RWMutex
    19  	nodes []*leastConnectionNode
    20  }
    21  
    22  type leastConnectionNode struct {
    23  	Node
    24  	inflight *gtype.Int
    25  }
    26  
    27  func NewSelectorLeastConnection() Selector {
    28  	return &selectorLeastConnection{
    29  		nodes: make([]*leastConnectionNode, 0),
    30  	}
    31  }
    32  
    33  func (s *selectorLeastConnection) Update(ctx context.Context, nodes Nodes) error {
    34  	intlog.Printf(ctx, `Update nodes: %s`, nodes.String())
    35  	var newNodes []*leastConnectionNode
    36  	for _, v := range nodes {
    37  		node := v
    38  		newNodes = append(newNodes, &leastConnectionNode{
    39  			Node:     node,
    40  			inflight: gtype.NewInt(),
    41  		})
    42  	}
    43  	s.mu.Lock()
    44  	defer s.mu.Unlock()
    45  	s.nodes = newNodes
    46  	return nil
    47  }
    48  
    49  func (s *selectorLeastConnection) Pick(ctx context.Context) (node Node, done DoneFunc, err error) {
    50  	s.mu.RLock()
    51  	defer s.mu.RUnlock()
    52  	var pickedNode *leastConnectionNode
    53  	if len(s.nodes) == 1 {
    54  		pickedNode = s.nodes[0]
    55  	} else {
    56  		for _, v := range s.nodes {
    57  			if pickedNode == nil {
    58  				pickedNode = v
    59  			} else if v.inflight.Val() < pickedNode.inflight.Val() {
    60  				pickedNode = v
    61  			}
    62  		}
    63  	}
    64  	pickedNode.inflight.Add(1)
    65  	done = func(ctx context.Context, di DoneInfo) {
    66  		pickedNode.inflight.Add(-1)
    67  	}
    68  	node = pickedNode.Node
    69  	intlog.Printf(ctx, `Picked node: %s`, node.Address())
    70  	return node, done, nil
    71  }