weighted_random_with_alias_method.go

/*
 * Copyright 2023 CloudWeGo 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 loadbalance

import (
	"context"

	"github.com/bytedance/gopkg/lang/fastrand"

	"github.com/cloudwego/kitex/pkg/discovery"
)

type AliasMethodPicker struct {
	instances []discovery.Instance
	weightSum int
	alias     []int
	prob      []float64
}

func newAliasMethodPicker(instances []discovery.Instance, weightSum int) Picker {
	picker := &AliasMethodPicker{
		instances: instances,
		weightSum: weightSum,
	}
	picker.init()
	return picker
}

// Alias Method need to init before use and after update instances
func (a *AliasMethodPicker) init() {
	n := len(a.instances)
	a.alias = make([]int, n)
	a.prob = make([]float64, n)

	totalWeight := a.weightSum

	scaledProb := make([]float64, n)
	small := make([]int, 0, n)
	large := make([]int, 0, n)

	for i, instance := range a.instances {
		scaledProb[i] = float64(instance.Weight()) * float64(n) / float64(totalWeight)
		if scaledProb[i] < 1.0 {
			small = append(small, i)
		} else {
			large = append(large, i)
		}
	}

	for len(small) > 0 && len(large) > 0 {
		l := small[len(small)-1]
		small = small[:len(small)-1]
		g := large[len(large)-1]
		large = large[:len(large)-1]

		a.prob[l] = scaledProb[l]
		a.alias[l] = g

		scaledProb[g] -= 1.0 - scaledProb[l]
		if scaledProb[g] < 1.0 {
			small = append(small, g)
		} else {
			large = append(large, g)
		}
	}

	for len(large) > 0 {
		g := large[len(large)-1]
		large = large[:len(large)-1]
		a.prob[g] = 1.0
	}

	for len(small) > 0 {
		l := small[len(small)-1]
		small = small[:len(small)-1]
		a.prob[l] = 1.0
	}
}

// Next implements the Picker interface.
func (a *AliasMethodPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
	i := fastrand.Intn(len(a.instances))
	if fastrand.Float64() < a.prob[i] {
		return a.instances[i]
	}
	return a.instances[a.alias[i]]
}