Repository: atlassian/gostatsd
Pull request #454 (2022-02-01 14:00:02)

New Relic backend: Added support for throttling ("Retry-After" header)
## Description
Under heavy loads using New Relic backend I've observed GoStatsD is not handling throttling correctly. In cases when New Relic reply with `StatusCode = 429`, we should read a header, sleep (potentially) and then retry.

### New Relic documentation
**API returns a 429 response for the remainder of the minute. The response will include a Retry-After header indicating how long to wait in seconds before resubmitting or sending new data.**


## Rationale
With this PR the code will try to use "Retry-After" (if available) header to sleep instead of backoff algorithm. In case of invalid or missing header it will fallback to backoff algorithm.
This change is intended to be backward-compatible.

# Changes
- [New Relic backend] added support to handle rate-limiting from the API by handling `Retry-After` response's header.
- [New Relic backend] added new Error for rate-limit case.
- [New Relic backend] added unit-test for new feature.

malderete: @tiedotguy, @MovieStoreGuy please can I get a review on this one?
Thanks in advance!

Patch:
diff --git a/.travis.yml b/.travis.yml
index 41f9559..26bf64a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -44,6 +44,7 @@ install:
   - make setup-ci
 
 script:
+  - make check-fmt
   - make build-all
   - make coveralls || true # Failing coveralls check should not break the build
   - make check
diff --git a/Makefile b/Makefile
index c396245..bee64aa 100644
--- a/Makefile
+++ b/Makefile
@@ -20,8 +20,7 @@ PROTOBUF_VERSION ?= 3.6.1
 
 setup: setup-ci
 	go install github.com/githubnemo/CompileDaemon \
-		github.com/jstemmer/go-junit-report \
-		golang.org/x/tools/cmd/goimports
+		github.com/jstemmer/go-junit-report
 
 tools/bin/protoc:
 	curl -L -O https://github.com/protocolbuffers/protobuf/releases/download/v$(PROTOBUF_VERSION)/protoc-$(PROTOBUF_VERSION)-linux-x86_64.zip
@@ -33,7 +32,7 @@ tools/bin/protoc:
 setup-ci: tools/bin/protoc
 	go install github.com/golangci/golangci-lint/cmd/golangci-lint
 
-build-cluster: fmt
+build-cluster:
 	go build -i -v -o build/bin/$(ARCH)/cluster $(GOBUILD_VERSION_ARGS) $(CLUSTER_PKG)
 
 pb/gostatsd.pb.go: pb/gostatsd.proto
@@ -44,22 +43,27 @@ pb/gostatsd.pb.go: pb/gostatsd.proto
 		$<
 	$(RM) protoc-gen-go
 
-build: pb/gostatsd.pb.go fmt
+build: pb/gostatsd.pb.go
 	go build -i -v -o build/bin/$(ARCH)/$(BINARY_NAME) $(GOBUILD_VERSION_ARGS) $(MAIN_PKG)
 
-build-race: fmt
+build-race:
 	go build -i -v -race -o build/bin/$(ARCH)/$(BINARY_NAME) $(GOBUILD_VERSION_ARGS) $(MAIN_PKG)
 
 build-all: pb/gostatsd.pb.go
 	go install -v ./...
 
-test-all: fmt cover test-race bench bench-race check
+test-all: check-fmt cover test-race bench bench-race check
 
-test-all-full: fmt cover test-race-full bench-full bench-race-full check
+test-all-full: check-fmt cover test-race-full bench-full bench-race-full check
 
