weighted_balancer_test.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"
	"fmt"
	"math"
	"strconv"
	"testing"

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

type balancerTestcase struct {
	Name    string
	factory func() Loadbalancer
}

var balancerTestcases = []*balancerTestcase{
	{Name: "weight_round_robin", factory: NewWeightedRoundRobinBalancer},
	{Name: "weight_random", factory: NewWeightedRandomBalancer},
	{Name: "weight_random_with_alias_method", factory: NewWeightedRandomWithAliasMethodBalancer},
	{Name: "interleaved_weighted_round_robin", factory: NewInterleavedWeightedRoundRobinBalancer},
}

func TestWeightedBalancer_GetPicker(t *testing.T) {
	for _, tc := range balancerTestcases {
		t.Run(tc.Name, func(t *testing.T) {
			balancer := tc.factory()
			// nil
			picker := balancer.GetPicker(discovery.Result{})
			test.Assert(t, picker != nil)
			dp, ok := picker.(*DummyPicker)
			test.Assert(t, ok && dp != nil)

			// invalid
			picker = balancer.GetPicker(discovery.Result{
				Instances: []discovery.Instance{
					discovery.NewInstance("tcp", "addr1", -10, nil),
					discovery.NewInstance("tcp", "addr2", -20, nil),
				},
			})
			test.Assert(t, picker != nil)
			dp, ok = picker.(*DummyPicker)
			test.Assert(t, ok && dp != nil)

			// one instance
			insList := []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
			}
			picker = balancer.GetPicker(discovery.Result{
				Instances: insList,
			})
			test.Assert(t, picker != nil)

			// multi instances
			insList = []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", 20, nil),
				discovery.NewInstance("tcp", "addr3", 30, nil),
			}
			picker = balancer.GetPicker(discovery.Result{
				Instances: insList,
			})
			test.Assert(t, picker != nil)

			// balanced instances
			insList = []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", 10, nil),
				discovery.NewInstance("tcp", "addr3", 10, nil),
			}
			picker = balancer.GetPicker(discovery.Result{
				Instances: insList,
			})
			test.Assert(t, picker != nil)
			test.Assert(t, balancer.Name() == tc.Name)
		})
	}
}

func TestWeightedPicker_Next(t *testing.T) {
	for _, tc := range balancerTestcases {
		t.Run(tc.Name, func(t *testing.T) {
			balancer := tc.factory()
			ctx := context.Background()
			// nil
			picker := balancer.GetPicker(discovery.Result{})
			ins := picker.Next(ctx, nil)
			test.Assert(t, ins == nil)

			// empty instance
			picker = balancer.GetPicker(discovery.Result{
				Instances: make([]discovery.Instance, 0),
			})
			ins = picker.Next(ctx, nil)
			test.Assert(t, ins == nil)

			// one instance
			insList := []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
			}
			for i := 0; i < 100; i++ {
				picker := balancer.GetPicker(discovery.Result{
					Instances: insList,
				})
				ins := picker.Next(ctx, nil)
				test.Assert(t, ins.Weight() == 10)
			}

			// multi instances, weightSum > 0
			insList = []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", 20, nil),
				discovery.NewInstance("tcp", "addr3", 50, nil),
				discovery.NewInstance("tcp", "addr4", 100, nil),
				discovery.NewInstance("tcp", "addr5", 200, nil),
				discovery.NewInstance("tcp", "addr6", 500, nil),
			}
			var weightSum int
			for _, ins := range insList {
				weight := ins.Weight()
				weightSum += weight
			}
			n := 10000000
			pickedStat := map[int]int{}
			for i := 0; i < n; i++ {
				picker := balancer.GetPicker(discovery.Result{
					Instances: insList,
					Cacheable: true,
				})
				ins := picker.Next(ctx, nil)
				weight := ins.Weight()
				if pickedCnt, ok := pickedStat[weight]; ok {
					pickedStat[weight] = pickedCnt + 1
				} else {
					pickedStat[weight] = 1
				}
			}

			for _, ins := range insList {
				weight := ins.Weight()
				expect := float64(weight) / float64(weightSum) * float64(n)
				actual := float64(pickedStat[weight])
				delta := math.Abs(expect - actual)
				test.Assertf(t, delta/expect < 0.01, "delta(%f)/expect(%f) = %f", delta, expect, delta/expect)
			}

			// weightSum = 0
			insList = []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", -10, nil),
			}
			picker = balancer.GetPicker(discovery.Result{
				Instances: insList,
			})
			test.Assert(t, picker.Next(ctx, nil) != nil)

			// weightSum < 0
			insList = []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", -20, nil),
			}
			picker = balancer.GetPicker(discovery.Result{
				Instances: insList,
			})
			test.Assert(t, picker.Next(ctx, nil) != nil)
		})
	}
}

