From c49565ce52c73cd2d474f7b8cedec17d47bd12d1 Mon Sep 17 00:00:00 2001
From: Daulet Zhanguzin <dzhanguzin@groq.com>
Date: Wed, 20 Nov 2024 18:07:29 -0800
Subject: [PATCH] feat: add support for outputting errors

---
 main.go      | 13 ++++++++++++-
 report.go    |  5 +++--
 requester.go | 12 ++++++++++--
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index 2aced4a..fd0c6f2 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	_ "net/http/pprof"
@@ -47,6 +48,7 @@ var (
 
 	autoOpenBrowser = kingpin.Flag("auto-open-browser", "Specify whether auto open browser to show web charts").Bool()
 	clean           = kingpin.Flag("clean", "Clean the histogram bar once its finished. Default is true").Default("true").NegatableBool()
+	outputErrors    = kingpin.Flag("output-errors", "Output errors to file").String()
 	summary         = kingpin.Flag("summary", "Only print the summary without realtime reports").Default("false").Bool()
 	pprofAddr       = kingpin.Flag("pprof", "Enable pprof at special address").Hidden().String()
 	url             = kingpin.Arg("url", "Request url").Required().String()
@@ -226,6 +228,15 @@ func main() {
 		}
 	}
 
+	errWriter := io.Discard
+	if *outputErrors != "" {
+		errWriter, err = os.Create(*outputErrors)
+		if err != nil {
+			errAndExit(err.Error())
+			return
+		}
+	}
+
 	clientOpt := ClientOpt{
 		url:       *url,
 		method:    *method,
@@ -248,7 +259,7 @@ func main() {
 		host:        *host,
 	}
 
-	requester, err := NewRequester(*concurrency, *requests, *duration, reqRate.Limit(), &clientOpt)
+	requester, err := NewRequester(*concurrency, *requests, *duration, reqRate.Limit(), errWriter, &clientOpt)
 	if err != nil {
 		errAndExit(err.Error())
 		return
diff --git a/report.go b/report.go
index 9ee5d4b..2f4d57b 100644
--- a/report.go
+++ b/report.go
@@ -1,11 +1,12 @@
 package main
 
 import (
-	"github.com/beorn7/perks/histogram"
-	"github.com/beorn7/perks/quantile"
 	"math"
 	"sync"
 	"time"
+
+	"github.com/beorn7/perks/histogram"
+	"github.com/beorn7/perks/quantile"
 )
 
 var quantiles = []float64{0.50, 0.75, 0.90, 0.95, 0.99, 0.999, 0.9999}
diff --git a/requester.go b/requester.go
index 6790b93..e44dcb3 100644
--- a/requester.go
+++ b/requester.go
@@ -98,6 +98,7 @@ type Requester struct {
 	clientOpt   *ClientOpt
 	httpClient  *fasthttp.HostClient
 	httpHeader  *fasthttp.RequestHeader
+	errWriter   io.Writer
 
 	recordChan chan *ReportRecord
 	closeOnce  sync.Once
@@ -131,7 +132,7 @@ type ClientOpt struct {
 	host        string
 }
 
-func NewRequester(concurrency int, requests int64, duration time.Duration, reqRate *rate.Limit, clientOpt *ClientOpt) (*Requester, error) {
+func NewRequester(concurrency int, requests int64, duration time.Duration, reqRate *rate.Limit, errWriter io.Writer, clientOpt *ClientOpt) (*Requester, error) {
 	maxResult := concurrency * 100
 	if maxResult > 8192 {
 		maxResult = 8192
@@ -141,6 +142,7 @@ func NewRequester(concurrency int, requests int64, duration time.Duration, reqRa
 		reqRate:     reqRate,
 		requests:    requests,
 		duration:    duration,
+		errWriter:   errWriter,
 		clientOpt:   clientOpt,
 		recordChan:  make(chan *ReportRecord, maxResult),
 	}
@@ -260,7 +262,13 @@ func (r *Requester) DoRequest(req *fasthttp.Request, resp *fasthttp.Response, rr
 		rr.error = err.Error()
 		return
 	}
-	err = resp.BodyWriteTo(io.Discard)
+
+	writeTo := io.Discard
+	if resp.StatusCode() >= 500 {
+		writeTo = r.errWriter
+		_, _ = r.errWriter.Write([]byte(fmt.Sprintf("\n%d %s\n", resp.StatusCode(), rr.cost)))
+	}
+	err = resp.BodyWriteTo(writeTo)
 	if err != nil {
 		rr.cost = time.Since(startTime) - t1
 		rr.error = err.Error()