interleaved_weighted_round_robin.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"
	"sync"

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

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

type iwrrNode struct {
	discovery.Instance
	remainder int

	next *iwrrNode
}

type iwrrQueue struct {
	head *iwrrNode
	tail *iwrrNode
}

type InterleavedWeightedRoundRobinPicker struct {
	current *iwrrQueue
	next    *iwrrQueue
	gcd     int

	lock sync.Mutex
}

func newInterleavedWeightedRoundRobinPicker(instances []discovery.Instance) Picker {
	iwrrp := new(InterleavedWeightedRoundRobinPicker)
	iwrrp.current = newIwrrQueue()
	iwrrp.next = newIwrrQueue()

	size := uint64(len(instances))
	offset := fastrand.Uint64n(size)
	gcd := 0
	for idx := uint64(0); idx < size; idx++ {
		ins := instances[(idx+offset)%size]
		gcd = gcdInt(gcd, ins.Weight())

		iwrrp.current.enqueue(&iwrrNode{
			Instance:  ins,
			remainder: ins.Weight(),
		})
	}

	iwrrp.gcd = gcd

	return iwrrp
}

func (ip *InterleavedWeightedRoundRobinPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
	ip.lock.Lock()
	defer ip.lock.Unlock()

	if ip.current.empty() {
		ip.current, ip.next = ip.next, ip.current
	}

	node := ip.current.dequeue()
	node.remainder -= ip.gcd

	if node.remainder > 0 {
		ip.current.enqueue(node)
	} else {
		node.remainder = node.Instance.Weight()
		ip.next.enqueue(node)
	}

	return node.Instance
}

func newIwrrQueue() *iwrrQueue {
	return &iwrrQueue{}
}

func (q *iwrrQueue) enqueue(node *iwrrNode) {
	node.next = nil
	tail := q.tail
	q.tail = node
	if tail == nil {
		q.head = node
	} else {
		tail.next = node
	}
}

func (q *iwrrQueue) dequeue() *iwrrNode {
	head := q.head
	next := head.next
	head.next = nil
	q.head = next
	if next == nil {
		q.tail = nil
	}
	return head
}

func (q *iwrrQueue) empty() bool {
	return q.head == nil
}