12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091 |
- /*
- *
- * Copyright 2019 gRPC authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
- package wrr
- import (
- "fmt"
- "sort"
- "sync"
- "google.golang.org/grpc/internal/grpcrand"
- )
- // weightedItem is a wrapped weighted item that is used to implement weighted random algorithm.
- type weightedItem struct {
- item interface{}
- weight int64
- accumulatedWeight int64
- }
- func (w *weightedItem) String() string {
- return fmt.Sprint(*w)
- }
- // randomWRR is a struct that contains weighted items implement weighted random algorithm.
- type randomWRR struct {
- mu sync.RWMutex
- items []*weightedItem
- // Are all item's weights equal
- equalWeights bool
- }
- // NewRandom creates a new WRR with random.
- func NewRandom() WRR {
- return &randomWRR{}
- }
- var grpcrandInt63n = grpcrand.Int63n
- func (rw *randomWRR) Next() (item interface{}) {
- rw.mu.RLock()
- defer rw.mu.RUnlock()
- if len(rw.items) == 0 {
- return nil
- }
- if rw.equalWeights {
- return rw.items[grpcrandInt63n(int64(len(rw.items)))].item
- }
- sumOfWeights := rw.items[len(rw.items)-1].accumulatedWeight
- // Random number in [0, sumOfWeights).
- randomWeight := grpcrandInt63n(sumOfWeights)
- // Item's accumulated weights are in ascending order, because item's weight >= 0.
- // Binary search rw.items to find first item whose accumulatedWeight > randomWeight
- // The return i is guaranteed to be in range [0, len(rw.items)) because randomWeight < last item's accumulatedWeight
- i := sort.Search(len(rw.items), func(i int) bool { return rw.items[i].accumulatedWeight > randomWeight })
- return rw.items[i].item
- }
- func (rw *randomWRR) Add(item interface{}, weight int64) {
- rw.mu.Lock()
- defer rw.mu.Unlock()
- accumulatedWeight := weight
- equalWeights := true
- if len(rw.items) > 0 {
- lastItem := rw.items[len(rw.items)-1]
- accumulatedWeight = lastItem.accumulatedWeight + weight
- equalWeights = rw.equalWeights && weight == lastItem.weight
- }
- rw.equalWeights = equalWeights
- rItem := &weightedItem{item: item, weight: weight, accumulatedWeight: accumulatedWeight}
- rw.items = append(rw.items, rItem)
- }
- func (rw *randomWRR) String() string {
- return fmt.Sprint(rw.items)
- }
|