google.golang.org/grpc@v1.62.1/internal/wrr/random.go (about)

     1  /*
     2   *
     3   * Copyright 2019 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  package wrr
    19  
    20  import (
    21  	"fmt"
    22  	"sort"
    23  
    24  	"google.golang.org/grpc/internal/grpcrand"
    25  )
    26  
    27  // weightedItem is a wrapped weighted item that is used to implement weighted random algorithm.
    28  type weightedItem struct {
    29  	item              any
    30  	weight            int64
    31  	accumulatedWeight int64
    32  }
    33  
    34  func (w *weightedItem) String() string {
    35  	return fmt.Sprint(*w)
    36  }
    37  
    38  // randomWRR is a struct that contains weighted items implement weighted random algorithm.
    39  type randomWRR struct {
    40  	items []*weightedItem
    41  	// Are all item's weights equal
    42  	equalWeights bool
    43  }
    44  
    45  // NewRandom creates a new WRR with random.
    46  func NewRandom() WRR {
    47  	return &randomWRR{}
    48  }
    49  
    50  var grpcrandInt63n = grpcrand.Int63n
    51  
    52  func (rw *randomWRR) Next() (item any) {
    53  	if len(rw.items) == 0 {
    54  		return nil
    55  	}
    56  	if rw.equalWeights {
    57  		return rw.items[grpcrandInt63n(int64(len(rw.items)))].item
    58  	}
    59  
    60  	sumOfWeights := rw.items[len(rw.items)-1].accumulatedWeight
    61  	// Random number in [0, sumOfWeights).
    62  	randomWeight := grpcrandInt63n(sumOfWeights)
    63  	// Item's accumulated weights are in ascending order, because item's weight >= 0.
    64  	// Binary search rw.items to find first item whose accumulatedWeight > randomWeight
    65  	// The return i is guaranteed to be in range [0, len(rw.items)) because randomWeight < last item's accumulatedWeight
    66  	i := sort.Search(len(rw.items), func(i int) bool { return rw.items[i].accumulatedWeight > randomWeight })
    67  	return rw.items[i].item
    68  }
    69  
    70  func (rw *randomWRR) Add(item any, weight int64) {
    71  	accumulatedWeight := weight
    72  	equalWeights := true
    73  	if len(rw.items) > 0 {
    74  		lastItem := rw.items[len(rw.items)-1]
    75  		accumulatedWeight = lastItem.accumulatedWeight + weight
    76  		equalWeights = rw.equalWeights && weight == lastItem.weight
    77  	}
    78  	rw.equalWeights = equalWeights
    79  	rItem := &weightedItem{item: item, weight: weight, accumulatedWeight: accumulatedWeight}
    80  	rw.items = append(rw.items, rItem)
    81  }
    82  
    83  func (rw *randomWRR) String() string {
    84  	return fmt.Sprint(rw.items)
    85  }