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)) +}