go/vt/callerid: speedup callerid context access

go/vt/callerid/callerid.go

@@ -26,14 +26,15 @@ import (
 )
 
 // The datatype for CallerID Context Keys
-type callerIDKey int
-
-var (
-	// internal Context key for immediate CallerID
-	immediateCallerIDKey callerIDKey
-	// internal Context key for effective CallerID
-	effectiveCallerIDKey callerIDKey = 1
-)
+type callerIDContextKey struct{}
+
+// callerIDContext is a tuple containing the entirety
+// of the "callerid" such that both pointers can be
+// stored together.
+type callerIDContext struct {
+	ef *vtrpcpb.CallerID
+	im *querypb.VTGateCallerID
+}
 
 // NewImmediateCallerID creates a querypb.VTGateCallerID initialized with username
 func NewImmediateCallerID(username string) *querypb.VTGateCallerID {
@@ -90,30 +91,23 @@ func GetSubcomponent(ef *vtrpcpb.CallerID) string {
 // NewContext adds the provided EffectiveCallerID(vtrpcpb.CallerID) and ImmediateCallerID(querypb.VTGateCallerID)
 // into the Context
 func NewContext(ctx context.Context, ef *vtrpcpb.CallerID, im *querypb.VTGateCallerID) context.Context {
-	ctx = context.WithValue(
-		context.WithValue(ctx, effectiveCallerIDKey, ef),
-		immediateCallerIDKey,
-		im,
-	)
-	return ctx
+	return context.WithValue(ctx, callerIDContextKey{}, callerIDContext{ef, im})
 }
 
 // EffectiveCallerIDFromContext returns the EffectiveCallerID(vtrpcpb.CallerID)
 // stored in the Context, if any
 func EffectiveCallerIDFromContext(ctx context.Context) *vtrpcpb.CallerID {
-	ef, ok := ctx.Value(effectiveCallerIDKey).(*vtrpcpb.CallerID)
-	if ok && ef != nil {
-		return ef
+	if cid, ok := ctx.Value(callerIDContextKey{}).(callerIDContext); ok {
+		return cid.ef
 	}
 	return nil
 }
 
 // ImmediateCallerIDFromContext returns the ImmediateCallerID(querypb.VTGateCallerID)
 // stored in the Context, if any
 func ImmediateCallerIDFromContext(ctx context.Context) *querypb.VTGateCallerID {
-	im, ok := ctx.Value(immediateCallerIDKey).(*querypb.VTGateCallerID)
-	if ok && im != nil {
-		return im
+	if cid, ok := ctx.Value(callerIDContextKey{}).(callerIDContext); ok {
+		return cid.im
 	}
 	return nil
 }

go/vt/callerid/callerid_test.go

@@ -0,0 +1,67 @@
+/*
+Copyright 2019 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package callerid_test
+
+import (
+	"context"
+	"testing"
+
+	"vitess.io/vitess/go/vt/callerid"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCallerIDContext(t *testing.T) {
+	ef := callerid.NewEffectiveCallerID("principal", "component", "subComponent")
+	im := callerid.NewImmediateCallerID("username")
+	ctx := callerid.NewContext(context.Background(), ef, im)
+
+	assert.Equal(t, ef, callerid.EffectiveCallerIDFromContext(ctx))
+	assert.Equal(t, im, callerid.ImmediateCallerIDFromContext(ctx))
+
+	assert.Nil(t, callerid.EffectiveCallerIDFromContext(context.Background()))
+	assert.Nil(t, callerid.ImmediateCallerIDFromContext(context.Background()))
+}
+
+func BenchmarkCallerIDContext(b *testing.B) {
+	ef := callerid.NewEffectiveCallerID("principal", "component", "subComponent")
+	im := callerid.NewImmediateCallerID("username")
+
+	b.Run("New", func(b *testing.B) {
+		b.ReportAllocs()
+		baseCtx := context.Background()
+		for range b.N {
+			_ = callerid.NewContext(baseCtx, ef, im)
+		}
+	})
+
+	b.Run("EffectiveCallerIDFromContext", func(b *testing.B) {
+		b.ReportAllocs()
+		ctx := callerid.NewContext(context.Background(), ef, im)
+		for range b.N {
+			_ = callerid.EffectiveCallerIDFromContext(ctx)
+		}
+	})
+
+	b.Run("ImmediateCallerIDFromContext", func(b *testing.B) {
+		b.ReportAllocs()
+		ctx := callerid.NewContext(context.Background(), ef, im)
+		for range b.N {
+			_ = callerid.EffectiveCallerIDFromContext(ctx)
+		}
+	})
+}