github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/interleaved_weighted_round_robin.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  	"sync"
    22  
    23  	"github.com/bytedance/gopkg/lang/fastrand"
    24  
    25  	"github.com/cloudwego/kitex/pkg/discovery"
    26  )
    27  
    28  type iwrrNode struct {
    29  	discovery.Instance
    30  	remainder int
    31  
    32  	next *iwrrNode
    33  }
    34  
    35  type iwrrQueue struct {
    36  	head *iwrrNode
    37  	tail *iwrrNode
    38  }
    39  
    40  type InterleavedWeightedRoundRobinPicker struct {
    41  	current *iwrrQueue
    42  	next    *iwrrQueue
    43  	gcd     int
    44  
    45  	lock sync.Mutex
    46  }
    47  
    48  func newInterleavedWeightedRoundRobinPicker(instances []discovery.Instance) Picker {
    49  	iwrrp := new(InterleavedWeightedRoundRobinPicker)
    50  	iwrrp.current = newIwrrQueue()
    51  	iwrrp.next = newIwrrQueue()
    52  
    53  	size := uint64(len(instances))
    54  	offset := fastrand.Uint64n(size)
    55  	gcd := 0
    56  	for idx := uint64(0); idx < size; idx++ {
    57  		ins := instances[(idx+offset)%size]
    58  		gcd = gcdInt(gcd, ins.Weight())
    59  
    60  		iwrrp.current.enqueue(&iwrrNode{
    61  			Instance:  ins,
    62  			remainder: ins.Weight(),
    63  		})
    64  	}
    65  
    66  	iwrrp.gcd = gcd
    67  
    68  	return iwrrp
    69  }
    70  
    71  func (ip *InterleavedWeightedRoundRobinPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
    72  	ip.lock.Lock()
    73  	defer ip.lock.Unlock()
    74  
    75  	if ip.current.empty() {
    76  		ip.current, ip.next = ip.next, ip.current
    77  	}
    78  
    79  	node := ip.current.dequeue()
    80  	node.remainder -= ip.gcd
    81  
    82  	if node.remainder > 0 {
    83  		ip.current.enqueue(node)
    84  	} else {
    85  		node.remainder = node.Instance.Weight()
    86  		ip.next.enqueue(node)
    87  	}
    88  
    89  	return node.Instance
    90  }
    91  
    92  func newIwrrQueue() *iwrrQueue {
    93  	return &iwrrQueue{}
    94  }
    95  
    96  func (q *iwrrQueue) enqueue(node *iwrrNode) {
    97  	node.next = nil
    98  	tail := q.tail
    99  	q.tail = node
   100  	if tail == nil {
   101  		q.head = node
   102  	} else {
   103  		tail.next = node
   104  	}
   105  }
   106  
   107  func (q *iwrrQueue) dequeue() *iwrrNode {
   108  	head := q.head
   109  	next := head.next
   110  	head.next = nil
   111  	q.head = next
   112  	if next == nil {
   113  		q.tail = nil
   114  	}
   115  	return head
   116  }
   117  
   118  func (q *iwrrQueue) empty() bool {
   119  	return q.head == nil
   120  }