Merge cockroachdb#137344

137344: ctxutil: introduce fast values r=RaduBerinde a=RaduBerinde

#### go.mod: update logtag dependency

See cockroachdb/logtags#4

Epic: none
Release note: None

#### serverident: unexport ServerIdentificationContextKey

Epic: none
Release note: None

#### ctxutil: introduce fast values

This commit introduces a custom context type that can be used to more
efficiently associate values to contexts.

Instead of using arbitrary objects as keys, the keys are a small set
of integers. Any key must be globally registered first, and there is a
limited number that can be registered.

The `fastValuesCtx` stores *all* values for these keys, so traversing
the context hierarchy is not necessary. We also provide a way to add
multiple values at the same time, which saves on allocations.

I looked at before and after profiles in sysbench. Before:
 - CPU usage:
  context.value 0.6%
  context.WithValue 0.3%
  total 0.9%
 - Allocs:
  context.WithValue: 3.5%

After:
 - CPU usage:
  context.value + ctxutil.FastValue 0.5%
  context.WithValue 0.1%
  ctxutil.WithFastValue(s) 0.1%
  total 0.7%
 - Allocs:
  context.WithValue: 0.2%
  context.WithFastValue: 0.6%
  ctxutil.FastValuesBuilder.Finish: 1.4%
  total 2.2%

I will investigate improving on this, perhaps with providing a
pre-allocated context in codepaths where this is possible.

Informs: cockroachdb#136581
Informs: cockroachdb#135904
Release note: None

Co-authored-by: Radu Berinde <[email protected]>

DEPS.bzl

@@ -1821,10 +1821,10 @@ def go_deps():
         name = "com_github_cockroachdb_logtags",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/cockroachdb/logtags",
