Add test and clean up extract code

Signed-off-by: Michel Hollands <[email protected]>

Author: MichelHollands

pkg/distributor/distributor.go

@@ -18,7 +18,6 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/instrument"
 	"github.com/weaveworks/common/user"
-	grpc_metadata "google.golang.org/grpc/metadata"
 
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	ingester_client "github.com/cortexproject/cortex/pkg/ingester/client"
@@ -541,7 +540,7 @@ func (d *Distributor) Push(ctx context.Context, req *client.WriteRequest) (*clie
 		}
 
 		// Get clientIP(s) from Context and add it to localCtx
-		localCtx = grpc_metadata.AppendToOutgoingContext(localCtx, util.IPAddressesKey, source)
+		localCtx = util.AddSourceToOutgoingContext(localCtx, source)
 
 		return d.send(localCtx, ingester, timeseries, metadata, req.Source)
 	}, func() { client.ReuseSlice(req.Timeseries) })

pkg/ingester/ingester_test.go

@@ -30,6 +30,7 @@ import (
 	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/ring"
+	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/chunkcompat"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/test"
@@ -468,8 +469,10 @@ func TestIngesterAppendOutOfOrderAndDuplicate(t *testing.T) {
 	require.Equal(t, errResp.code, 400)
 
 	// Same timestamp as previous sample, but different value.
+	ctx = util.AddSourceToIncomingContext(ctx, "1.2.3.4")
 	err = ing.append(ctx, userID, m, 1, 1, client.API, nil)
 	require.Contains(t, err.Error(), "sample with repeated timestamp but different value")
+	require.Contains(t, err.Error(), "1.2.3.4")
 	errResp, ok = err.(*validationError)
 	require.True(t, ok)
 	require.Equal(t, errResp.code, 400)

pkg/util/extract_forwarded.go

@@ -9,8 +9,8 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-// IPAddressesKey is key for the GRPC metadata where the IP addresses are stored
-const IPAddressesKey = "ipaddresseskey"
+// ipAddressesKey is key for the GRPC metadata where the IP addresses are stored
+const ipAddressesKey = "ipaddresseskey"
 
 // extractHost returns the Host IP address without any port information
 func extractHost(address string) string {
@@ -53,7 +53,7 @@ func GetSourceFromOutgoingCtx(ctx context.Context) string {
 	if !ok {
 		return ""
 	}
-	ipAddresses, ok := md[IPAddressesKey]
+	ipAddresses, ok := md[ipAddressesKey]
 	if !ok {
 		return ""
 	}
@@ -66,9 +66,26 @@ func GetSourceFromIncomingCtx(ctx context.Context) string {
 	if !ok {
 		return ""
 	}
-	ipAddresses, ok := md[IPAddressesKey]
+	ipAddresses, ok := md[ipAddressesKey]
 	if !ok {
 		return ""
 	}
 	return ipAddresses[0]
 }
+
+// AddSourceToOutgoingContext adds the given source to the GRPC context
+func AddSourceToOutgoingContext(ctx context.Context, source string) context.Context {
+	if source != "" {
+		ctx = metadata.AppendToOutgoingContext(ctx, ipAddressesKey, source)
+	}
+	return ctx
+}
+
+// AddSourceToIncomingContext adds the given source to the GRPC context
+func AddSourceToIncomingContext(ctx context.Context, source string) context.Context {
+	if source != "" {
+		md := metadata.Pairs(ipAddressesKey, source)
+		ctx = metadata.NewIncomingContext(ctx, md)
+	}
+	return ctx
+}

pkg/util/push/push.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"net/http"
 
-	"google.golang.org/grpc/metadata"
-
 	"github.com/go-kit/kit/log/level"
 	"github.com/weaveworks/common/httpgrpc"
 
@@ -20,11 +18,9 @@ func Handler(cfg distributor.Config, push func(context.Context, *client.WriteReq
 		// Extract X-Forwarder-For header
 		ctx := r.Context()
 		source := util.GetSource(r)
-		logger := util.WithContext(ctx, util.Logger)
-		if source != "" {
-			ctx = metadata.AppendToOutgoingContext(ctx, util.IPAddressesKey, source)
-		}
+		ctx = util.AddSourceToOutgoingContext(ctx, source)
 
+		logger := util.WithContext(ctx, util.Logger)
 		compressionType := util.CompressionTypeFor(r.Header.Get("X-Prometheus-Remote-Write-Version"))
 		var req client.PreallocWriteRequest
 		_, err := util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), cfg.MaxRecvMsgSize, &req, compressionType)