prand: add Container
This commit is contained in:
parent
2a4b263955
commit
79213c6bbd
2 changed files with 100 additions and 0 deletions
75
pkg/prand/prand.go
Normal file
75
pkg/prand/prand.go
Normal 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
25
pkg/prand/prand_test.go
Normal 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
|
||||
})
|
||||
}
|
Loading…
Reference in a new issue