wrr.go

/*
 *
 * Copyright 2020 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 testutils

import (
	"fmt"
	"sync"

	"google.golang.org/grpc/internal/wrr"
)

// testWRR is a deterministic WRR implementation.
//
// The real implementation does random WRR. testWRR makes the balancer behavior
// deterministic and easier to test.
//
// With {a: 2, b: 3}, the Next() results will be {a, a, b, b, b}.
type testWRR struct {
	itemsWithWeight []struct {
		item   any
		weight int64
	}
	length int

	mu    sync.Mutex
	idx   int   // The index of the item that will be picked
	count int64 // The number of times the current item has been picked.
}

// NewTestWRR return a WRR for testing. It's deterministic instead of random.
func NewTestWRR() wrr.WRR {
	return &testWRR{}
}

func (twrr *testWRR) Add(item any, weight int64) {
	twrr.itemsWithWeight = append(twrr.itemsWithWeight, struct {
		item   any
		weight int64
	}{item: item, weight: weight})
	twrr.length++
}

func (twrr *testWRR) Next() any {
	twrr.mu.Lock()
	iww := twrr.itemsWithWeight[twrr.idx]
	twrr.count++
	if twrr.count >= iww.weight {
		twrr.idx = (twrr.idx + 1) % twrr.length
		twrr.count = 0
	}
	twrr.mu.Unlock()
	return iww.item
}

func (twrr *testWRR) String() string {
	return fmt.Sprint(twrr.itemsWithWeight)
}