loadtest.go

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

type Config struct {
	Users       int
	URI         string
	Concurrency int
}

type Request struct{}

func (request *Request) get(config Config) (*Response, error) {
	start := time.Now()
	resp, err := http.Get(config.URI)
	if err != nil {
		fmt.Println("\nWTF")
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatalln(err)
	}

	elapsed := time.Since(start)
	return &Response{Status: resp.Status, Time: elapsed, Content: string(body)}, nil
}

type Response struct {
	Status  string
	Time    time.Duration
	Content string
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s: ", os.Args[0])
	fmt.Fprintf(os.Stderr, "Load test an API")
	flag.PrintDefaults()
}

// worker function that executes requests for user count and sends results across the channel
func worker(cid int, config Config, request Request, results chan<- Response, errors chan<- error) {
	for u := 1; u <= config.Users; u++ {

		start := time.Now()
		resp, err := http.Get(config.URI)
		if err != nil {
			fmt.Println("\nWTF")
		}

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			errors <- err
			return
		}

		elapsed := time.Since(start)
		results <- Response{Status: resp.Status, Time: elapsed, Content: string(body)}
	}
}

func main() {
	userCount := flag.Int("users", 100, "User Count")
	concurrency := flag.Int("concurrency", 100, `Number of concurrent requests per 
	user e.g. with 10 users and concurrency 100, then total of 1000 requests will be made`)
	uri := flag.String("uri", "http://localhost:8080", "URI to load test")

	flag.Usage = usage
	flag.Parse()

	// load up config and data
	config := Config{}
	config.Users = *userCount
	config.Concurrency = *concurrency
	config.URI = *uri

	request := Request{}

	// setup the response channel
	results := make(chan Response)
	errors := make(chan error)

	start := time.Now()

	for c := 1; c <= config.Concurrency; c++ {
		go worker(c, config, request, results, errors)
	}

	fmt.Printf("Start load test >>> %s\n", config.URI)
	success := 0
	failed := 0
	var totalTime time.Duration
	total := config.Concurrency * config.Users
	responses := 0
	for r := 1; r <= total; r++ {
		response := <-results
		if response.Status == "200 OK" {
			success++
		} else {
			failed++
		}
		// check if there is a response body, then increment the responses
		if strings.Contains(response.Content, "Body") {
			responses++
		}
		totalTime += response.Time
		fmt.Printf("\r%.0f%%", (float32(success+failed)/float32(total))*100)
	}

	elapsed := time.Since(start)
	fmt.Printf("\rDone!\n")
	fmt.Printf("Concurrency Level:    %d\n", config.Concurrency)
	fmt.Printf("Users:                %d\n", config.Users)
	fmt.Printf("Completed requests:   %d\n", success)
	fmt.Printf("Failed requests:      %d\n", failed)
	fmt.Printf("Total time:           %s\n", elapsed)
	fmt.Printf("Time per request:     %s\n", (totalTime / time.Duration(total)))
	ms := float64(elapsed) / float64(time.Millisecond)

	fmt.Printf("Requests per second:  %d\n", int(float64(success)/ms*1000))

	fmt.Println("-----------------------------------")
	fmt.Printf("Total sent:       %d\n", total)
	fmt.Printf("Total recv        %d\n", responses)
}