go.uber.org/yarpc@v1.72.1/peer/roundrobin/peerring.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package roundrobin
    22  
    23  import (
    24  	"container/ring"
    25  	"sync"
    26  
    27  	"go.uber.org/yarpc/api/peer"
    28  	"go.uber.org/yarpc/api/transport"
    29  	"go.uber.org/yarpc/peer/abstractlist"
    30  )
    31  
    32  // Option configures the peer list implementation constructor.
    33  type Option interface {
    34  	apply(*options)
    35  }
    36  
    37  type options struct{}
    38  
    39  // NewImplementation creates a new round-robin abstractlist.Implementation.
    40  //
    41  // Use this constructor instead of NewList, when wanting to do custom peer
    42  // connection management.
    43  func NewImplementation(opts ...Option) abstractlist.Implementation {
    44  	return &peerRing{}
    45  }
    46  
    47  type subscriber struct {
    48  	peer peer.StatusPeer
    49  	node *ring.Ring
    50  }
    51  
    52  func (s *subscriber) UpdatePendingRequestCount(int) {}
    53  
    54  // peerRing provides a safe way to interact (Add/Remove/Get) with a potentially
    55  // changing list of peer objects
    56  type peerRing struct {
    57  	nextNode *ring.Ring
    58  
    59  	m sync.RWMutex
    60  }
    61  
    62  // Add a peer.StatusPeer to the end of the peerRing, if the ring is empty it
    63  // initializes the nextNode marker
    64  func (pr *peerRing) Add(p peer.StatusPeer, _ peer.Identifier) abstractlist.Subscriber {
    65  	pr.m.Lock()
    66  	defer pr.m.Unlock()
    67  
    68  	sub := &subscriber{peer: p}
    69  	newNode := ring.New(1)
    70  	newNode.Value = sub
    71  	sub.node = newNode
    72  
    73  	if pr.nextNode == nil {
    74  		// Empty ring, add the first node
    75  		pr.nextNode = newNode
    76  	} else {
    77  		// Push the node to the ring
    78  		pushBeforeNode(pr.nextNode, newNode)
    79  	}
    80  	return sub
    81  }
    82  
    83  // Remove the peer from the ring. Use the subscriber to address the node of the
    84  // ring directly.
    85  func (pr *peerRing) Remove(p peer.StatusPeer, _ peer.Identifier, s abstractlist.Subscriber) {
    86  	pr.m.Lock()
    87  	defer pr.m.Unlock()
    88  
    89  	sub, ok := s.(*subscriber)
    90  	if !ok {
    91  		// Don't panic.
    92  		return
    93  	}
    94  
    95  	node := sub.node
    96  	if isLastRingNode(node) {
    97  		pr.nextNode = nil
    98  	} else {
    99  		if pr.isNextNode(node) {
   100  			pr.nextNode = pr.nextNode.Next()
   101  		}
   102  		popNodeFromRing(node)
   103  	}
   104  }
   105  
   106  func (pr *peerRing) isNextNode(node *ring.Ring) bool {
   107  	return pr.nextNode == node
   108  }
   109  
   110  // Choose returns the next peer in the ring, or nil if there is no peer in the ring
   111  // after it has the next peer, it increments the nextPeer marker in the ring
   112  func (pr *peerRing) Choose(_ *transport.Request) peer.StatusPeer {
   113  	pr.m.Lock()
   114  	defer pr.m.Unlock()
   115  
   116  	if pr.nextNode == nil {
   117  		return nil
   118  	}
   119  
   120  	p := getPeerForRingNode(pr.nextNode)
   121  	pr.nextNode = pr.nextNode.Next()
   122  
   123  	return p
   124  }
   125  
   126  func getPeerForRingNode(rNode *ring.Ring) peer.StatusPeer {
   127  	return rNode.Value.(*subscriber).peer
   128  }
   129  
   130  func isLastRingNode(rNode *ring.Ring) bool {
   131  	return rNode.Next() == rNode
   132  }
   133  
   134  func popNodeFromRing(rNode *ring.Ring) {
   135  	rNode.Prev().Unlink(1)
   136  }
   137  
   138  func pushBeforeNode(curNode, newNode *ring.Ring) {
   139  	curNode.Prev().Link(newNode)
   140  }