-        sha256 = "ca7776f47e5fecb4c495490a679036bfc29d95bd7625290cfdb9abb0baf97476",
-        strip_prefix = "github.com/cockroachdb/[email protected]20230118201751-21c54148d20b",
+        sha256 = "ff8b3d36873cfff78669000d194cbaa96520549a62408c52b538380d953b99d7",
+        strip_prefix = "github.com/cockroachdb/[email protected]20241205023844-89a8856d99be",
         urls = [
-            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/logtags/com_github_cockroachdb_logtags-v0.0.0-20230118201751-21c54148d20b.zip",
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/logtags/com_github_cockroachdb_logtags-v0.0.0-20241205023844-89a8856d99be.zip",
         ],
     )
     go_repository(

build/bazelutil/distdir_files.bzl

@@ -352,7 +352,7 @@ DISTDIR_FILES = {
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/go-test-teamcity/com_github_cockroachdb_go_test_teamcity-v0.0.0-20191211140407-cff980ad0a55.zip": "bac30148e525b79d004da84d16453ddd2d5cd20528e9187f1d7dac708335674b",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/gogoproto/com_github_cockroachdb_gogoproto-v1.3.3-0.20241118145159-46874edb1b83.zip": "1382d1ef3a015855f8268c595e2afeef7ae7374b7e27c2181b27a42c2fb639a7",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/gostdlib/com_github_cockroachdb_gostdlib-v1.19.0.zip": "c4d516bcfe8c07b6fc09b8a9a07a95065b36c2855627cb3514e40c98f872b69e",
-    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/logtags/com_github_cockroachdb_logtags-v0.0.0-20230118201751-21c54148d20b.zip": "ca7776f47e5fecb4c495490a679036bfc29d95bd7625290cfdb9abb0baf97476",
+    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/logtags/com_github_cockroachdb_logtags-v0.0.0-20241205023844-89a8856d99be.zip": "ff8b3d36873cfff78669000d194cbaa96520549a62408c52b538380d953b99d7",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/metamorphic/com_github_cockroachdb_metamorphic-v0.0.0-20231108215700-4ba948b56895.zip": "28c8cf42192951b69378cf537be5a9a43f2aeb35542908cc4fe5f689505853ea",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/pebble/com_github_cockroachdb_pebble-v0.0.0-20241206160845-91e091aa3637.zip": "dfd82d8e65c0b4cf44629de5ad07104a327c80d2cd515a387fc008f6051df040",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/cockroachdb/redact/com_github_cockroachdb_redact-v1.1.5.zip": "11b30528eb0dafc8bc1a5ba39d81277c257cbe6946a7564402f588357c164560",

go.mod

@@ -133,7 +133,7 @@ require (
 	github.com/cockroachdb/errors v1.11.3
 	github.com/cockroachdb/go-test-teamcity v0.0.0-20191211140407-cff980ad0a55
 	github.com/cockroachdb/gostdlib v1.19.0
-	github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b
+	github.com/cockroachdb/logtags v0.0.0-20241205023844-89a8856d99be
 	github.com/cockroachdb/pebble v0.0.0-20241206160845-91e091aa3637
 	github.com/cockroachdb/redact v1.1.5
 	github.com/cockroachdb/returncheck v0.0.0-20200612231554-92cdbca611dd

go.sum

@@ -541,8 +541,8 @@ github.com/cockroachdb/gogoproto v1.3.3-0.20241118145159-46874edb1b83/go.mod h1:
 github.com/cockroachdb/gostdlib v1.19.0 h1:cSISxkVnTlWhTkyple/T6NXzOi5659FkhxvUgZv+Eb0=
 github.com/cockroachdb/gostdlib v1.19.0/go.mod h1:+dqqpARXbE/gRDEhCak6dm0l14AaTymPZUKMfURjBtY=
 github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f/go.mod h1:Vz9DsVWQQhf3vs21MhPMZpMGSht7O/2vFW2xusFUVOs=
-github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b h1:r6VH0faHjZeQy818SGhaone5OnYfxFR/+AzdY3sf5aE=
-github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b/go.mod h1:Vz9DsVWQQhf3vs21MhPMZpMGSht7O/2vFW2xusFUVOs=
+github.com/cockroachdb/logtags v0.0.0-20241205023844-89a8856d99be h1:74/l4OkGUedTm/IzF3IUzecZ4y4ClCoQZkLN0X8sMQg=
+github.com/cockroachdb/logtags v0.0.0-20241205023844-89a8856d99be/go.mod h1:Vz9DsVWQQhf3vs21MhPMZpMGSht7O/2vFW2xusFUVOs=
 github.com/cockroachdb/metamorphic v0.0.0-20231108215700-4ba948b56895 h1:XANOgPYtvELQ/h4IrmPAohXqe2pWA8Bwhejr3VQoZsA=
 github.com/cockroachdb/metamorphic v0.0.0-20231108215700-4ba948b56895/go.mod h1:aPd7gM9ov9M8v32Yy5NJrDyOcD8z642dqs+F0CeNXfA=
 github.com/cockroachdb/pebble v0.0.0-20241206160845-91e091aa3637 h1:DHpvNh7WubKNL0yv0nIK318JbMxQo5/CZA1C26KBGJc=

pkg/base/serverident/BUILD.bazel

@@ -5,4 +5,5 @@ go_library(
     srcs = ["server_ident.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/base/serverident",
     visibility = ["//visibility:public"],
+    deps = ["//pkg/util/ctxutil"],
 )

pkg/base/serverident/server_ident.go

@@ -5,38 +5,44 @@
 
 package serverident
 
-import "context"
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/util/ctxutil"
+)
 
 // SystemTenantID is the string representation of
 // roachpb.SystemTenantID. Injected at initialization to avoid
 // an import dependency cycle. See SetSystemTenantID.
 var SystemTenantID string
 
-// ServerIdentificationContextKey is the type of a context.Value key
-// used to carry ServerIdentificationPayload values.
-type ServerIdentificationContextKey struct{}
+// ServerIdentificationContextKey is the fast value key used to annotate a
+// context with a ServerIdentificationPayload.
+var ServerIdentificationContextKey = ctxutil.RegisterFastValueKey()
 
 // ContextWithServerIdentification returns a context annotated with the provided
 // server identity. Use ServerIdentificationFromContext(ctx) to retrieve it from
 // the ctx later.
 func ContextWithServerIdentification(
 	ctx context.Context, serverID ServerIdentificationPayload,
 ) context.Context {
-	return context.WithValue(ctx, ServerIdentificationContextKey{}, serverID)
+	return ctxutil.WithFastValue(ctx, ServerIdentificationContextKey, serverID)
 }
 
 // ServerIdentificationFromContext retrieves the server identity put in the
 // context by ContextWithServerIdentification.
 func ServerIdentificationFromContext(ctx context.Context) ServerIdentificationPayload {
-	r := ctx.Value(ServerIdentificationContextKey{})
+	r := ctxutil.FastValue(ctx, ServerIdentificationContextKey)
 	if r == nil {
 		return nil
 	}
+	// TODO(radu): an interface-to-interface conversion is not great in a hot
+	// path. Maybe the type should be just a func instead of an interface.
 	return r.(ServerIdentificationPayload)
 }
 
 // ServerIdentificationPayload is the type of a context.Value payload
-// associated with a ServerIdentificationContextKey.
+// associated with a serverIdentificationContextKey.
 type ServerIdentificationPayload interface {
 	// ServerIdentityString retrieves an identifier corresponding to the
 	// given retrieval key. If there is no value known for a given key,
@@ -57,7 +63,7 @@ const (
 	IdentifyInstanceID
 	// IdentifyTenantID retrieves the tenant ID of the server.
 	IdentifyTenantID
-	// IdentifyTenantLabel retrieves the tenant name of the server.
+	// IdentifyTenantName retrieves the tenant name of the server.
 	IdentifyTenantName
 )
 

pkg/ccl/sqlproxyccl/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgwirebase",
         "//pkg/util/cache",
+        "//pkg/util/ctxutil",
         "//pkg/util/grpcutil",
         "//pkg/util/httputil",
         "//pkg/util/log",

pkg/ccl/sqlproxyccl/server.go

@@ -14,6 +14,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/util/cache"
+	"github.com/cockroachdb/cockroach/pkg/util/ctxutil"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/httputil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -453,20 +454,20 @@ func (s *Server) shouldLogError(
 	return limiter.ShouldLog()
 }
 
-// requestTagsContextKey is the type of a context.Value key used to carry the
-// request tags map in a context.Context object.
-type requestTagsContextKey struct{}
+// requestTagsContextKey is the fast value key used to carry the request tags
+// map in a context.Context object.
+var requestTagsContextKey = ctxutil.RegisterFastValueKey()
 
 // contextWithRequestTags returns a context annotated with the provided request
 // tags map. Use requestTagsFromContext(ctx) to retrieve it back.
 func contextWithRequestTags(ctx context.Context, reqTags map[string]interface{}) context.Context {
-	return context.WithValue(ctx, requestTagsContextKey{}, reqTags)
+	return ctxutil.WithFastValue(ctx, requestTagsContextKey, reqTags)
 }
 
 // requestTagsFromContext retrieves the request tags map stored in the context
 // via contextWithRequestTags.
 func requestTagsFromContext(ctx context.Context) map[string]interface{} {
-	r := ctx.Value(requestTagsContextKey{})
+	r := ctxutil.FastValue(ctx, requestTagsContextKey)
 	if r == nil {
 		return nil
 	}

pkg/kv/kvserver/kvflowcontrol/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
         "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/util/admission/admissionpb",
+        "//pkg/util/ctxutil",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_dustin_go_humanize//:go-humanize",
     ],

pkg/kv/kvserver/kvflowcontrol/kvflowcontrol.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+	"github.com/cockroachdb/cockroach/pkg/util/ctxutil"
 	"github.com/cockroachdb/redact"
 	"github.com/dustin/go-humanize"
 )
@@ -436,7 +437,7 @@ func (s Stream) SafeFormat(p redact.SafePrinter, verb rune) {
 	p.Printf("t%s/s%s", tenantSt, s.StoreID.String())
 }
 
-type raftAdmissionMetaKey struct{}
+var raftAdmissionMetaKey = ctxutil.RegisterFastValueKey()
 
 // ContextWithMeta returns a Context wrapping the supplied raft admission meta,
 // if any.
@@ -445,15 +446,15 @@ type raftAdmissionMetaKey struct{}
 // #104154.
 func ContextWithMeta(ctx context.Context, meta *kvflowcontrolpb.RaftAdmissionMeta) context.Context {
 	if meta != nil {
-		ctx = context.WithValue(ctx, raftAdmissionMetaKey{}, meta)
+		ctx = ctxutil.WithFastValue(ctx, raftAdmissionMetaKey, meta)
 	}
 	return ctx
 }
 
 // MetaFromContext returns the raft admission meta embedded in the Context, if
 // any.
 func MetaFromContext(ctx context.Context) *kvflowcontrolpb.RaftAdmissionMeta {
-	val := ctx.Value(raftAdmissionMetaKey{})
+	val := ctxutil.FastValue(ctx, raftAdmissionMetaKey)
 	h, ok := val.(*kvflowcontrolpb.RaftAdmissionMeta)
 	if !ok {
 		return nil

pkg/server/testserver.go

@@ -2509,14 +2509,14 @@ func TestingMakeLoggingContexts(
 	appTenantID roachpb.TenantID,
 ) (sysContext, appContext context.Context) {
 	ctxSysTenant := context.Background()
-	ctxSysTenant = context.WithValue(ctxSysTenant, serverident.ServerIdentificationContextKey{}, &idProvider{
+	ctxSysTenant = serverident.ContextWithServerIdentification(ctxSysTenant, &idProvider{
 		tenantID:   roachpb.SystemTenantID,
 		clusterID:  &base.ClusterIDContainer{},
 		serverID:   &base.NodeIDContainer{},
 		tenantName: roachpb.NewTenantNameContainer("system"),
 	})
 	ctxAppTenant := context.Background()
-	ctxAppTenant = context.WithValue(ctxAppTenant, serverident.ServerIdentificationContextKey{}, &idProvider{
+	ctxAppTenant = serverident.ContextWithServerIdentification(ctxAppTenant, &idProvider{
 		tenantID:   appTenantID,
 		clusterID:  &base.ClusterIDContainer{},
 		serverID:   &base.NodeIDContainer{},

pkg/sql/BUILD.bazel

@@ -547,6 +547,7 @@ go_library(
         "//pkg/util/collatedstring",
         "//pkg/util/ctxgroup",
         "//pkg/util/ctxlog",
+        "//pkg/util/ctxutil",
         "//pkg/util/duration",
         "//pkg/util/encoding",
         "//pkg/util/encoding/csv",

pkg/sql/conn_executor.go

@@ -70,6 +70,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/cancelchecker"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxlog"
+	"github.com/cockroachdb/cockroach/pkg/util/ctxutil"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/fsm"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -4714,38 +4715,34 @@ func (ps connExPrepStmtsAccessor) DeleteAll(ctx context.Context) {
 	)
 }
 
-// contextStatementKey is an empty type for the handle associated with the
-// statement value (see context.Value).
-type contextStatementKey struct{}
+var contextStatementKey = ctxutil.RegisterFastValueKey()
 
 // withStatement adds a SQL statement to the provided context. The statement
 // will then be included in crash reports which use that context.
 func withStatement(ctx context.Context, stmt tree.Statement) context.Context {
-	return context.WithValue(ctx, contextStatementKey{}, stmt)
+	return ctxutil.WithFastValue(ctx, contextStatementKey, stmt)
 }
 
 // statementFromCtx returns the statement value from a context, or nil if unset.
 func statementFromCtx(ctx context.Context) tree.Statement {
-	stmt := ctx.Value(contextStatementKey{})
+	stmt := ctxutil.FastValue(ctx, contextStatementKey)
 	if stmt == nil {
 		return nil
 	}
 	return stmt.(tree.Statement)
 }
 
-// contextGistKey is an empty type for the handle associated with the
-// gist value (see context.Value).
-type contextPlanGistKey struct{}
+var contextPlanGistKey = ctxutil.RegisterFastValueKey()
 
 func withPlanGist(ctx context.Context, gist string) context.Context {
 	if gist == "" {
 		return ctx
 	}
-	return context.WithValue(ctx, contextPlanGistKey{}, gist)
+	return ctxutil.WithFastValue(ctx, contextPlanGistKey, gist)
 }
 
 func planGistFromCtx(ctx context.Context) string {
-	val := ctx.Value(contextPlanGistKey{})
+	val := ctxutil.FastValue(ctx, contextPlanGistKey)
 	if val != nil {
 		return val.(string)
 	}

pkg/util/ctxutil/BUILD.bazel

@@ -7,19 +7,23 @@ go_library(
         "canceler_1_21_bazel.go",
         "context.go",
         "doc.go",
+        "fast_value.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/util/ctxutil",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/util/buildutil",
-        "//pkg/util/log",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_logtags//:logtags",
     ],
 )
 
 go_test(
     name = "ctxutil_test",
-    srcs = ["context_test.go"],
+    srcs = [
+        "context_test.go",
+        "fast_value_test.go",
+    ],
     embed = [":ctxutil"],
     deps = [
         "//pkg/util/leaktest",

pkg/util/ctxutil/context.go

@@ -7,10 +7,10 @@ package ctxutil
 
 import (
 	"context"
+	"fmt"
 	_ "unsafe" // Must import unsafe to enable linkname.
 
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
-	"github.com/cockroachdb/cockroach/pkg/util/log"
 )
 
 // WhenDoneFunc is the callback invoked by context when it becomes done.
@@ -49,7 +49,7 @@ func WhenDone(parent context.Context, done WhenDoneFunc) bool {
 	// But, be safe and loudly fail tests in case somebody introduces strange
 	// context implementation.
 	if buildutil.CrdbTestBuild && !CanDirectlyDetectCancellation(parent) {
-		log.Fatalf(parent, "expected context that supports direct cancellation detection, found %T", parent)
+		panic(fmt.Sprintf("expected context that supports direct cancellation detection, found %T", parent))
 	}
 
 	propagateCancel(parent, done)