github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/weighted_random_with_alias_method.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package loadbalance
    18  
    19  import (
    20  	"context"
    21  
    22  	"github.com/bytedance/gopkg/lang/fastrand"
    23  
    24  	"github.com/cloudwego/kitex/pkg/discovery"
    25  )
    26  
    27  type AliasMethodPicker struct {
    28  	instances []discovery.Instance
    29  	weightSum int
    30  	alias     []int
    31  	prob      []float64
    32  }
    33  
    34  func newAliasMethodPicker(instances []discovery.Instance, weightSum int) Picker {
    35  	picker := &AliasMethodPicker{
    36  		instances: instances,
    37  		weightSum: weightSum,
    38  	}
    39  	picker.init()
    40  	return picker
    41  }
    42  
    43  // Alias Method need to init before use and after update instances
    44  func (a *AliasMethodPicker) init() {
    45  	n := len(a.instances)
    46  	a.alias = make([]int, n)
    47  	a.prob = make([]float64, n)
    48  
    49  	totalWeight := a.weightSum
    50  
    51  	scaledProb := make([]float64, n)
    52  	small := make([]int, 0, n)
    53  	large := make([]int, 0, n)
    54  
    55  	for i, instance := range a.instances {
    56  		scaledProb[i] = float64(instance.Weight()) * float64(n) / float64(totalWeight)
    57  		if scaledProb[i] < 1.0 {
    58  			small = append(small, i)
    59  		} else {
    60  			large = append(large, i)
    61  		}
    62  	}
    63  
    64  	for len(small) > 0 && len(large) > 0 {
    65  		l := small[len(small)-1]
    66  		small = small[:len(small)-1]
    67  		g := large[len(large)-1]
    68  		large = large[:len(large)-1]
    69  
    70  		a.prob[l] = scaledProb[l]
    71  		a.alias[l] = g
    72  
    73  		scaledProb[g] -= 1.0 - scaledProb[l]
    74  		if scaledProb[g] < 1.0 {
    75  			small = append(small, g)
    76  		} else {
    77  			large = append(large, g)
    78  		}
    79  	}
    80  
    81  	for len(large) > 0 {
    82  		g := large[len(large)-1]
    83  		large = large[:len(large)-1]
    84  		a.prob[g] = 1.0
    85  	}
    86  
    87  	for len(small) > 0 {
    88  		l := small[len(small)-1]
    89  		small = small[:len(small)-1]
    90  		a.prob[l] = 1.0
    91  	}
    92  }
    93  
    94  // Next implements the Picker interface.
    95  func (a *AliasMethodPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
    96  	i := fastrand.Intn(len(a.instances))
    97  	if fastrand.Float64() < a.prob[i] {
    98  		return a.instances[i]
    99  	}
   100  	return a.instances[a.alias[i]]
   101  }