Merge pull request #635 from uber/set-shardkey-in-tchannel-client

Set shardkey in tchannel client

runtime/context.go

@@ -41,6 +41,7 @@ const (
 	endpointKey           = contextFieldKey("endpoint")
 	requestUUIDKey        = contextFieldKey("requestUUID")
 	routingDelegateKey    = contextFieldKey("rd")
+	shardKey              = contextFieldKey("sk")
 	endpointRequestHeader = contextFieldKey("endpointRequestHeader")
 	requestLogFields      = contextFieldKey("requestLogFields")
 	scopeTags             = contextFieldKey("scopeTags")
@@ -147,6 +148,22 @@ func GetRoutingDelegateFromCtx(ctx context.Context) string {
 	return ""
 }
 
+// WithShardKey adds the tchannel shard key information in the
+// request context.
+func WithShardKey(ctx context.Context, sk string) context.Context {
+	return context.WithValue(ctx, shardKey, sk)
+}
+
+// GetShardKeyFromCtx returns the tchannel shardkey info
+// extracted from context.
+func GetShardKeyFromCtx(ctx context.Context) string {
+	if val := ctx.Value(shardKey); val != nil {
+		sk, _ := val.(string)
+		return sk
+	}
+	return ""
+}
+
 // WithLogFields returns a new context with the given log fields attached to context.Context
 func WithLogFields(ctx context.Context, newFields ...zap.Field) context.Context {
 	return context.WithValue(ctx, requestLogFields, accumulateLogFields(ctx, newFields))

runtime/context_test.go

@@ -141,6 +141,24 @@ func TestGetRoutingDelegateFromCtx(t *testing.T) {
 	assert.Equal(t, expected, rd)
 }
 
+func TestWithShardKey(t *testing.T) {
+	expected := "myshardkey"
+	ctx := WithShardKey(context.TODO(), expected)
+	sk := ctx.Value(shardKey)
+	shardKey, ok := sk.(string)
+
+	assert.True(t, ok)
+	assert.Equal(t, shardKey, expected)
+}
+
+func TestGetShardKeyFromCtx(t *testing.T) {
+	expected := "myshardkey"
+	ctx := WithShardKey(context.TODO(), expected)
+	sk := GetShardKeyFromCtx(ctx)
+
+	assert.Equal(t, expected, sk)
+}
+
 func TestContextLogger(t *testing.T) {
 	zapLoggerCore, logs := observer.New(zap.DebugLevel)
 	zapLogger := zap.New(zapLoggerCore)

runtime/tchannel_client.go

@@ -80,6 +80,7 @@ type TChannelClient struct {
 	timeout           time.Duration
 	timeoutPerAttempt time.Duration
 	routingKey        *string
+	shardKey          *string
 	metrics           ContextMetrics
 	contextExtractor  ContextExtractor
 
@@ -231,6 +232,12 @@ func (c *TChannelClient) call(
 	if rd != "" {
 		ctxBuilder.SetRoutingDelegate(rd)
 	}
+
+	sk := GetShardKeyFromCtx(ctx)
+	if sk != "" {
+		ctxBuilder.SetShardKey(sk)
+	}
+
 	ctx, cancel := ctxBuilder.Build()
 	defer cancel()