From afc68a3944a0e867e9ef5667869cf65b0b14b124 Mon Sep 17 00:00:00 2001
From: zhuwei <wei.zhu@sinobbd.com>
Date: Wed, 21 Mar 2018 14:03:20 +0800
Subject: [PATCH] internal/util: new UniqRands() implementation

using the Fisher-Yates shuffle algorithm

faster and more consistent runtime when
quantity is near maxval (aka l is near n)
---
 internal/util/rand.go      | 31 +++++++++++++++++--------------
 internal/util/util_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 14 deletions(-)
 create mode 100644 internal/util/util_test.go

diff --git a/internal/util/rand.go b/internal/util/rand.go
index 7f6604f53..c444098ae 100644
--- a/internal/util/rand.go
+++ b/internal/util/rand.go
@@ -4,19 +4,22 @@ import (
 	"math/rand"
 )
 
-func UniqRands(l int, n int) []int {
-	set := make(map[int]struct{})
-	nums := make([]int, 0, l)
-	for {
-		num := rand.Intn(n)
-		if _, ok := set[num]; !ok {
-			set[num] = struct{}{}
-			nums = append(nums, num)
-		}
-		if len(nums) == l {
-			goto exit
-		}
+func UniqRands(quantity int, maxval int) []int {
+	if maxval < quantity {
+		quantity = maxval
 	}
-exit:
-	return nums
+
+	intSlice := make([]int, maxval)
+	for i := 0; i < maxval; i++ {
+		intSlice[i] = i
+	}
+
+	for i := 0; i < quantity; i++ {
+		j := rand.Int()%maxval + i
+		// swap
+		intSlice[i], intSlice[j] = intSlice[j], intSlice[i]
+		maxval--
+
+	}
+	return intSlice[0:quantity]
 }
diff --git a/internal/util/util_test.go b/internal/util/util_test.go
new file mode 100644
index 000000000..9a55db621
--- /dev/null
+++ b/internal/util/util_test.go
@@ -0,0 +1,36 @@
+package util
+
+import (
+	"testing"
+
+	"github.com/nsqio/nsq/internal/test"
+)
+
+func BenchmarkUniqRands5of5(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(5, 5)
+	}
+}
+func BenchmarkUniqRands20of20(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(20, 20)
+	}
+}
+
+func BenchmarkUniqRands20of50(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(20, 50)
+	}
+}
+
+func TestUniqRands(t *testing.T) {
+	var x []int
+	x = UniqRands(3, 10)
+	test.Equal(t, 3, len(x))
+
+	x = UniqRands(10, 5)
+	test.Equal(t, 5, len(x))
+
+	x = UniqRands(10, 20)
+	test.Equal(t, 10, len(x))
+}