snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/wrr/edf.go

@@ -0,0 +1,92 @@
+/*
+ *
+ * Copyright 2019 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 wrr
+
+import (
+	"container/heap"
+	"sync"
+)
+
+// edfWrr is a struct for EDF weighted round robin implementation.
+type edfWrr struct {
+	lock               sync.Mutex
+	items              edfPriorityQueue
+	currentOrderOffset uint64
+	currentTime        float64
+}
+
+// NewEDF creates Earliest Deadline First (EDF)
+// (https://en.wikipedia.org/wiki/Earliest_deadline_first_scheduling) implementation for weighted round robin.
+// Each pick from the schedule has the earliest deadline entry selected. Entries have deadlines set
+// at current time + 1 / weight, providing weighted round robin behavior with O(log n) pick time.
+func NewEDF() WRR {
+	return &edfWrr{}
+}
+
+// edfEntry is an internal wrapper for item that also stores weight and relative position in the queue.
+type edfEntry struct {
+	deadline    float64
+	weight      int64
+	orderOffset uint64
+	item        any
+}
+
+// edfPriorityQueue is a heap.Interface implementation for edfEntry elements.
+type edfPriorityQueue []*edfEntry
+
+func (pq edfPriorityQueue) Len() int { return len(pq) }
+func (pq edfPriorityQueue) Less(i, j int) bool {
+	return pq[i].deadline < pq[j].deadline || pq[i].deadline == pq[j].deadline && pq[i].orderOffset < pq[j].orderOffset
+}
+func (pq edfPriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }
+
+func (pq *edfPriorityQueue) Push(x any) {
+	*pq = append(*pq, x.(*edfEntry))
+}
+
+func (pq *edfPriorityQueue) Pop() any {
+	old := *pq
+	*pq = old[0 : len(old)-1]
+	return old[len(old)-1]
+}
+
+func (edf *edfWrr) Add(item any, weight int64) {
+	edf.lock.Lock()
+	defer edf.lock.Unlock()
+	entry := edfEntry{
+		deadline:    edf.currentTime + 1.0/float64(weight),
+		weight:      weight,
+		item:        item,
+		orderOffset: edf.currentOrderOffset,
+	}
+	edf.currentOrderOffset++
+	heap.Push(&edf.items, &entry)
+}
+
+func (edf *edfWrr) Next() any {
+	edf.lock.Lock()
+	defer edf.lock.Unlock()
+	if len(edf.items) == 0 {
+		return nil
+	}
+	item := edf.items[0]
+	edf.currentTime = item.deadline
+	item.deadline = edf.currentTime + 1.0/float64(item.weight)
+	heap.Fix(&edf.items, 0)
+	return item.item
+}

internal/wrr/edf_test.go

@@ -0,0 +1,35 @@
+/*
+ *
+ * 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 wrr
+
+import (
+	"testing"
+)
+
+func (s) TestEDFOnEndpointsWithSameWeight(t *testing.T) {
+	wrr := NewEDF()
+	wrr.Add("1", 1)
+	wrr.Add("2", 1)
+	wrr.Add("3", 1)
+	expected := []string{"1", "2", "3", "1", "2", "3", "1", "2", "3", "1", "2", "3"}
+	for i := 0; i < len(expected); i++ {
+		item := wrr.Next().(string)
+		if item != expected[i] {
+			t.Errorf("wrr Next=%s, want=%s", item, expected[i])
+		}
+	}
+}

internal/wrr/random.go

@@ -0,0 +1,84 @@
+/*
+ *
+ * Copyright 2019 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 wrr
+
+import (
+	"fmt"
+	rand "math/rand/v2"
+	"sort"
+)
+
+// weightedItem is a wrapped weighted item that is used to implement weighted random algorithm.
+type weightedItem struct {
+	item              any
+	weight            int64
+	accumulatedWeight int64
+}
+
+func (w *weightedItem) String() string {
+	return fmt.Sprint(*w)
+}
+
+// randomWRR is a struct that contains weighted items implement weighted random algorithm.
+type randomWRR struct {
+	items []*weightedItem
+	// Are all item's weights equal
+	equalWeights bool
+}
+
+// NewRandom creates a new WRR with random.
+func NewRandom() WRR {
+	return &randomWRR{}
+}
+
+var randInt64n = rand.Int64N
+
+func (rw *randomWRR) Next() (item any) {
+	if len(rw.items) == 0 {
+		return nil
+	}
+	if rw.equalWeights {
+		return rw.items[randInt64n(int64(len(rw.items)))].item
+	}
+
+	sumOfWeights := rw.items[len(rw.items)-1].accumulatedWeight
+	// Random number in [0, sumOfWeights).
+	randomWeight := randInt64n(sumOfWeights)
+	// Item's accumulated weights are in ascending order, because item's weight >= 0.
+	// Binary search rw.items to find first item whose accumulatedWeight > randomWeight
+	// The return i is guaranteed to be in range [0, len(rw.items)) because randomWeight < last item's accumulatedWeight
+	i := sort.Search(len(rw.items), func(i int) bool { return rw.items[i].accumulatedWeight > randomWeight })
+	return rw.items[i].item
+}
+
+func (rw *randomWRR) Add(item any, weight int64) {
+	accumulatedWeight := weight
+	equalWeights := true
+	if len(rw.items) > 0 {
+		lastItem := rw.items[len(rw.items)-1]
+		accumulatedWeight = lastItem.accumulatedWeight + weight
+		equalWeights = rw.equalWeights && weight == lastItem.weight
+	}
+	rw.equalWeights = equalWeights
+	rItem := &weightedItem{item: item, weight: weight, accumulatedWeight: accumulatedWeight}
+	rw.items = append(rw.items, rItem)
+}
+
+func (rw *randomWRR) String() string {
+	return fmt.Sprint(rw.items)
+}

internal/wrr/wrr.go

@@ -0,0 +1,32 @@
+/*
+ *
+ * Copyright 2019 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 wrr contains the interface and common implementations of wrr
+// algorithms.
+package wrr
+
+// WRR defines an interface that implements weighted round robin.
+type WRR interface {
+	// Add adds an item with weight to the WRR set. Add must be only called
+	// before any calls to Next.
+	Add(item any, weight int64)
+	// Next returns the next picked item.
+	//
+	// Next needs to be thread safe. Add may not be called after any call to
+	// Next.
+	Next() any
+}

internal/wrr/wrr_test.go

@@ -0,0 +1,193 @@
+/*
+ *
+ * Copyright 2019 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 wrr
+
+import (
+	"errors"
+	"math"
+	rand "math/rand/v2"
+	"strconv"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+const iterCount = 10000
+
+func equalApproximate(a, b float64) error {
+	opt := cmp.Comparer(func(x, y float64) bool {
+		delta := math.Abs(x - y)
+		mean := math.Abs(x+y) / 2.0
+		return delta/mean < 0.05
+	})
+	if !cmp.Equal(a, b, opt) {
+		return errors.New(cmp.Diff(a, b))
+	}
+	return nil
+}
+
+func testWRRNext(t *testing.T, newWRR func() WRR) {
+	tests := []struct {
+		name    string
+		weights []int64
+	}{
+		{
+			name:    "1-1-1",
+			weights: []int64{1, 1, 1},
+		},
+		{
+			name:    "1-2-3",
+			weights: []int64{1, 2, 3},
+		},
+		{
+			name:    "5-3-2",
+			weights: []int64{5, 3, 2},
+		},
+		{
+			name:    "17-23-37",
+			weights: []int64{17, 23, 37},
+		},
+		{
+			name:    "no items",
+			weights: []int64{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := newWRR()
+			if len(tt.weights) == 0 {
+				if next := w.Next(); next != nil {
+					t.Fatalf("w.Next returns non nil value:%v when there is no item", next)
+				}
+				return
+			}
+
+			var sumOfWeights int64
+			for i, weight := range tt.weights {
+				w.Add(i, weight)
+				sumOfWeights += weight
+			}
+
+			results := make(map[int]int)
+			for i := 0; i < iterCount; i++ {
+				results[w.Next().(int)]++
+			}
+
+			wantRatio := make([]float64, len(tt.weights))
+			for i, weight := range tt.weights {
+				wantRatio[i] = float64(weight) / float64(sumOfWeights)
+			}
+			gotRatio := make([]float64, len(tt.weights))
+			for i, count := range results {
+				gotRatio[i] = float64(count) / iterCount
+			}
+
+			for i := range wantRatio {
+				if err := equalApproximate(gotRatio[i], wantRatio[i]); err != nil {
+					t.Errorf("%v not equal %v", i, err)
+				}
+			}
+		})
+	}
+}
+
+func (s) TestRandomWRRNext(t *testing.T) {
+	testWRRNext(t, NewRandom)
+}
+
+func (s) TestEdfWrrNext(t *testing.T) {
+	testWRRNext(t, NewEDF)
+}
+
+func BenchmarkRandomWRRNext(b *testing.B) {
+	for _, n := range []int{100, 500, 1000} {
+		b.Run("equal-weights-"+strconv.Itoa(n)+"-items", func(b *testing.B) {
+			w := NewRandom()
+			sumOfWeights := n
+			for i := 0; i < n; i++ {
+				w.Add(i, 1)
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				for i := 0; i < sumOfWeights; i++ {
+					w.Next()
+				}
+			}
+		})
+	}
+
+	var maxWeight int64 = 1024
+	for _, n := range []int{100, 500, 1000} {
+		b.Run("random-weights-"+strconv.Itoa(n)+"-items", func(b *testing.B) {
+			w := NewRandom()
+			var sumOfWeights int64
+			for i := 0; i < n; i++ {
+				weight := rand.Int64N(maxWeight + 1)
+				w.Add(i, weight)
+				sumOfWeights += weight
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				for i := 0; i < int(sumOfWeights); i++ {
+					w.Next()
+				}
+			}
+		})
+	}
+
+	itemsNum := 200
+	heavyWeight := int64(itemsNum)
+	lightWeight := int64(1)
+	heavyIndices := []int{0, itemsNum / 2, itemsNum - 1}
+	for _, heavyIndex := range heavyIndices {
+		b.Run("skew-weights-heavy-index-"+strconv.Itoa(heavyIndex), func(b *testing.B) {
+			w := NewRandom()
+			var sumOfWeights int64
+			for i := 0; i < itemsNum; i++ {
+				var weight int64
+				if i == heavyIndex {
+					weight = heavyWeight
+				} else {
+					weight = lightWeight
+				}
+				sumOfWeights += weight
+				w.Add(i, weight)
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				for i := 0; i < int(sumOfWeights); i++ {
+					w.Next()
+				}
+			}
+		})
+	}
+}
+
+func init() {
+	r := rand.New(rand.NewPCG(0, 0))
+	randInt64n = r.Int64N
+}