Merge pull request #9 from krithika369/add_context_return_val

Add Response Labels

.gitignore

@@ -11,5 +11,5 @@
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 
-# Dependency directories (remove the comment below to include it)
-# vendor/
+# Dependency directories
+vendor/

README.md

@@ -271,19 +271,21 @@ func (s *OddEvenRoutingStrategy) SelectRoute(
 	ctx context.Context,
 	req fiber.Request,
 	routes map[string]fiber.Component,
-) (fiber.Component, []fiber.Component, error) {
+) (fiber.Component, []fiber.Component, fiber.Labels, error) {
 	sessionIdStr := ""
 	if sessionHeader, ok := req.Header()["X-Session-ID"]; ok {
 		sessionIdStr = sessionHeader[0]
 	}
+    // Metadata that can be propagated upstream for logging / debugging
+    labels := fiber.NewLabelsMap()
 
 	if sessionID, err := strconv.Atoi(sessionIdStr); err != nil {
 		return nil, nil, err
 	} else {
 		if sessionID % 2 != 0 {
-			return routes["route-a"], []fiber.Component{}, nil
+			return routes["route-a"], []fiber.Component{}, labels.WithLabel("Match-Type", "even"), nil
 		} else {
-			return routes["route-b"], []fiber.Component{}, nil
+			return routes["route-b"], []fiber.Component{}, labels.WithLabel("Match-Type", "odd"), nil
 		}
 	}
 }

eager_router.go