func TestWeightedPicker_NoMoreInstance(t *testing.T) {
	for _, tc := range balancerTestcases {
		t.Run(tc.Name, func(t *testing.T) {
			balancer := tc.factory()
			ctx := context.Background()

			// multi instances, weightSum > 0
			insList := []discovery.Instance{
				discovery.NewInstance("tcp", "addr1", 10, nil),
				discovery.NewInstance("tcp", "addr2", 20, nil),
				discovery.NewInstance("tcp", "addr3", 50, nil),
				discovery.NewInstance("tcp", "addr4", 100, nil),
				discovery.NewInstance("tcp", "addr5", 200, nil),
				discovery.NewInstance("tcp", "addr6", 500, nil),
			}

			picker := balancer.GetPicker(discovery.Result{
				Instances: insList,
			})

			for i := 0; i < len(insList); i++ {
				ins := picker.Next(ctx, nil)
				test.Assert(t, ins != nil)
			}

			ins := picker.Next(ctx, nil)
			test.Assert(t, ins != nil)
		})
	}
}

func makeNInstances(n, weight int) (res []discovery.Instance) {
	for i := 0; i < n; i++ {
		res = append(res, discovery.NewInstance("tcp", fmt.Sprintf("addr[%d]-weight[%d]", i, weight), weight, nil))
	}
	return
}

func BenchmarkWeightedPicker(b *testing.B) {
	ctx := context.Background()

	for _, tc := range balancerTestcases {
		b.Run(tc.Name, func(b *testing.B) {
			balancer := tc.factory()
			n := 10
			for i := 0; i < 4; i++ {
				b.Run(fmt.Sprintf("%dins", n), func(b *testing.B) {
					inss := makeNInstances(n, 10)
					e := discovery.Result{
						Cacheable: true,
						CacheKey:  fmt.Sprintf("test-%d", n),
						Instances: inss,
					}
					picker := balancer.GetPicker(e)
					picker.Next(ctx, nil)
					if r, ok := picker.(internal.Reusable); ok {
						r.Recycle()
					}
					b.ReportAllocs()
					b.ResetTimer()
					for i := 0; i < b.N; i++ {
						picker := balancer.GetPicker(e)
						picker.Next(ctx, nil)
						if r, ok := picker.(internal.Reusable); ok {
							r.Recycle()
						}
					}
				})
				n *= 10
			}
		})
	}
}

func BenchmarkGetPicker(b *testing.B) {
	insList := genInstList()
	for _, tc := range balancerTestcases {
		b.Run(tc.Name, func(b *testing.B) {
			balancer := tc.factory()
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				balancer.GetPicker(discovery.Result{
					Instances: insList,
				})
			}
		})
	}
}

func BenchmarkWeightedPicker_Next(b *testing.B) {
	ctx := context.Background()

	makeNInstances := func(n, base int, seq []int) (res []discovery.Instance) {
		for i := 0; i < n; i++ {
			weight := seq[i%len(seq)] * base
			res = append(res, discovery.NewInstance("tcp", fmt.Sprintf("addr[%d]-weight[%d]", i, weight), weight, nil))
		}
		return
	}

	for _, tc := range balancerTestcases {
		b.Run(tc.Name, func(b *testing.B) {
			balancer := tc.factory()
			n := 10
			for i := 0; i < 4; i++ {
				b.Run(fmt.Sprintf("%dins", n), func(b *testing.B) {
					inss := makeNInstances(n, 1000, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
					e := discovery.Result{
						Cacheable: true,
						CacheKey:  fmt.Sprintf("test-%d", n),
						Instances: inss,
					}
					picker := balancer.GetPicker(e)
					b.ResetTimer()
					for i := 0; i < b.N; i++ {
						picker.Next(ctx, nil)
					}
				})
				n *= 10
			}
		})
	}
}

func genInstList() []discovery.Instance {
	n := 1000
	insList := make([]discovery.Instance, 0, n)
	for i := 0; i < n; i++ {
		insList = append(insList, discovery.NewInstance("tcp", "addr"+strconv.Itoa(i), 10, nil))
	}
	return insList
}