Stress test demo code for charm.io

stressballer.go

package main

import (
	"context"
	"flag"
	"fmt"
	"sync"
	"time"

	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/chromedp"
)

type Result struct {
	Workerid	int
	Document	string
	Start	time.Time
	Stop	time.Time
	Error error
}

type Test struct {
	URL string
	Cookies string
	N int
	Delay time.Duration
	FailureDuration time.Duration
	Results map[int]Result
	mu sync.Mutex
}

func NewTest(url, cookies string, n int, delay, failureDuration time.Duration) *Test {
	return &Test{
		URL: url,
		Cookies: cookies,
		N: n,
		Delay: delay,
		FailureDuration: failureDuration,
		Results: make(map[int]Result),
	}
}

func (t Test) Run() error {
	var (
		wg sync.WaitGroup
	)
	
	for i:=1;i<t.N+1;i++ {
		<-time.After(t.Delay);
		wg.Add(1);
		go func(num int) {
			tick := time.Now();
			defer wg.Done()
			// 30 second timeout
			parent, cancel := context.WithTimeout(context.Background(), t.FailureDuration);
			defer cancel()
			ctx, cancel := chromedp.NewContext(parent);
			defer cancel()
	
			var page string
			err := chromedp.Run(ctx,
				setHeaders(
					t.URL,
					map[string]interface{}{
						"Cookie": t.Cookies,
					},
					&page,
				),
			)
			if err != nil {
				fmt.Printf("worker %d failed: %v\n", num, err);
			}
			t.mu.Lock()
			t.Results[num] = Result{
				Workerid: num,
				Document: page,
				Start: tick,
				Stop: time.Now(),
				Error: err,
			}
			fmt.Printf("worker %d took %v seconds (%d/%d)\n", num, time.Now().Sub(tick).Round(100*time.Millisecond).Seconds(),len(t.Results),t.N);
			t.mu.Unlock()
		}(i)
	}
	wg.Wait()
	return nil
}

// setHeaders returns a task list that sets the passed headers.
func setHeaders(host string, headers map[string]interface{}, res *string) chromedp.Tasks {
	return chromedp.Tasks{
		network.Enable(),
		network.SetExtraHTTPHeaders(network.Headers(headers)),
		chromedp.Navigate(host),
		chromedp.OuterHTML("html", res),
	}
}

func main() {
	var (
		//  ---- production machine ----
		url = flag.String("url", "https://domain.to.test/brands/some_brand", "URL to load")
		authcookies = flag.String("domain.to.test", "csrftoken=SUPERSECRETTOKEN;sessionid=SUPERSECRETSESSIONID;","value of the csrftoken cookie obtained from a valid request")
		n = flag.Int("n", 20, "number of simultaneous requests")
		delay = flag.Duration("delay", time.Millisecond*100, "delay between requests")
		failureDuration = flag.Duration("failAfter", 15*time.Second, "amount of time to wait before indicating failure")
	)
	flag.Parse()
	t := NewTest(*url, *authcookies, *n, *delay, *failureDuration)
	t.Run()
	var errCnt int
	for _, result := range t.Results {
		if result.Error != nil {
			errCnt++
		}
	}
	fmt.Printf("test encountered %d errors", errCnt)
}