-fmt:
-	gofmt -w=true -s $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*")
-	goimports -w=true -d -local github.com/atlassian/gostatsd $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*")
+check-fmt:
+	@# Since gofmt and goimports dont return 1 on chamges, this !() stuff will trigger a build failure if theres any problems.
+	! (gofmt -l -s $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*") | grep .)
+	! (go run golang.org/x/tools/cmd/goimports -l -local github.com/atlassian/gostatsd $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*") | grep .)
+
+fix-fmt:
+	gofmt -w -s $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*")
+	go run golang.org/x/tools/cmd/goimports -w -l -local github.com/atlassian/gostatsd $$(find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./pb/*")
 
 test-full: pb/gostatsd.pb.go
 	go test ./...
diff --git a/go.mod b/go.mod
index 66a6f67..a8a53a1 100644
--- a/go.mod
+++ b/go.mod
@@ -184,3 +184,5 @@ require (
 )
 
 replace github.com/gogo/protobuf => github.com/gogo/protobuf v1.3.2
+
+replace github.com/denis-tingajkin/go-header => github.com/denis-tingaikin/go-header v0.4.2
diff --git a/go.sum b/go.sum
index 3398dae..a24a3b4 100644
--- a/go.sum
+++ b/go.sum
@@ -175,8 +175,8 @@ github.com/davecgh/go-spew v0.0.0-20161028175848-04cdfd42973b/go.mod h1:J7Y8YcW2
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/denis-tingajkin/go-header v0.4.2 h1:jEeSF4sdv8/3cT/WY8AgDHUoItNSoEZ7qg9dX7pc218=
-github.com/denis-tingajkin/go-header v0.4.2/go.mod h1:eLRHAVXzE5atsKAnNRDB90WHCFFnBUn4RN0nRcs1LJA=
+github.com/denis-tingaikin/go-header v0.4.2 h1:RZWCc6Cu6h5hgVTyy1/WMovrwgiWgstk6LjzkyxHpT8=
+github.com/denis-tingaikin/go-header v0.4.2/go.mod h1:eLRHAVXzE5atsKAnNRDB90WHCFFnBUn4RN0nRcs1LJA=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
 github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZgvJUkLughtfhJv5dyTYa91l1fOUCrgjqmcifM=
diff --git a/internal/lexer/lexer_fuzz.go b/internal/lexer/lexer_fuzz.go
index 202eaa3..fb8ab47 100644
--- a/internal/lexer/lexer_fuzz.go
+++ b/internal/lexer/lexer_fuzz.go
@@ -1,4 +1,4 @@
-// +build gofuzz
+//go:build gofuzz
 
 package lexer
 
diff --git a/metric_consolidator.go b/metric_consolidator.go
index 39923a2..8dfe306 100644
--- a/metric_consolidator.go
+++ b/metric_consolidator.go
@@ -38,21 +38,21 @@ func NewMetricConsolidator(spots int, flushInterval time.Duration, sink chan<- [
 	return mc
 }
 
-	func (mc *MetricConsolidator) Run(ctx context.Context) {
-		clck := clock.FromContext(ctx)
-		t := clck.NewTicker(mc.flushInterval)
-		defer t.Stop()
+func (mc *MetricConsolidator) Run(ctx context.Context) {
+	clck := clock.FromContext(ctx)
+	t := clck.NewTicker(mc.flushInterval)
+	defer t.Stop()
 
-		for {
-			select {
-			case <-ctx.Done():
-				mc.Flush()
-				return
-			case <-t.C:
-				mc.Flush()
-			}
+	for {
+		select {
+		case <-ctx.Done():
+			mc.Flush()
+			return
+		case <-t.C:
+			mc.Flush()
 		}
 	}
+}
 
 // Drain will collect all the MetricMaps in the MetricConsolidator and return them.
 func (mc *MetricConsolidator) Drain() []*MetricMap {
diff --git a/pkg/backends/newrelic/newrelic.go b/pkg/backends/newrelic/newrelic.go
index 953373c..0a0a74a 100644
--- a/pkg/backends/newrelic/newrelic.go
+++ b/pkg/backends/newrelic/newrelic.go
@@ -5,6 +5,7 @@ import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -129,7 +130,6 @@ type NRMetric struct {
 
 // SendMetricsAsync flushes the metrics to New Relic, preparing payload synchronously but doing the send asynchronously.
 func (n *Client) SendMetricsAsync(ctx context.Context, metrics *gostatsd.MetricMap, cb gostatsd.SendCallback) {
-
 	counter := 0
 	results := make(chan error)
 	now := float64(clock.FromContext(ctx).Now().Unix())
@@ -295,6 +295,16 @@ func (n *Client) Name() string {
 	return BackendName
 }
 
+// RetryAfterError indicates if we should wait before retrying
+type RetryAfterError struct {
+	Duration time.Duration
+}
+
+// Error returns the error message
+func (r RetryAfterError) Error() string {
+	return fmt.Sprintf("Too many request retry-after=%+v", r.Duration)
+}
+
 func (n *Client) post(ctx context.Context, buffer *bytes.Buffer, data interface{}) error {
 	post, err := n.constructPost(ctx, buffer, data)
 	if err != nil {
@@ -307,13 +317,24 @@ func (n *Client) post(ctx context.Context, buffer *bytes.Buffer, data interface{
 	b.Clock = clck
 	b.Reset()
 	b.MaxElapsedTime = n.maxRequestElapsedTime
+
+	var next time.Duration
+	var retryAfterErr *RetryAfterError
 	for {
 		if err = post(); err == nil {
 			atomic.AddUint64(&n.batchesSent, 1)
 			return nil
 		}
 
-		next := b.NextBackOff()
+		next = b.NextBackOff()
+		if next != backoff.Stop && errors.As(err, &retryAfterErr) {
+			next = time.Duration(math.Max(float64(next), float64(retryAfterErr.Duration)))
+			if n.maxRequestElapsedTime > 0 {
+				// When maxRequestElapsedTime is specified put an upper-bound limit on the duration
+				next = time.Duration(math.Min(float64(next), float64(n.maxRequestElapsedTime)))
+			}
+		}
+
 		if next == backoff.Stop {
 			atomic.AddUint64(&n.batchesDropped, 1)
 			return fmt.Errorf("[%s] %v", BackendName, err)
@@ -407,13 +428,24 @@ func (n *Client) postWrapper(ctx context.Context, json []byte, dataType string)
 		if err != nil {
 			return fmt.Errorf("error POSTing: %s", err.Error())
 		}
+
 		defer resp.Body.Close()
 		body := io.LimitReader(resp.Body, maxResponseSize)
+
 		if resp.StatusCode < http.StatusOK || resp.StatusCode > http.StatusNoContent {
 			b, _ := ioutil.ReadAll(body)
+
+			if resp.StatusCode == http.StatusTooManyRequests {
+				retryAfterSecs, err := strconv.Atoi(resp.Header.Get("Retry-After"))
+				if err == nil && retryAfterSecs > 0 {
+					return &RetryAfterError{Duration: time.Duration(retryAfterSecs) * time.Second}
+				}
+			}
+
 			n.logger.WithFields(logrus.Fields{
-				"status": resp.StatusCode,
-				"body":   string(b),
+				"status":  resp.StatusCode,
+				"body":    string(b),
+				"address": address,
 			}).Info("request failed")
 			return fmt.Errorf("received bad status code %d", resp.StatusCode)
 		}
diff --git a/tools.go b/tools.go
index 49c5e13..0e695a9 100644
--- a/tools.go
+++ b/tools.go
@@ -1,4 +1,4 @@
-// +build tools
+//go:build tools
 
 package gostatsd
 


Test patch:
diff --git a/pkg/backends/newrelic/newrelic_test.go b/pkg/backends/newrelic/newrelic_test.go
index 24888f5..b93ba96 100644
--- a/pkg/backends/newrelic/newrelic_test.go
+++ b/pkg/backends/newrelic/newrelic_test.go
@@ -80,6 +80,95 @@ func TestRetries(t *testing.T) {
 	ch <- struct{}{}
 }
 
+func TestRetryAfter(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name                  string
+		retryAfterHeader      string
+		maxRequestElapsedTime time.Duration
+		expectedDuration      int64
+	}{
+		{
+			name:                  "Positive valid RetryAfter less than maxRequestElapsedTime",
+			retryAfterHeader:      "42",
+			maxRequestElapsedTime: 60 * time.Second,
+			expectedDuration:      42,
+		},
+		{
+			name:                  "Positive valid RetryAfter greather than maxRequestElapsedTime",
+			retryAfterHeader:      "42",
+			maxRequestElapsedTime: 2 * time.Second,
+			expectedDuration:      2,
+		},
+		{
+			name:                  "Negative invalid RetryAfter",
+			retryAfterHeader:      "-10",
+			maxRequestElapsedTime: 2 * time.Second,
+			expectedDuration:      0,
+		},
+		{
+			name:                  "Non-numeric invalid RetryAfter",
+			retryAfterHeader:      "wrong",
+			maxRequestElapsedTime: 2 * time.Second,
+			expectedDuration:      0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var requestNum uint32
+			var requestTimes []int64
+
+			mux := http.NewServeMux()
+			clck := clock.NewMock(time.Unix(0, 0))
+
+			mux.HandleFunc("/v1/data", func(w http.ResponseWriter, r *http.Request) {
+				requestTimes = append(requestTimes, clck.Now().Unix())
+				defer r.Body.Close()
+				n := atomic.AddUint32(&requestNum, 1)
+				data, err := ioutil.ReadAll(r.Body)
+
+				assert.NoError(t, err)
+				assert.NotEmpty(t, data)
+				if n == 1 {
+					// Return error on first request to trigger a retry after a second
+					w.Header().Add("Retry-After", tt.retryAfterHeader)
+					w.WriteHeader(http.StatusTooManyRequests)
+				}
+			})
+			ts := httptest.NewServer(mux)
+			defer ts.Close()
+
+			v := viper.New()
+			v.SetDefault("transport.default.client-timeout", 1*time.Second)
+			p := transport.NewTransportPool(logrus.New(), v)
+
+			client, err := NewClient("default", ts.URL+"/v1/data", "", "GoStatsD", "", "", "", "metric_name", "metric_type",
+				"metric_per_second", "metric_value", "samples_min", "samples_max", "samples_count",
+				"samples_mean", "samples_median", "samples_std_dev", "samples_sum", "samples_sum_squares", "agent",
+				defaultMetricsPerBatch, defaultMaxRequests, tt.maxRequestElapsedTime, 1*time.Second, gostatsd.TimerSubtypes{}, logrus.New(), p)
+
+			require.NoError(t, err)
+			res := make(chan []error, 1)
+			ctx := clock.Context(context.Background(), clck)
+			ch := make(chan struct{})
+			go advanceTime(clck, ch)
+			client.SendMetricsAsync(ctx, twoCounters(), func(errs []error) {
+				res <- errs
+			})
+			errs := <-res
+			for _, err := range errs {
+				assert.NoError(t, err)
+			}
+			assert.EqualValues(t, 2, requestNum)
+			assert.Len(t, requestTimes, 2)
+			assert.Equal(t, int64(tt.expectedDuration), requestTimes[1]-requestTimes[0])
+			ch <- struct{}{}
+		})
+	}
+}
+
 func TestSendMetricsInMultipleBatches(t *testing.T) {
 	t.Parallel()
 	var requestNum uint32