From 3f6e5c04e0bdccf2e8274ad49dd6e4166ebf5aac Mon Sep 17 00:00:00 2001
From: Vishal Raj <vishal.raj@elastic.co>
Date: Tue, 24 Oct 2023 11:24:19 +0800
Subject: [PATCH] Allow option to ignore http errors for loadgen (#42)

---
 cmd/apmbench/bench.go                      |  1 +
 internal/loadgen/config/config.go          |  2 ++
 internal/loadgen/eventhandler.go           |  2 ++
 internal/loadgen/eventhandler/handler.go   |  6 +++++-
 internal/loadgen/eventhandler/transport.go | 24 ++++++++++------------
 5 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/cmd/apmbench/bench.go b/cmd/apmbench/bench.go
index 668c02d..308ce52 100644
--- a/cmd/apmbench/bench.go
+++ b/cmd/apmbench/bench.go
@@ -202,6 +202,7 @@ func newEventHandler(tb testing.TB, p string, l *rate.Limiter) *eventhandler.Han
 		URL:               loadgencfg.Config.ServerURL.String(),
 		Token:             loadgencfg.Config.SecretToken,
 		APIKey:            loadgencfg.Config.APIKey,
+		IgnoreErrors:      loadgencfg.Config.IgnoreErrors,
 		RewriteIDs:        loadgencfg.Config.RewriteIDs,
 		RewriteTimestamps: loadgencfg.Config.RewriteTimestamps,
 		Headers:           loadgencfg.Config.Headers,
diff --git a/internal/loadgen/config/config.go b/internal/loadgen/config/config.go
index ab28c03..e0034d5 100644
--- a/internal/loadgen/config/config.go
+++ b/internal/loadgen/config/config.go
@@ -21,6 +21,7 @@ var Config struct {
 	APIKey                    string
 	Secure                    bool
 	EventRate                 RateFlag
+	IgnoreErrors              bool
 	RewriteIDs                bool
 	RewriteTimestamps         bool
 	RewriteServiceNames       bool
@@ -110,6 +111,7 @@ func init() {
 		},
 	)
 	flag.Var(&Config.EventRate, "event-rate", "Event rate in format of {burst}/{interval}. For example, 200/5s, <= 0 values evaluate to Inf (default 0/s)")
+	flag.BoolVar(&Config.IgnoreErrors, "ignore-errors", false, "Ignore HTTP errors while sending events")
 
 	rewriteNames := map[string]*bool{
 		"service.name":        &Config.RewriteServiceNames,
diff --git a/internal/loadgen/eventhandler.go b/internal/loadgen/eventhandler.go
index c2162e8..825a544 100644
--- a/internal/loadgen/eventhandler.go
+++ b/internal/loadgen/eventhandler.go
@@ -29,6 +29,7 @@ type EventHandlerParams struct {
 	APIKey                    string
 	Limiter                   *rate.Limiter
 	Rand                      *rand.Rand
+	IgnoreErrors              bool
 	RewriteIDs                bool
 	RewriteServiceNames       bool
 	RewriteServiceNodeNames   bool
@@ -56,6 +57,7 @@ func NewEventHandler(p EventHandlerParams) (*eventhandler.Handler, error) {
 		Storage:                   events,
 		Limiter:                   p.Limiter,
 		Rand:                      p.Rand,
+		IgnoreErrors:              p.IgnoreErrors,
 		RewriteIDs:                p.RewriteIDs,
 		RewriteServiceNames:       p.RewriteServiceNames,
 		RewriteServiceNodeNames:   p.RewriteServiceNodeNames,
diff --git a/internal/loadgen/eventhandler/handler.go b/internal/loadgen/eventhandler/handler.go
index e55097c..aa92e18 100644
--- a/internal/loadgen/eventhandler/handler.go
+++ b/internal/loadgen/eventhandler/handler.go
@@ -95,6 +95,10 @@ type Config struct {
 	// it must not be invoked concurrently by any other goroutines.
 	Rand *rand.Rand
 
+	// IgnoreErrors when set to true ignores HTTP errors while sending
+	// events using the event handler.
+	IgnoreErrors bool
+
 	// RewriteTimestamps controls whether event timestamps are rewritten
 	// during replay.
 	//
@@ -342,7 +346,7 @@ func (h *Handler) sendBatch(
 		if err := s.w.Close(); err != nil {
 			return err
 		}
-		if err := h.config.Transport.SendV2Events(ctx, &s.w.buf); err != nil {
+		if err := h.config.Transport.SendV2Events(ctx, &s.w.buf, h.config.IgnoreErrors); err != nil {
 			return err
 		}
 		s.w.Reset()
diff --git a/internal/loadgen/eventhandler/transport.go b/internal/loadgen/eventhandler/transport.go
index 15f499a..7eb4334 100644
--- a/internal/loadgen/eventhandler/transport.go
+++ b/internal/loadgen/eventhandler/transport.go
@@ -6,7 +6,6 @@ package eventhandler
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -37,7 +36,7 @@ func NewTransport(c *http.Client, srvURL, token, apiKey string, headers map[stri
 }
 
 // SendV2Events sends the reader contents to `/intake/v2/events` as a batch.
-func (t *Transport) SendV2Events(ctx context.Context, r io.Reader) error {
+func (t *Transport) SendV2Events(ctx context.Context, r io.Reader, ignoreErrs bool) error {
 	req, err := http.NewRequestWithContext(ctx, "POST", t.intakeV2URL, r)
 	if err != nil {
 		return err
@@ -46,27 +45,26 @@ func (t *Transport) SendV2Events(ctx context.Context, r io.Reader) error {
 	// set it to `-1` just like the agents would.
 	req.ContentLength = -1
 	req.Header = t.intakeHeaders
-	return t.sendEvents(req, r)
+	return t.sendEvents(req, r, ignoreErrs)
 }
 
-func (t *Transport) sendEvents(req *http.Request, r io.Reader) error {
+func (t *Transport) sendEvents(req *http.Request, r io.Reader, ignoreErrs bool) error {
 	res, err := t.client.Do(req)
 	if err != nil {
 		return err
 	}
 	defer res.Body.Close()
 
-	switch res.StatusCode {
-	case http.StatusOK, http.StatusAccepted:
-		return nil
+	if !ignoreErrs {
+		switch res.StatusCode / 100 {
+		case 4:
+			return fmt.Errorf("unexpected client error: %d", res.StatusCode)
+		case 5:
+			return fmt.Errorf("unexpected server error: %d", res.StatusCode)
+		}
 	}
 
-	msg := fmt.Sprintf("unexpected apm server response %d", res.StatusCode)
-	b, err := io.ReadAll(res.Body)
-	if err != nil {
-		return errors.New(msg)
-	}
-	return fmt.Errorf(msg+": %s", string(b))
+	return nil
 }
 
 func getAuthHeader(token string, apiKey string) string {