diff --git a/pkg/prand/prand.go b/pkg/prand/prand.go new file mode 100644 index 0000000..5dd6a64 --- /dev/null +++ b/pkg/prand/prand.go @@ -0,0 +1,75 @@ +// Package prand allows parallel access to randomness based on indices or +// infohashes. +package prand + +import ( + "encoding/binary" + "math/rand" + "sync" + "time" + + "github.com/chihaya/chihaya/bittorrent" +) + +type lockableRand struct { + *rand.Rand + *sync.Mutex +} + +// Container is a container for sources of random numbers that can be locked +// individually. +type Container struct { + rands []lockableRand +} + +// NewSeeded returns a new Container with num sources that are seeeded with +// seed. +func NewSeeded(num int, seed int64) *Container { + toReturn := Container{ + rands: make([]lockableRand, num), + } + + for i := 0; i < num; i++ { + toReturn.rands[i].Rand = rand.New(rand.NewSource(seed)) + toReturn.rands[i].Mutex = &sync.Mutex{} + } + + return &toReturn +} + +// New returns a new Container with num sources that are seeded with the current +// time. +func New(num int) *Container { + return NewSeeded(num, time.Now().UnixNano()) +} + +// Get locks and returns the nth source. +// +// Get panics if n is not a valid index for this Container. +func (s *Container) Get(n int) *rand.Rand { + r := s.rands[n] + r.Lock() + return r.Rand +} + +// GetByInfohash locks and returns a source derived from the infohash. +func (s *Container) GetByInfohash(ih bittorrent.InfoHash) *rand.Rand { + u := int(binary.BigEndian.Uint32(ih[:4])) % len(s.rands) + return s.Get(u) +} + +// Return returns the nth source to be available again. +// +// Return panics if n is not a valid index for this Container. +// Return also panics if the nth source is unlocked already. +func (s *Container) Return(n int) { + s.rands[n].Unlock() +} + +// ReturnByInfohash returns the source derived from the infohash. +// +// ReturnByInfohash panics if the source is unlocked already. +func (s *Container) ReturnByInfohash(ih bittorrent.InfoHash) { + u := int(binary.BigEndian.Uint32(ih[:4])) % len(s.rands) + s.Return(u) +} diff --git a/pkg/prand/prand_test.go b/pkg/prand/prand_test.go new file mode 100644 index 0000000..2a550dd --- /dev/null +++ b/pkg/prand/prand_test.go @@ -0,0 +1,25 @@ +package prand + +import ( + "math/rand" + "sync/atomic" + "testing" +) + +func BenchmarkContainer_GetReturn(b *testing.B) { + c := New(1024) + a := uint64(0) + + b.ResetTimer() + b.RunParallel(func(p *testing.PB) { + i := int(atomic.AddUint64(&a, 1)) + var r *rand.Rand + + for p.Next() { + r = c.Get(i) + c.Return(i) + } + + _ = r + }) +}