Go JSON HTTP load testing

load.go

// This is a program to load test the OAuth Token API in AlphaAuth
// It's written in Go to allow it to take advantage of parallelism and run many requests at the same time
// MRI Ruby doesn't handle parallelism well so isn't very appropriate for this task.

// Install Go v1.1.2
// Build with "go build load.go" to build a native binary for your platform. Go builds statically linked binaries, so you don't
// need the go runtime installed where the app is run (but you do need to build it for the target architecture)
// ./load -h for command line options

// Currently this is only good for testing oauth token API
// It shouldn't be hard to modify the code to support other API endpoints. The main LoadTester is completely independent of
// what's being tested and only needs a single payload function along with the proper URL to test an API.
// That modification to main, to switch out a different "bodyBuilder" func should be the only thing needed to support a different API.
// Maybe a map of a string[(Url, BodyBuilder func)] or similar with a command line parameter to change which one to use?

package main

import (
  "bytes"
  "encoding/json"
  "flag"
  "fmt"
  "io"
  "math/rand"
  "net/http"
  "runtime"
  "sync"
  "time"
)

var loadTester LoadTester
var help bool
var maxClientApp int

func init() {
  // Setup LoadTester and parse command line args
  loadTester = LoadTester{RequestType: "application/json"}

  flag.BoolVar(&help, "h", false, "help")
  flag.IntVar(&loadTester.Count, "n", 1000, "number of requests")
  flag.IntVar(&loadTester.Concurrent, "c", runtime.NumCPU() + 1, "number of concurrent requests")
  flag.StringVar(&loadTester.Url, "u", "http://127.0.0.1:5000/oauth/token", "url")

  // This one is specific to our OAuth implementation
  flag.IntVar(&maxClientApp, "m", 1000, "max number of client app")
  flag.Parse()
}

func main() {
  if help {
    flag.Usage()
    return
  }

  fmt.Printf("Concurrent: %v\n", loadTester.Concurrent)
  runtime.GOMAXPROCS(loadTester.Concurrent + 2)

  loadTester.RunAll(OauthBodyBuilder)
}

// Generic Stuff

type LoadTester struct {
  Count int
  Concurrent int
  Url string
  RequestType string
}

func (lt *LoadTester) RunAll(bodyBuilder func() io.Reader) {
  runChan := make(chan int, lt.Concurrent)
  resultChan := make(chan Result)

  var wg sync.WaitGroup

  success_cnt := 0
  failure_cnt := 0
  total_dur := time.Duration(0)


  // Run the stuff
  dur := duration(func() {

    // setup to handle responses
    go func() {
      for {
        r := <-resultChan
        total_dur += r.Duration
        if 200 == r.StatusCode {
          success_cnt += 1
        } else {
          fmt.Printf("Error: %v; %v\n", r.StatusCode, r.Err.Error())
          failure_cnt += 1
        }
        wg.Done()
      }
    }()

    // setup to handle running requests
    wg.Add(lt.Count)
    go func() {
      for i:=0; i < lt.Count; i++ {
        <-runChan
        fmt.Printf(".")
        go func() {
          resultChan <- lt.ExecutePost(bodyBuilder)
          runChan<- 1
        }()
      }
    }()

    // tell N number of requests to run, but this limits the concurrency
    for i := 0; i < lt.Concurrent; i ++ {
      runChan<- 1
    }

    wg.Wait()
  })


  fmt.Printf("\n")
  fmt.Printf("Success: %v\n", success_cnt)
  fmt.Printf("Failure: %v\n", failure_cnt)
  fmt.Printf("Average: %v\n", (total_dur) / (time.Duration(success_cnt + failure_cnt) * time.Millisecond))
  fmt.Printf("Elapsed time: %v\n", dur.Seconds())

}

func (lt *LoadTester) ExecutePost(bodyBuilder func() io.Reader) Result {
  var resp *http.Response
  var err error
  dur := duration(func() {
    body := bodyBuilder()
    resp, err = http.Post(fmt.Sprintf(lt.Url), lt.RequestType, body)
  })

  if err != nil {
    return Result{dur, -1, err}
  }
  defer resp.Body.Close()
  return Result{dur, resp.StatusCode, nil}

}

type Result struct {
  Duration time.Duration
  StatusCode int
  Err error
}

func duration(f func()) time.Duration {
  start := time.Now()
  f()
  return time.Now().Sub(start)
}

// OAuth Specific Stuff

func OauthBodyBuilder() io.Reader {
  n := rand.Intn(maxClientApp) + 1  // rand does 0 - 999, we want 1 - 1000
  rqJson := NewOauthReq(n).Json()
  return bytes.NewReader(rqJson)
}

type OauthReq struct {
  num int
  ClientId string `json:"client_id"`
  ClientSecret string `json:"client_secret"`
  GrantType string `json:"grant_type"`
  Scope string `json:"scope"`
  Format string `json:"format"`
}

func NewOauthReq(n int) (*OauthReq) {
  appName := fmt.Sprintf("app_%d", n)
  secret := fmt.Sprintf("secret_%d", n)
  return &OauthReq{n, appName, secret, "client_credentials", "stub", "json"}
}

func (req *OauthReq) Json() (rqJson []byte) {
  rqJson, _ = json.Marshal(req)
  return
}