Repository: atlassian/gostatsd
Pull request #261 (2019-09-17 05:29:02)

Adding support for custom headers when sending requests.
In order to forward requests to different target groups, adding the ability to add additional custom headers gives that opportunity at no cost. 

Since we do not have cluster ready yet, running a singular instance has proven problematic since we can't reason who is sending us data and rely on IP addresses to investigate further. Adding custom headers allows us to inspect these at the ALB and shard noisy neighbours onto their own deployment stack.


MovieStoreGuy: I believe everything should be cleaned up @aelse , would you mind giving it a once threw? 

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b67acf9..e8e81a1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+18.0.0
+------
+- Support for `custom-headers` for http_forwarder_v2
+
 17.2.0
 ------
 - Support `log-raw-metric` configuration
diff --git a/README.md b/README.md
index 79dad7e..331c60f 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,8 @@ following configuration options:
 - `flush-interval`: duration for how long to batch metrics before flushing. Should be an order of magnitude less than
   the upstream flush interval. Defaults to `1s`.
 - `log-raw-metric`: print metrics received from network to stdout in JSON format.
+- `custom-headers` : a map of strings that are added to each request sent to allow for additional network routing / request inspection. 
+  Not required, default is empty. Example: `--custom-headers='{"region" : "us-east-1", "service" : "event-producer"}'`
 
 Configuring HTTP servers
 ------------------------
diff --git a/pkg/statsd/handler_http_forwarder_v2.go b/pkg/statsd/handler_http_forwarder_v2.go
index 1077532..be48a48 100644
--- a/pkg/statsd/handler_http_forwarder_v2.go
+++ b/pkg/statsd/handler_http_forwarder_v2.go
@@ -55,6 +55,7 @@ type HttpForwarderHandlerV2 struct {
 	consolidatedMetrics   <-chan []*gostatsd.MetricMap
 	eventWg               sync.WaitGroup
 	compress              bool
+	headers               map[string]string
 }
 
 // NewHttpForwarderHandlerV2FromViper returns a new http API client.
@@ -81,11 +82,12 @@ func NewHttpForwarderHandlerV2FromViper(logger logrus.FieldLogger, v *viper.Vipe
 		subViper.GetDuration("client-timeout"),
 		subViper.GetDuration("max-request-elapsed-time"),
 		subViper.GetDuration("flush-interval"),
+		subViper.GetStringMapString("custom-headers"),
 	)
 }
 
 // NewHttpForwarderHandlerV2 returns a new handler which dispatches metrics over http to another gostatsd server.
-func NewHttpForwarderHandlerV2(logger logrus.FieldLogger, apiEndpoint, network string, consolidatorSlots, maxRequests int, compress, enableHttp2 bool, clientTimeout, maxRequestElapsedTime time.Duration, flushInterval time.Duration) (*HttpForwarderHandlerV2, error) {
+func NewHttpForwarderHandlerV2(logger logrus.FieldLogger, apiEndpoint, network string, consolidatorSlots, maxRequests int, compress, enableHttp2 bool, clientTimeout, maxRequestElapsedTime time.Duration, flushInterval time.Duration, xheaders map[string]string) (*HttpForwarderHandlerV2, error) {
 	if apiEndpoint == "" {
 		return nil, fmt.Errorf("api-endpoint is required")
 	}
@@ -117,6 +119,20 @@ func NewHttpForwarderHandlerV2(logger logrus.FieldLogger, apiEndpoint, network s
 		"flush-interval":           flushInterval,
 	}).Info("created HttpForwarderHandler")
 
+	// Default set of headers used for the forwarder
+	// Once these values are set, modifying the map is illadvised
+	// due to the fact that map is just a reference to memory.
+	headers := map[string]string{
+		"Content-Type": "application/x-protobuf",
+		"User-Agent":   "gostatsd (http forwarder)",
+	}
+
+	// Adding extra headers to the default block of headers to emit.
+	for k, v := range xheaders {
+		k = http.CanonicalHeaderKey(k)
+		headers[k] = v
+	}
+
 	dialer := &net.Dialer{
 		Timeout:   5 * time.Second,
 		KeepAlive: 30 * time.Second,
@@ -162,6 +178,7 @@ func NewHttpForwarderHandlerV2(logger logrus.FieldLogger, apiEndpoint, network s
 			Transport: transport,
 			Timeout:   clientTimeout,
 		},
+		headers: headers,
 	}, nil
 }
 
@@ -359,6 +376,18 @@ func (hfh *HttpForwarderHandlerV2) post(ctx context.Context, message proto.Messa
 	}
 }
 
+// debug rendering
+/*
+func (hh *HttpForwarderHandlerV2) serializeText(message proto.Message) ([]byte, error) {
+       buf := &bytes.Buffer{}
+       err := proto.MarshalText(buf, message)
+       if err != nil {
+               return nil, err
+       }
+       return buf.Bytes(), nil
+}
+*/
+
 func (hfh *HttpForwarderHandlerV2) serialize(message proto.Message) ([]byte, error) {
 	buf, err := proto.Marshal(message)
 	if err != nil {
@@ -368,18 +397,6 @@ func (hfh *HttpForwarderHandlerV2) serialize(message proto.Message) ([]byte, err
 	return buf, nil
 }
 
-// debug rendering
-/*
-func (hh *HttpForwarderHandlerV2) serializeText(message proto.Message) ([]byte, error) {
-	buf := &bytes.Buffer{}
-	err := proto.MarshalText(buf, message)
-	if err != nil {
-		return nil, err
-	}
-	return buf.Bytes(), nil
-}
-*/
-
 func (hfh *HttpForwarderHandlerV2) serializeAndCompress(message proto.Message) ([]byte, error) {
 	raw, err := hfh.serialize(message)
 	if err != nil {
@@ -419,19 +436,15 @@ func (hfh *HttpForwarderHandlerV2) constructPost(ctx context.Context, logger log
 	}
 
 	return func() error {
-		headers := map[string]string{
-			"Content-Type":     "application/x-protobuf",
-			"Content-Encoding": encoding,
-			"User-Agent":       "gostatsd (http forwarder)",
-		}
 		req, err := http.NewRequest("POST", path, bytes.NewReader(body))
 		if err != nil {
 			return fmt.Errorf("unable to create http.Request: %v", err)
 		}
 		req = req.WithContext(ctx)
-		for header, v := range headers {
+		for header, v := range hfh.headers {
 			req.Header.Set(header, v)
 		}
+		req.Header.Set("Content-Encoding", encoding)
 		resp, err := hfh.client.Do(req)
 		if err != nil {
 			return fmt.Errorf("error POSTing: %v", err)


Test patch:
diff --git a/pkg/web/http_receiver_v2_test.go b/pkg/web/http_receiver_v2_test.go
index f5ac86b..d4667c9 100644
--- a/pkg/web/http_receiver_v2_test.go
+++ b/pkg/web/http_receiver_v2_test.go
@@ -50,6 +50,7 @@ func TestForwardingEndToEndV2(t *testing.T) {
 		10*time.Second,
 		10*time.Second,
 		10*time.Millisecond,
+		nil,
 	)
 	require.NoError(t, err)