Initial version of the selenium stresser
1 files changed, 77 insertions(+), 115 deletions(-)

M main.go
M main.go +77 -115
@@ 1,11 1,12 @@ 
-// Run some code on play.golang.org and display the result
 package main
 
 import (
-	"bufio"
+	"encoding/base64"
 	"encoding/csv"
-	"io"
+	"fmt"
 	"log"
+	"math/rand"
+	"net/http"
 	"os"
 	"sync"
 	"time"

          
@@ 21,114 22,40 @@ func init() {
 }
 
 type Result struct {
-	URL      string
 	Duration time.Duration
 	Error    error
 }
 
-type Client struct {
-	Capabilities selenium.Capabilities
-	URL          string
-	startTime    time.Time
-	driver       selenium.WebDriver
-}
-
-func (c *Client) reset() {
-	if c.driver != nil {
-		c.driver.Quit()
-		c.driver = nil
-	}
-}
-
-func (c *Client) connect() selenium.WebDriver {
-	if c.driver == nil {
-		// Try to connect to the endpoint
-		for {
-			var err error
-
-			c.driver, err = selenium.NewRemote(c.Capabilities, c.URL)
-			if err != nil {
-				logger.Println("Error connecting to", c.URL, err)
-				time.Sleep(30 * time.Second)
-			} else {
-				c.startTime = time.Now()
-				break
-			}
-		}
-	}
-
-	return c.driver
-}
-
-//
-// Download url re-using the previous Selenium session if it ran for less than
-// 10 minutes.
-//
-func (c *Client) Get(url string) (duration time.Duration, err error) {
-	// Re-connect every 10 minutes
-	if c.driver != nil && time.Since(c.startTime) > 10*time.Minute {
-		c.reset()
-	}
-
-	var driver = c.connect()
-
-	var start = time.Now()
-	err = driver.Get(url)
-	duration = time.Since(start)
-
-	if err != nil {
-		// There was an error loading the URL, destroy the web driver just
-		// in case and report the error.
-		c.reset()
-	}
-
-	return
-}
-
 func worker(
 	caps selenium.Capabilities,
 	seleniumUrl string,
+	url string,
+	duration time.Duration,
+	output chan<- Result,
 	waitGroup *sync.WaitGroup,
-	input <-chan string,
-	output chan<- Result,
 ) {
 	defer waitGroup.Done()
 
-	var client = Client{Capabilities: caps, URL: seleniumUrl}
-
-	for url := range input {
-		var duration, err = client.Get(url)
-		output <- Result{
-			URL:      url,
-			Duration: duration,
-			Error:    err,
-		}
+	var driver, err = selenium.NewRemote(caps, seleniumUrl)
+	if err != nil {
+		logger.Println("Error connecting to", seleniumUrl, err)
+		return
 	}
 
-	client.reset()
-}
+	var startTime = time.Now()
 
-//
-// Read input line by line and send it to the returned channel. Once there's
-// nothing left to read closes the channel.
-//
-func ScanLines(input io.Reader) <-chan string {
-	var c = make(chan string)
+	for time.Since(startTime) < duration {
+		var n = time.Now()
+		err = driver.Get(url)
+		var d = time.Since(n)
 
-	go func() {
-		var scanner = bufio.NewScanner(input)
-
-		for scanner.Scan() {
-			c <- scanner.Text()
+		if err != nil {
+			output <- Result{Duration: d, Error: err}
+		} else {
+			output <- Result{Duration: d}
 		}
-
-		if err := scanner.Err(); err != nil {
-			logger.Println("reading input:", err)
-		}
-		close(c)
-	}()
-
-	return c
+	}
+	driver.Quit()
 }
 
 //

          
@@ 138,12 65,18 @@ func parse_arguments() (
 	url string,
 	capabilities selenium.Capabilities,
 	workers int,
-	progress bool,
+	listen string,
+	remoteUrl string,
+	size int,
+	duration uint,
 ) {
 	var args struct {
-		Workers      uint              `short:"w" default:"4" description:"Number of workers"`
+		Listen       string            `long:"listen" default:":8080" description:"address to listen on"`
+		RemoteURL    string            `long:"remoteurl" default:"http://localhost:8080" description:"URL the selenium side uses for download"`
+		Size         int               `long:"size" default:"1000" description:"size of the downloaded files"`
+		Duration     uint              `long:"duration" default:"60" description:"how long the stress test will run"`
+		Workers      uint              `long:"worker" default:"4" description:"Number of workers"`
 		Capabilities map[string]string `short:"c" default:"browserName:firefox" description:"Selenium capabilities"`
-		Progress     bool              `short:"p" description:"Display progress on stderr"`
 		Arg          struct {
 			URL string `description:"URL to the Selenium server"`
 		} `positional-args:"yes" required:"yes"`

          
@@ 162,7 95,10 @@ func parse_arguments() (
 
 	url = args.Arg.URL
 	workers = int(args.Workers)
-	progress = args.Progress
+	listen = args.Listen
+	remoteUrl = args.RemoteURL
+	size = args.Size
+	duration = args.Duration
 	capabilities = make(selenium.Capabilities, len(args.Capabilities))
 	for key, value := range args.Capabilities {
 		capabilities[key] = value

          
@@ 171,10 107,46 @@ func parse_arguments() (
 	return
 }
 
+const blockSize = 100
+
 func main() {
-	var url, capabilities, workers, progress = parse_arguments()
+	var url, capabilities, workers, listen, remoteUrl, size, duration = parse_arguments()
+
+	//
+	// Start HTTP server in the background
+	//
+	go func() {
+		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+			var start = time.Now()
+			// Write file size first
+			var written, err = fmt.Fprintf(w, "Size: %d\n", size)
+			if err != nil {
+				return
+			}
+			var left = size
+			left -= written
 
-	var input = ScanLines(os.Stdin)
+			// FIXME we currently write by blocks of 100 bytes, which means the
+			// size will rounded to the closest upper divisible by 100 value.
+			// For example: a file size of 301 will be rounded to 400.
+			var e = base64.RawStdEncoding
+			var randData = make([]byte, blockSize)
+			var buffer = make([]byte, e.EncodedLen(blockSize))
+			for left > 0 {
+				rand.Read(randData)
+				e.Encode(buffer, randData)
+				n, err := w.Write(buffer)
+				if err != nil {
+					return
+				}
+				left -= n
+			}
+			logger.Println("handled request in ", time.Since(start))
+		})
+
+		log.Fatal(http.ListenAndServe(listen, nil))
+	}()
+
 	// Buffer results to avoid blocking workers
 	var output = make(chan Result, 10)
 	var wg sync.WaitGroup

          
@@ 183,7 155,7 @@ func main() {
 
 	wg.Add(workers)
 	for i := 0; i < workers; i++ {
-		go worker(capabilities, url, &wg, input, output)
+		go worker(capabilities, url, remoteUrl, time.Duration(duration)*time.Second, output, &wg)
 	}
 
 	// Wait for all worker to finish, then close output to end the program

          
@@ 197,21 169,11 @@ func main() {
 	for o := range output {
 		var e = ""
 		if o.Error != nil {
-			logger.Println("Error for", o.URL, ":", o.Error)
+			logger.Println("Worker errored: ", o.Error.Error())
 			e = o.Error.Error()
 		}
-		if progress {
-			if e == "" {
-				os.Stderr.Write([]byte{'.'})
-			} else {
-				os.Stderr.Write([]byte{'x'})
-			}
-		}
-		writer.Write([]string{o.URL, e, o.Duration.String()})
+		writer.Write([]string{e, fmt.Sprintf("%d", o.Duration / time.Millisecond)})
 		writer.Flush()
-	}
-
-	if progress {
-		os.Stderr.Write([]byte{'\n'})
+		logger.Println("request took ", o.Duration)
 	}
 }