@@ -63,7 +63,7 @@ func (fanIn *eagerRouterFanIn) Aggregate(
 ) Response {
 	// use routing strategy to fetch primary route and fallbacks
 	// publish the ordered routes into a channel
-	routesOrderCh, errCh := fanIn.strategy.getRoutesOrder(ctx, req, fanIn.router.GetRoutes())
+	routesOrderCh := fanIn.strategy.getRoutesOrder(ctx, req, fanIn.router.GetRoutes())
 
 	out := make(chan Response, 1)
 	go func() {
@@ -77,6 +77,9 @@ func (fanIn *eagerRouterFanIn) Aggregate(
 			// would be initialized from a routesOrderCh channel
 			routes []Component
 
+			// response labels
+			labels Labels = NewLabelsMap()
+
 			// index of current primary route
 			currentRouteIdx int
 
@@ -93,18 +96,17 @@ func (fanIn *eagerRouterFanIn) Aggregate(
 				} else {
 					responseCh = nil
 				}
-			case orderedRoutes, ok := <-routesOrderCh:
+			case routesOrderResponse, ok := <-routesOrderCh:
 				if ok {
-					routes = orderedRoutes
+					labels = routesOrderResponse.Labels
+					if routesOrderResponse.Err != nil {
+						masterResponse = NewErrorResponse(errors.NewFiberError(req.Protocol(), routesOrderResponse.Err))
+					} else {
+						routes = routesOrderResponse.Components
+					}
 				} else {
 					routesOrderCh = nil
 				}
-			case err, ok := <-errCh:
-				if ok {
-					masterResponse = NewErrorResponse(errors.NewFiberError(req.Protocol(), err))
-				} else {
-					errCh = nil
-				}
 			case <-ctx.Done():
 				if routes == nil {
 					// timeout exceeded, but no routes received. Sending error response
@@ -139,7 +141,7 @@ func (fanIn *eagerRouterFanIn) Aggregate(
 				}
 			}
 		}
-		out <- masterResponse
+		out <- masterResponse.WithLabels(labels)
 	}()
 
 	return <-out

extras/random_routing_strategy.go

@@ -3,6 +3,7 @@ package extras
 import (
 	"context"
 	"math/rand"
+	"strconv"
 
 	"github.com/gojek/fiber"
 )
@@ -19,8 +20,11 @@ func (s *RandomRoutingStrategy) SelectRoute(
 	_ context.Context,
 	_ fiber.Request,
 	routes map[string]fiber.Component,
-) (route fiber.Component, fallbacks []fiber.Component, err error) {
+) (route fiber.Component, fallbacks []fiber.Component, labels fiber.Labels, err error) {
 	idx := rand.Intn(len(routes))
+	// Add idx to attribute map for logging / debugging upstream
+	labels = fiber.NewLabelsMap().WithLabel("idx", strconv.Itoa(idx))
+
 	for _, child := range routes {
 		if idx == 0 {
 			route = child
@@ -29,5 +33,5 @@ func (s *RandomRoutingStrategy) SelectRoute(
 		}
 		idx--
 	}
-	return route, fallbacks, nil
+	return route, fallbacks, labels, nil
 }

go.mod

@@ -1,21 +1,32 @@
 module github.com/gojek/fiber
 
-go 1.14
+go 1.18
 
 require (
 	github.com/ghodss/yaml v1.0.0
 	github.com/google/go-cmp v0.5.8
-	github.com/kr/text v0.2.0 // indirect
-	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	github.com/opentracing/opentracing-go v1.1.0
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/stretchr/testify v1.8.0
 	go.uber.org/zap v1.17.0
+	google.golang.org/grpc v1.48.0
+	google.golang.org/protobuf v1.28.0
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/kr/text v0.2.0 // indirect
+	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/objx v0.4.0 // indirect
+	go.uber.org/atomic v1.7.0 // indirect
+	go.uber.org/multierr v1.6.0 // indirect
 	golang.org/x/net v0.0.0-20220722155237-a158d28d115b // indirect
 	golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f // indirect
+	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/genproto v0.0.0-20220519153652-3a47de7e79bd // indirect
-	google.golang.org/grpc v1.48.0
-	google.golang.org/protobuf v1.28.0
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -113,12 +113,9 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f h1:v4INt8xihDGvnrfjMDVXGxw9wrfxYyCjk0KbXjhR55s=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
-golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

grpc/response.go

@@ -27,10 +27,35 @@ func (r *Response) StatusCode() int {
 	return int(r.Status.Code())
 }
 
+// Label returns all the values associated with the given key, in the response metadata.
+// If the key does not exist, an empty slice will be returned.
+func (r *Response) Label(key string) []string {
+	return r.Metadata.Get(key)
+}
+
+// WithLabel appends the given value(s) to the key, in the response metadata.
+// If the key does not already exist, a new key will be created.
+// The modified response is returned.
+func (r *Response) WithLabel(key string, values ...string) fiber.Response {
+	r.Metadata.Append(key, values...)
+	return r
+}
+
+// WithLabels does the same thing as WithLabel but over a collection of key-values.
+func (r *Response) WithLabels(labels fiber.Labels) fiber.Response {
+	for _, key := range labels.Keys() {
+		values := labels.Label(key)
+		r.Metadata.Append(key, values...)
+	}
+	return r
+}
+
 func (r *Response) BackendName() string {
-	return strings.Join(r.Metadata.Get("backend"), ",")
+	return strings.Join(r.Label("backend"), ",")
 }
 
+// WithBackendName sets the given backend name in the response metadata.
+// The modified response is returned.
 func (r *Response) WithBackendName(backendName string) fiber.Response {
 	r.Metadata.Set("backend", backendName)
 	return r

grpc/response_test.go

@@ -1,9 +1,11 @@
-package grpc
+package grpc_test
 
 import (
 	"log"
 	"testing"
 
+	"github.com/gojek/fiber"
+	"github.com/gojek/fiber/grpc"
 	testproto "github.com/gojek/fiber/internal/testdata/gen/testdata/proto"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc/codes"
@@ -15,16 +17,16 @@ import (
 func TestResponse_Backend(t *testing.T) {
 	tests := []struct {
 		name        string
-		res         Response
-		want        Response
+		res         grpc.Response
+		want        grpc.Response
 		backendName string
 	}{
 		{
 			name: "ok",
-			res: Response{
+			res: grpc.Response{
 				Metadata: map[string][]string{},
 			},
-			want: Response{
+			want: grpc.Response{
 				Metadata: metadata.New(map[string]string{"backend": "testing"}),
 			},
 			backendName: "testing",
@@ -42,21 +44,21 @@ func TestResponse_Backend(t *testing.T) {
 func TestResponse_Status(t *testing.T) {
 	tests := []struct {
 		name            string
-		res             Response
+		res             grpc.Response
 		expectedCode    int
 		expectedSuccess bool
 	}{
 		{
 			name: "ok",
-			res: Response{
+			res: grpc.Response{
 				Status: *status.New(codes.OK, ""),
 			},
 			expectedCode:    0,
 			expectedSuccess: true,
 		},
 		{
 			name: "ok",
-			res: Response{
+			res: grpc.Response{
 				Status: *status.New(codes.InvalidArgument, ""),
 			},
 			expectedCode:    3,
@@ -83,12 +85,12 @@ func TestResponse_Payload(t *testing.T) {
 	responseByte, _ := proto.Marshal(response)
 	tests := []struct {
 		name     string
-		req      Response
+		req      grpc.Response
 		expected []byte
 	}{
 		{
 			name: "",
-			req: Response{
+			req: grpc.Response{
 				Message: responseByte,
 			},
 			expected: responseByte,
@@ -100,3 +102,98 @@ func TestResponse_Payload(t *testing.T) {
 		})
 	}
 }
+
+func TestResponse_Label(t *testing.T) {
+	tests := map[string]struct {
+		response fiber.Response
+		key      string
+		expected []string
+	}{
+		"empty labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{},
+			},
+			key: "dummy-key",
+		},
+		"non-empty labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{"key": []string{"v1", "v2"}},
+			},
+			key:      "key",
+			expected: []string{"v1", "v2"},
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			values := tt.response.Label(tt.key)
+			assert.Equal(t, tt.expected, values)
+		})
+	}
+}
+
+func TestResponse_WithLabel(t *testing.T) {
+	tests := map[string]struct {
+		response fiber.Response
+		key      string
+		values   []string
+		expected []string
+	}{
+		"new labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{"key": []string{"v1", "v2"}},
+			},
+			key:      "k1",
+			values:   []string{"v1", "v2"},
+			expected: []string{"v1", "v2"},
+		},
+		"append labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{"k1": []string{"v1", "v2"}},
+			},
+			key:      "k1",
+			values:   []string{"v3"},
+			expected: []string{"v1", "v2", "v3"},
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			newLabels := tt.response.WithLabel(tt.key, tt.values...)
+			assert.Equal(t, tt.expected, newLabels.Label(tt.key))
+		})
+	}
+}
+
+func TestHTTPResponse_WithLabels(t *testing.T) {
+	tests := map[string]struct {
+		response fiber.Response
+		labels   fiber.Labels
+		key      string
+		expected []string
+	}{
+		"new labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{"key": []string{"v1", "v2"}},
+			},
+			labels:   fiber.LabelsMap{"k1": []string{"v1", "v2"}},
+			key:      "k1",
+			expected: []string{"v1", "v2"},
+		},
+		"append labels": {
+			response: &grpc.Response{
+				Metadata: map[string][]string{"k1": []string{"v1", "v2"}},
+			},
+			labels:   fiber.LabelsMap{"k1": []string{"v3"}},
+			key:      "k1",
+			expected: []string{"v1", "v2", "v3"},
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			newLabels := tt.response.WithLabels(tt.labels)
+			assert.Equal(t, tt.expected, newLabels.Label(tt.key))
+		})
+	}
+}