From 3d46c2fe810b0d1d87611648b7700f03c67b5f68 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tomas@soundcloud.com>
Date: Tue, 13 Aug 2013 15:01:41 +0200
Subject: [PATCH] Implement Clients with throttling and metrics

---
 client.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 main.go   | 41 +++++++++++++++++------------------------
 urls.go   | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 24 deletions(-)
 create mode 100644 client.go
 create mode 100644 urls.go

diff --git a/client.go b/client.go
new file mode 100644
index 0000000..5c79c38
--- /dev/null
+++ b/client.go
@@ -0,0 +1,49 @@
+package main
+
+import (
+	"net/http"
+	"time"
+)
+
+// Client is an http.Client with rate limiting and time series instrumentation.
+type Client struct {
+	cli      http.Client
+	qps      uint
+	codes    []uint64
+	timings  []time.Duration
+	bytesOut []int64
+	bytesIn  []int64
+}
+
+func NewClient(qps uint) *Client {
+	return &Client{
+		cli:      http.Client{},
+		qps:      qps,
+		codes:    []uint64{},
+		timings:  []time.Duration{},
+		bytesOut: []int64{},
+		bytesIn:  []int64{},
+	}
+}
+
+// Drill loops over the passed reqs channel and executes each request.
+// It is throttled to the qps specified in the initializer
+func (c *Client) Drill(reqs chan *http.Request) {
+	throttle := time.Tick(time.Duration(1e9 / c.qps))
+	for req := range reqs {
+		<-throttle
+		go c.Do(req)
+	}
+}
+
+// Do executes the passed http.Request and saves some metrics
+// (timings, bytesIn, bytesOut, codes)
+func (c *Client) Do(req *http.Request) (*http.Response, error) {
+	began := time.Now()
+	resp, err := c.cli.Do(req)
+	c.timings = append(c.timings, time.Since(began))
+	c.bytesOut = append(c.bytesOut, req.ContentLength)
+	c.bytesIn = append(c.bytesIn, resp.ContentLength)
+	c.codes[resp.StatusCode]++
+	return resp, err
+}
diff --git a/main.go b/main.go
index ae73750..486db6f 100644
--- a/main.go
+++ b/main.go
@@ -1,15 +1,12 @@
 package main
 
 import (
-	"bytes"
 	"flag"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"math"
 	"math/rand"
 	"net/http"
-	"net/url"
 	"time"
 )
 
@@ -24,13 +21,6 @@ func main() {
 		requests = flag.Uint("requests", 5000, "Number of requests to do")
 		duration = flag.Duration("duration", 10*time.Second, "Maximum duration of execution")
 	)
-
-	var (
-		urls    []*url.URL
-		clients []*http.Client
-		err     error
-	)
-
 	flag.Parse()
 
 	// Validate QPS argument
@@ -39,16 +29,23 @@ func main() {
 	}
 	// Magic formula that assumes each client can
 	// sustain 500 QPS under normal circumstances
-	clients = make([]*http.Client, int(math.Ceil(float64(*qps)/500.0)))
+	clients := make([]*Client, int(math.Ceil(float64(*qps)/500.0)))
+	qpsClient := *qps / uint(len(clients))
+	for i := 0; i < len(clients); i++ {
+		clients[i] = NewClient(qpsClient)
+	}
 
 	// Parse URLs file
-	if urls, err = readURLsFromFile(*urlsFile); err != nil {
+	urls, err := NewURLsFromFile(*urlsFile)
+	if err != nil {
 		log.Fatal(err)
 	}
 
 	// Parse mode argument
+	random := false
 	if *mode == "random" {
 		rand.Seed(time.Now().UnixNano())
+		random = true
 	} else if *mode != "sequential" {
 		log.Fatal("Unknown mode %s", *mode)
 	}
@@ -60,21 +57,17 @@ func main() {
 
 	fmt.Printf("Hitting %d URLs in %s mode for %s with %d requests and %d clients.",
 		len(urls), *mode, duration.String(), *requests, len(clients))
-}
 
-func readURLsFromFile(filename string) ([]*url.URL, error) {
-	lines, err := ioutil.ReadFile(filename)
-	if err != nil {
-		return []*url.URL{}, err
+	reqs := make(chan *http.Request, *requests)
+	for _, client := range clients {
+		go client.Drill(reqs)
 	}
-
-	var urls []*url.URL
-	for _, line := range bytes.Split(lines, []byte("\n")) {
-		uri, err := url.Parse(string(line))
+	for _, index := range urls.Iter(random) {
+		url := urls[index]
+		req, err := http.NewRequest("GET", url.String(), nil)
 		if err != nil {
-			return []*url.URL{}, fmt.Errorf("Failed to parse URI (%s): %s", line, err)
+			log.Fatal("Bad request: %s", err)
 		}
-		urls = append(urls, uri)
+		reqs <- req
 	}
-	return urls, nil
 }
diff --git a/urls.go b/urls.go
new file mode 100644
index 0000000..f1207b7
--- /dev/null
+++ b/urls.go
@@ -0,0 +1,38 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"math/rand"
+	"net/url"
+)
+
+type URLs []*url.URL
+
+func NewURLsFromFile(filename string) (urls URLs, err error) {
+	lines, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return URLs{}, err
+	}
+
+	for _, line := range bytes.Split(lines, []byte("\n")) {
+		uri, err := url.Parse(string(line))
+		if err != nil {
+			return URLs{}, fmt.Errorf("Failed to parse URI (%s): %s", line, err)
+		}
+		urls = append(urls, uri)
+	}
+	return urls, nil
+}
+
+func (urls URLs) Iter(random bool) []int {
+	if random {
+		return rand.Perm(len(urls))
+	}
+	iter := make([]int, len(urls))
+	for i := 0; i < len(urls); i++ {
+		iter[i] = i
+	}
+	return iter
+}