couchbase-bench.go

package main

import (
 "github.com/couchbase/go-couchbase"
 "flag"
 "log"
 "math/rand"
 "fmt"
 "time"
 "strconv"
 "sync"
)

type User struct {
   Age int `json:"age"`
   Name string `json:"name"`
   Email string `json:"email"`
   Password string `json:"password"`
   CityId int `json:"city_id"`
}

func worker(linkChan chan int64, wg *sync.WaitGroup, processor func(int64)) {
  // Decreasing internal counter for wait-group as soon as goroutine finishes
  defer wg.Done()

  for id := range linkChan {
    processor(id)
  }

}

func getUser() User {
   return User{rand.Intn(90)+1, "Some really long name, like the wind", "thewind@gmail.com", "jfajoweihf0283048hg", rand.Intn(1e5)+1}
}

const BUCKET = "__test_bucket__"

func Set(numDocs int, threads int, bucket *couchbase.Bucket) {

  lCh := make(chan int64)
  wg := new(sync.WaitGroup)

  start := time.Now().UnixNano() / int64(time.Millisecond)

  // Adding routines to workgroup and running then
  for i := 0; i < threads; i++ {
    wg.Add(1)
    go worker(lCh, wg, func(id int64) {
       err := bucket.Set("users:" + strconv.FormatInt(id, 10), 0, getUser())
       if err != nil {
          log.Println(err)
       }
    })
  }

  // var wg sync.WaitGroup
  // Processing all links by spreading them to `free` goroutines
  count := int64(0)

  for i := 0; i < numDocs; i++ {
    count += 1
    lCh <- count
  }

  // Closing channel (waiting in goroutines won't continue any more)
  close(lCh)

  // Waiting for all goroutines to finish (otherwise they die as main routine dies)
  wg.Wait()

  end := time.Now().UnixNano() / int64(time.Millisecond)
  log.Println( fmt.Sprintf("%1.2f insert OPS (ops/second)", float64(1000) * (float64(count) / (float64(end) - float64(start))) ))


}

func Get(numDocs int, threads int, bucket *couchbase.Bucket) {

  lCh := make(chan int64)
  wg := new(sync.WaitGroup)

  start := time.Now().UnixNano() / int64(time.Millisecond)

  // Adding routines to workgroup and running then
  for i := 0; i < threads; i++ {
    wg.Add(1)
    go worker(lCh, wg, func(id int64) {
       user := new(User)
       key := "users:" + fmt.Sprintf("%d", id)
       err := bucket.Get(key, &user)

       if err != nil {
          log.Println(key, err)
       }
    })
  }

  // var wg sync.WaitGroup
  // Processing all links by spreading them to `free` goroutines
  count := int64(0)

  for i := 0; i < numDocs; i++ {
    count += 1
    lCh <- int64(rand.Intn(numDocs))
  }

  // Closing channel (waiting in goroutines won't continue any more)
  close(lCh)

  // Waiting for all goroutines to finish (otherwise they die as main routine dies)
  wg.Wait()

  end := time.Now().UnixNano() / int64(time.Millisecond)
  log.Println( fmt.Sprintf("%1.2f GET OPS (ops/second)", float64(1000) * (float64(count) / (float64(end) - float64(start))) ))

}

func main() {
  c, err := couchbase.Connect("http://localhost:8091/")
  if err != nil {
      log.Fatalf("Error connecting:  %v", err)
  }

  pool, err := c.GetPool("default")
  if err != nil {
      log.Fatalf("Error getting pool:  %v", err)
  }

  bucket, err := pool.GetBucket(BUCKET)
  if err != nil {
      log.Fatalf("Error getting bucket:  %v", err)
  }

  requests := flag.Int("r", 20000, "number of requests to make")
  threads := flag.Int("c", 20, "number of threads")
  doGet := flag.Bool("get", false, "Get test to Couchbase")

  flag.Parse()

  log.Println(fmt.Sprintf("Making %d requests using %d threads.", *requests, *threads))

  Set(*requests, *threads, bucket)

  if *doGet {
    log.Println("Making get requests")
    Get(*requests, *threads, bucket)
  }

}