load.go

package main

import (
  "flag"
  "fmt"
  "io/ioutil"
  "net/http"
  "runtime"
  "sync"
  "time"
)

var help bool
var count int
var concurrent int
var url string

var client *http.Client

func init() {
  client = &http.Client{}

  flag.BoolVar(&help, "h", false, "help")
  flag.IntVar(&count, "n", 1000, "number of requests")
  flag.IntVar(&concurrent, "c", runtime.NumCPU() + 1, "number of concurrent requests")
  flag.StringVar(&url, "u", "http://127.0.0.1:5000/", "url")
  flag.Parse()
}

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

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

  runChan := make(chan int, concurrent)
  resultChan := make(chan Result)

  var wg sync.WaitGroup

  success_cnt := 0
  failure_cnt := 0
  var durations []time.Duration
  var min_dur time.Duration
  var max_dur time.Duration

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

    // setup to handle responses
    go func() {
      for {
        r := <-resultChan
        durations = append(durations, r.Duration)
        min_dur = min(min_dur, r.Duration)
        max_dur = max(max_dur, r.Duration)

        // 200s and 300s are success in HTTP
        if r.StatusCode < 400 {
          success_cnt += 1
        } else {
          fmt.Printf("Error: %v; %v\n", r.StatusCode, r.ErrOrBody())
          failure_cnt += 1
        }
        wg.Done()
      }
    }()

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

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

    wg.Wait()
  })

  fmt.Printf("\n")
  fmt.Printf("Success: %v\nFailure: %v\n", success_cnt, failure_cnt)
  fmt.Printf("Min: %v\nMax: %v\n", min_dur, max_dur)
  fmt.Printf("Mean: %v\n", avg(durations))
  fmt.Printf("Elapsed time: %v\n", dur.Seconds())

}

func avg(durs []time.Duration) time.Duration {
  total := float64(0)
  for _, d := range durs {
    total += d.Seconds()
  }
  return time.Duration((total / float64(len(durs))) * float64(time.Second))
}

func min(a time.Duration, b time.Duration) time.Duration {
  if a != 0 && a < b {
    return a
  }
  return b
}

func max(a time.Duration, b time.Duration) time.Duration {
  if a > b {
    return a
  }
  return b
}

func Execute() Result {
  var resp *http.Response
  var err error
  dur := duration(func() {
    resp, err = http.Get(url)
  })

  if err != nil {
    return Result{dur, -1, err, ""}
  }
  defer resp.Body.Close()
  var body string
  if b, err := ioutil.ReadAll(resp.Body); err == nil {
    body = string(b)
  } else {
    body = ""
  }

  return Result{dur, resp.StatusCode, nil, body}
}

type Result struct {
  Duration time.Duration
  StatusCode int
  Err error
  Body string
}
func (r *Result) ErrOrBody() string {
  if nil != r.Err {
    return r.Err.Error()
  } else {
    return r.Body
  }
}

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