prand: add Container

This commit is contained in:
Leo Balduf 2016-11-07 13:22:27 +01:00
parent 2a4b263955
commit 79213c6bbd
2 changed files with 100 additions and 0 deletions

75
pkg/prand/prand.go Normal file
View file

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

25
pkg/prand/prand_test.go Normal file
View file

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