package main
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"log"
"net/http"
"os"
"strconv"
"sync"
"sync/atomic"
"time"
)
var query string = `
query {
me {
id
canonicalName
sshKeys {
cursor
results {
fingerprint
}
}
}
}
`
var errors sync.Map
func doReqs(wg *sync.WaitGroup, tok string, nsuccess, nfail *int64, duration *int64) {
defer wg.Done()
start := time.Now()
defer func() {
elapsed := time.Now().Sub(start)
atomic.AddInt64(duration, elapsed.Milliseconds())
}()
client := http.Client{}
body, err := json.Marshal(&struct{
Query string `json:"query"`
}{
Query: query,
})
if err != nil {
panic(err)
}
req, err := http.NewRequest("POST", "http://localhost:5100/query",
bytes.NewReader(body))
if err != nil {
panic(err)
}
req.Header.Add("Authorization", "Bearer " + tok)
req.Header.Add("Content-Type", "application/json")
if resp, err := client.Do(req); err != nil {
atomic.AddInt64(nfail, 1)
errors.Store(fmt.Sprintf("%v", err), nil)
} else {
defer resp.Body.Close()
if resp.StatusCode == 200 {
atomic.AddInt64(nsuccess, 1)
} else {
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
return
}
errors.Store(string(body), nil)
atomic.AddInt64(nfail, 1)
}
}
}
func main() {
batch, _ := strconv.Atoi(os.Args[2])
log.Printf("Sending batch of %d requests", batch)
var nsuccess, nfail, duration int64
var wg sync.WaitGroup
wg.Add(batch)
for i := 0; i < batch; i += 1 {
go doReqs(&wg, os.Args[1], &nsuccess, &nfail, &duration)
}
time.Sleep(1 * time.Second)
wg.Wait()
log.Printf("%d success, %d failure; %dms avg",
nsuccess, nfail, duration / int64(batch))
errors.Range(func(key, value interface{}) bool {
body, _ := key.(string)
log.Printf("Sample error: %s", body)
return true
})
}