From 2526e0993628f8ee530a5289ee902ca57c0330bb Mon Sep 17 00:00:00 2001
From: Matthew Rothenberg <mrothenberg@gmail.com>
Date: Fri, 19 Jun 2020 20:52:37 -0400
Subject: [PATCH 1/2] bench: add parallel perf test for lock contention

we expect to see lock contention stemming from the usage of global rand.
---
 weightedrand_test.go | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/weightedrand_test.go b/weightedrand_test.go
index e444ea8..020ed90 100644
--- a/weightedrand_test.go
+++ b/weightedrand_test.go
@@ -127,6 +127,21 @@ func BenchmarkPick(b *testing.B) {
 	}
 }
 
+func BenchmarkPickParallel(b *testing.B) {
+	for n := BMminChoices; n <= BMmaxChoices; n *= 10 {
+		b.Run(strconv.Itoa(n), func(b *testing.B) {
+			choices := mockChoices(n)
+			chooser := NewChooser(choices...)
+			b.ResetTimer()
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					chooser.Pick()
+				}
+			})
+		})
+	}
+}
+
 func mockChoices(n int) []Choice {
 	choices := make([]Choice, 0, n)
 	for i := 0; i < n; i++ {

From 7a6357de7f904df3c725af84a533b17cc634616f Mon Sep 17 00:00:00 2001
From: Matthew Rothenberg <mrothenberg@gmail.com>
Date: Fri, 19 Jun 2020 21:32:09 -0400
Subject: [PATCH 2/2] feat: parallel safe Pick to avoid rand contention

---
 weightedrand.go      | 18 ++++++++++++++++++
 weightedrand_test.go | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/weightedrand.go b/weightedrand.go
index 6647f47..f3edbc8 100644
--- a/weightedrand.go
+++ b/weightedrand.go
@@ -49,8 +49,26 @@ func NewChooser(cs ...Choice) Chooser {
 }
 
 // Pick returns a single weighted random Choice.Item from the Chooser.
+//
+// Utilizes global rand as the source of randomness -- you will likely want to
+// seed it.
 func (chs Chooser) Pick() interface{} {
 	r := rand.Intn(chs.max) + 1
 	i := sort.SearchInts(chs.totals, r)
 	return chs.data[i].Item
 }
+
+// PickSource returns a single weighted random Choice.Item from the Chooser,
+// utilizing the provided *rand.Rand source rs for randomness.
+//
+// The primary use-case for this is avoid lock contention from the global random
+// source if utilizing Chooser(s) from multiple goroutines in extremely
+// high-throughput situations.
+//
+// It is the responsibility of the caller to ensure the provided rand.Source is
+// safe from thread safety issues.
+func (chs Chooser) PickSource(rs *rand.Rand) interface{} {
+	r := rs.Intn(chs.max) + 1
+	i := sort.SearchInts(chs.totals, r)
+	return chs.data[i].Item
+}
diff --git a/weightedrand_test.go b/weightedrand_test.go
index 020ed90..0b7ed30 100644
--- a/weightedrand_test.go
+++ b/weightedrand_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math/rand"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 )
@@ -60,6 +61,35 @@ func TestChooser_Pick(t *testing.T) {
 	verifyFrequencyCounts(t, counts, choices)
 }
 
+// TestChooser_PickSource is the same test methodology as TestChooser_Pick, but
+// here we use the PickSource method and access the same chooser concurrently
+// from multiple different goroutines, each providing its own source of
+// randomness.
+func TestChooser_PickSource(t *testing.T) {
+	choices := mockFrequencyChoices(t, testChoices)
+	chooser := NewChooser(choices...)
+	t.Log("totals in chooser", chooser.totals)
+
+	counts1 := make(map[int]int)
+	counts2 := make(map[int]int)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	checker := func(counts map[int]int) {
+		defer wg.Done()
+		rs := rand.New(rand.NewSource(time.Now().UTC().UnixNano()))
+		for i := 0; i < testIterations/2; i++ {
+			c := chooser.PickSource(rs)
+			counts[c.(int)]++
+		}
+	}
+	go checker(counts1)
+	go checker(counts2)
+	wg.Wait()
+
+	verifyFrequencyCounts(t, counts1, choices)
+	verifyFrequencyCounts(t, counts2, choices)
+}
+
 // Similar to what is used in randutil test, but in randomized order to avoid
 // any issues with algorithms that are accidentally dependant on presorted data.
 func mockFrequencyChoices(t *testing.T, n int) []Choice {
@@ -134,8 +164,9 @@ func BenchmarkPickParallel(b *testing.B) {
 			chooser := NewChooser(choices...)
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
+				rs := rand.New(rand.NewSource(time.Now().UTC().UnixNano()))
 				for pb.Next() {
-					chooser.Pick()
+					chooser.PickSource(rs)
 				}
 			})
 		})