feat: parallel safe Pick to avoid rand contention

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

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