go.uber.org/yarpc@v1.72.1/peer/tworandomchoices/tworandomchoices.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 tworandomchoices
    22  
    23  import (
    24  	"math/rand"
    25  	"sync"
    26  	"time"
    27  
    28  	"go.uber.org/atomic"
    29  	"go.uber.org/yarpc/api/peer"
    30  	"go.uber.org/yarpc/api/transport"
    31  	"go.uber.org/yarpc/peer/abstractlist"
    32  )
    33  
    34  type twoRandomChoicesList struct {
    35  	subscribers []*subscriber
    36  	random      *rand.Rand
    37  
    38  	m sync.RWMutex
    39  }
    40  
    41  // Option configures the peer list implementation constructor.
    42  type Option interface {
    43  	apply(*options)
    44  }
    45  
    46  type options struct{}
    47  
    48  // NewImplementation creates a new fewest pending heap
    49  // abstractlist.Implementation.
    50  //
    51  // Use this constructor instead of NewList, when wanting to do custom peer
    52  // connection management.
    53  func NewImplementation(opts ...Option) abstractlist.Implementation {
    54  	return newTwoRandomChoicesList(10, rand.NewSource(time.Now().UnixNano()))
    55  }
    56  
    57  func newTwoRandomChoicesList(cap int, source rand.Source) *twoRandomChoicesList {
    58  	return &twoRandomChoicesList{
    59  		subscribers: make([]*subscriber, 0, cap),
    60  		random:      rand.New(source),
    61  	}
    62  }
    63  
    64  func (l *twoRandomChoicesList) Add(peer peer.StatusPeer, _ peer.Identifier) abstractlist.Subscriber {
    65  	l.m.Lock()
    66  	defer l.m.Unlock()
    67  
    68  	index := len(l.subscribers)
    69  	l.subscribers = append(l.subscribers, &subscriber{
    70  		index: index,
    71  		peer:  peer,
    72  	})
    73  	return l.subscribers[index]
    74  }
    75  
    76  func (l *twoRandomChoicesList) Remove(peer peer.StatusPeer, _ peer.Identifier, ps abstractlist.Subscriber) {
    77  	l.m.Lock()
    78  	defer l.m.Unlock()
    79  
    80  	sub, ok := ps.(*subscriber)
    81  	if !ok || len(l.subscribers) == 0 {
    82  		return
    83  	}
    84  	index := sub.index
    85  	last := len(l.subscribers) - 1
    86  	l.subscribers[index] = l.subscribers[last]
    87  	l.subscribers[index].index = index
    88  	l.subscribers = l.subscribers[0:last]
    89  }
    90  
    91  func (l *twoRandomChoicesList) Choose(_ *transport.Request) peer.StatusPeer {
    92  	// Usage of a wite lock because r.random.Intn is not thread safe
    93  	// see: https://golang.org/pkg/math/rand/
    94  	l.m.Lock()
    95  	defer l.m.Unlock()
    96  
    97  	numSubs := len(l.subscribers)
    98  	if numSubs == 0 {
    99  		return nil
   100  	}
   101  	if numSubs == 1 {
   102  		return l.subscribers[0].peer
   103  	}
   104  	i := l.random.Intn(numSubs)
   105  	j := i + 1 + l.random.Intn(numSubs-1)
   106  	if j >= numSubs {
   107  		j -= numSubs
   108  	}
   109  	if l.subscribers[i].pending.Load() > l.subscribers[j].pending.Load() {
   110  		i = j
   111  	}
   112  	return l.subscribers[i].peer
   113  }
   114  
   115  type subscriber struct {
   116  	index   int
   117  	peer    peer.StatusPeer
   118  	pending atomic.Int32
   119  }
   120  
   121  var _ abstractlist.Subscriber = (*subscriber)(nil)
   122  
   123  func (s *subscriber) UpdatePendingRequestCount(pendingRequestCount int) {
   124  	s.pending.Store(int32(pendingRequestCount))
   125  }