func taskValue(r *rand.Rand, cmd common.Op) []byte {
if cmd == common.Set || cmd == common.Add || cmd == common.Replace {
// Random length between 1k and 10k
valLen := r.Intn(9*1024) + 1024
return common.RandData(r, valLen, true)
}
return nil
}
func cmdGenerator(tasks chan<- *common.Task, taskGens *sync.WaitGroup, numTasks int, cmd common.Op) {
r := rand.New(rand.NewSource(common.RandSeed()))
for i := 0; i < numTasks; i++ {
task := taskPool.Get().(*common.Task)
task.Cmd = cmd
task.Key = common.RandData(r, f.KeyLength, false)
task.Value = taskValue(r, cmd)
tasks <- task
}
taskGens.Done()
}
func cmdGenerator(tasks chan<- *common.Task, taskGens *sync.WaitGroup, numTasks int, cmd common.Op) {
r := rand.New(rand.NewSource(common.RandSeed()))
for i := 0; i < numTasks; i++ {
tasks <- &common.Task{
Cmd: cmd,
Key: common.RandData(r, f.KeyLength, false),
Value: taskValue(r, cmd),
}
}
taskGens.Done()
}
func main() {
var prot common.Prot
if f.Binary {
var b binprot.BinProt
prot = b
} else {
var t textprot.TextProt
prot = t
}
wg := &sync.WaitGroup{}
wg.Add(f.NumWorkers)
for i := 0; i < f.NumWorkers; i++ {
go func(prot common.Prot, wg *sync.WaitGroup) {
conn, err := common.Connect("localhost", f.Port)
if err != nil {
panic("Couldn't connect")
}
r := rand.New(rand.NewSource(common.RandSeed()))
rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
// 0 to 100k data
for i := 0; i < 102400; i++ {
key := common.RandData(r, f.KeyLength, false)
value := common.RandData(nil, i, true)
prot.Set(rw, key, value)
prot.Get(rw, key)
}
fmt.Println("Done.")
wg.Done()
}(prot, wg)
}
wg.Wait()
}
func worker(numops int, prot common.Prot, wg *sync.WaitGroup) {
conn, err := common.Connect(f.Host, f.Port)
if err != nil {
panic(err.Error())
}
rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
r := rand.New(rand.NewSource(time.Now().UnixNano()))
for i := 0; i < numops; i++ {
curOpCount := atomic.AddUint64(opCount, 1)
if curOpCount%10000 == 0 {
log.Println(curOpCount)
}
// Generate key anew each time
key := common.RandData(r, f.KeyLength, false)
// Value size between 5k and 20k
valLen := r.Intn(15*1024) + 5*1024
value := common.RandData(nil, valLen, true)
// continue on even if there's errors here
if err := prot.Set(rw, key, value); err != nil {
log.Println("Error during set:", err.Error())
conn.Close()
conn, err = common.Connect(f.Host, f.Port)
if err != nil {
panic(err.Error())
}
rw = bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
}
}
wg.Done()
}
func worker(prot common.Prot, rw *bufio.ReadWriter, keys chan []byte, wg *sync.WaitGroup) {
r := rand.New(rand.NewSource(time.Now().UnixNano()))
for key := range keys {
curOpCount := atomic.AddUint64(opCount, 1)
if curOpCount%10000 == 0 {
log.Println(curOpCount)
}
valLen := r.Intn(9*1024) + 1024
value := common.RandData(nil, valLen, true)
// continue on even if there's errors here
if err := prot.Set(rw, key, value); err != nil {
log.Println("Error during set:", err.Error())
if err == io.EOF {
log.Println("End of file. Aborting!")
wg.Done()
return
}
}
opaque := r.Int()
// pass the test if the data matches
ret, err := prot.GetWithOpaque(rw, key, opaque)
if err != nil {
log.Println("Error getting data for key", string(key), ":", err.Error())
if err == io.EOF {
log.Println("End of file. Aborting!")
wg.Done()
return
}
continue
}
if !bytes.Equal(value, ret) {
log.Println("Data returned from server does not match!",
"\nData len sent:", len(value),
"\nData len recv:", len(ret))
//log.Println("sent:", string(value))
//log.Println("got:", string(ret))
}
}
wg.Done()
}