Merge cockroachdb#137632

137632: ctxutil: bundle up fastValueCtx allocs in Batch r=RaduBerinde a=RaduBerinde

This commit adds `fastValuesContainer` which preallocates multiple
`fastValueCtx` objects. Contexts allocated from a container retain a
reference to the container, allowing it to be used by new descendant
contexts.

We use this explicitly in `Batch()` where we annotate a context three
times (once at the node level, then at the store level, then at the
replica level).

```
name                                   old time/op    new time/op    delta
Sysbench/KV/3node/oltp_read_write-10     2.41ms ± 1%    2.37ms ± 2%  -1.71%  (p=0.001 n=8+9)
Sysbench/SQL/3node/oltp_read_write-10    5.15ms ± 4%    5.19ms ± 4%    ~     (p=0.631 n=10+10)

name                                   old alloc/op   new alloc/op   delta
Sysbench/KV/3node/oltp_read_write-10     1.43MB ± 0%    1.43MB ± 0%  +0.15%  (p=0.001 n=10+10)
Sysbench/SQL/3node/oltp_read_write-10    2.49MB ± 1%    2.50MB ± 0%  +0.37%  (p=0.024 n=9+9)

name                                   old allocs/op  new allocs/op  delta
Sysbench/KV/3node/oltp_read_write-10      6.21k ± 0%     6.14k ± 0%  -1.19%  (p=0.000 n=9+9)
Sysbench/SQL/3node/oltp_read_write-10     10.5k ± 0%     10.5k ± 0%  -0.48%  (p=0.000 n=10+10)
```

Informs: cockroachdb#135904
Release note: None

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

pkg/server/node.go

@@ -1817,7 +1817,7 @@ func (n *Node) Batch(ctx context.Context, args *kvpb.BatchRequest) (*kvpb.BatchR
 	// NB: Node.Batch is called directly for "local" calls. We don't want to
 	// carry the associated log tags forward as doing so makes adding additional
 	// log tags more expensive and makes local calls differ from remote calls.
-	ctx = n.storeCfg.AmbientCtx.ResetAndAnnotateCtx(ctx)
+	ctx = n.storeCfg.AmbientCtx.ResetAndAnnotateCtxPrealloc(ctx)
 
 	comparisonResult := n.getLocalityComparison(ctx, args.GatewayNodeID)
 	n.metrics.updateCrossLocalityMetricsOnBatchRequest(comparisonResult, int64(args.Size()))

pkg/util/ctxutil/fast_value.go

@@ -44,21 +44,10 @@ func FastValue(ctx context.Context, key FastValueKey) any {
 // This is a more efficient alternative to using context.WithValue() and
 // Context.Value().
 func WithFastValue(parent context.Context, key FastValueKey, val any) context.Context {
-	c := &fastValuesCtx{Context: parent}
-	// Because we want the new node to be able to produce values for all keys, we
-	// need to copy all fast values from the parent context (i.e. the closest
-	// fastValuesCtx ancestor).
-	if p, ok := parent.(*fastValuesCtx); ok {
-		// Connect to the grandparent context, since the parent won't do anything.
-		// Note that (because of exactly this code), the grandparent context won't
-		// be a fastValuesCtx.
-		c.Context = p.Context
-		c.values = p.values
-	} else if v := parent.Value(fastValuesAncestorKey{}); v != nil {
-		c.values = v.(*fastValuesCtx).values
-	}
-	c.values[key] = val
-	return c
+	ancestor, parent := findFastValuesCtxAncestor(parent)
+	ctx := newFastValuesCtx(ancestor, parent)
+	ctx.values[key] = val
+	return ctx
 }
 
 // WithFastValues starts the process of creating a new context that sets a
@@ -73,54 +62,53 @@ func WithFastValue(parent context.Context, key FastValueKey, val any) context.Co
 //	b.Set(key2, val2)
 //	ctx = b.Finish()
 func WithFastValues(parent context.Context) FastValuesBuilder {
-	var bld FastValuesBuilder
-	bld.parent = parent
-	// Because we want the new node to be able to produce values for all keys, we
-	// need to copy all fast values from the parent context (i.e. the closest
-	// fastValuesCtx ancestor).
-	if p, ok := parent.(*fastValuesCtx); ok {
-		// Connect to the grandparent context, since the parent won't do anything.
-		// Note that (because of exactly this code), the grandparent context won't
-		// be a fastValuesCtx.
-		bld.parent = p.Context
-		bld.values = p.values
-	} else if v := parent.Value(fastValuesAncestorKey{}); v != nil {
-		bld.values = v.(*fastValuesCtx).values
+	ancestor, parent := findFastValuesCtxAncestor(parent)
+	return FastValuesBuilder{
+		ctx: newFastValuesCtx(ancestor, parent),
+	}
+}
+
+// WithFastValuesPrealloc is like WithFastValues, but preallocates multiple
+// contexts that can later be reused on subsequent WithFastValue(s) calls to
+// save on allocations.
+func WithFastValuesPrealloc(parent context.Context) FastValuesBuilder {
+	ancestor, parent := findFastValuesCtxAncestor(parent)
+	container := &fastValuesContainer{}
+	container.used.Store(1)
+	ctx := &container.buf[0]
+	ctx.init(parent, container, ancestor)
+	return FastValuesBuilder{
+		ctx: ctx,
 	}
-	return bld
 }
 
 // FastValuesBuilder contains multiple values; used for WithFastValues.
 type FastValuesBuilder struct {
-	parent context.Context
-	values [MaxFastValues]any
+	ctx *fastValuesCtx
 }
 
 // Get gets the value for the key in the context being built. If this key was
 // set before, it returns the last value passed to Set(); otherwise it returns
 // the value in the context that was passed to WithFastValues().
 func (b *FastValuesBuilder) Get(key FastValueKey) any {
-	return b.values[key]
+	return b.ctx.values[key]
 }
 
 // Set sets the value for the key in the context being built.
 func (b *FastValuesBuilder) Set(key FastValueKey, val any) {
-	b.values[key] = val
+	b.ctx.values[key] = val
 }
 
 // Finish constructs the context with the values set by Set().
 //
 // The FastValuesBuilder must not be used again.
 func (b *FastValuesBuilder) Finish() context.Context {
-	if buildutil.CrdbTestBuild && b.parent == nil {
+	if buildutil.CrdbTestBuild && b.ctx == nil {
 		panic("invalid use of FastValuesBuilder")
 	}
-	c := &fastValuesCtx{
-		Context: b.parent,
-		values:  b.values,
-	}
-	b.parent = nil
-	return c
+	ctx := b.ctx
+	b.ctx = nil
+	return ctx
 }
 
 // RegisterFastValueKey creates a key that can be used with WithFastValue(). This
@@ -142,17 +130,101 @@ var numFastValues atomic.Uint32
 type fastValuesCtx struct {
 	context.Context
 	values [MaxFastValues]any
+	// container is set if this node was allocated using a fastValuesContainer
+	// that had more space. The container can be used to allocate new descendant
+	// nodes.
+	container *fastValuesContainer
+}
+
+func (ctx *fastValuesCtx) init(
+	parent context.Context, container *fastValuesContainer, ancestor *fastValuesCtx,
+) {
+	ctx.Context = parent
+	if ancestor != nil {
+		// Because we want the new node to be able to produce values for all fast
+		// keys, we need to copy all fast values from the parent context (i.e. the
+		// closest fastValuesCtx ancestor).
+		ctx.values = ancestor.values
+	}
+	ctx.container = container
 }
 
 // fastValuesAncestorKey is used to retrieve the closest fastValuesCtx ancestor.
 type fastValuesAncestorKey struct{}
 
 // Value is part of the context.Context interface.
-func (c *fastValuesCtx) Value(key any) any {
+func (ctx *fastValuesCtx) Value(key any) any {
 	if _, ok := key.(fastValuesAncestorKey); ok {
-		return c
+		return ctx
+	}
+	return ctx.Context.Value(key)
+}
+
+// newFastValuesCtx returns a new context with the given parent and the values
+// from the ancestor.
+//
+// If the ancestor has an associated container that has free slots, it is used
+// to avoid a new allocation.
+func newFastValuesCtx(ancestor *fastValuesCtx, parent context.Context) *fastValuesCtx {
+	var ctx *fastValuesCtx
+	var container *fastValuesContainer
+	if ancestor != nil && ancestor.container != nil {
+		// Try to use the container.
+		var more bool
+		ctx, more = ancestor.container.get()
+		if more {
+			container = ancestor.container
+		} else if ctx == nil {
+			ctx = &fastValuesCtx{}
+		}
+	} else {
+		ctx = &fastValuesCtx{}
+	}
+	ctx.init(parent, container, ancestor)
+	return ctx
+}
+
+// findFastValuesCtxAncestor returns the closest fastValuesCtx ancestor of the
+// given parent context.
+//
+// The returned newParent is the same with parent except when parent is a
+// fastValueCtx, in which case newParent is that parent's parent. This is used
+// to avoid unnecessarily stacking fastValueCtx nodes.
+func findFastValuesCtxAncestor(
+	parent context.Context,
+) (_ *fastValuesCtx, newParent context.Context) {
+	if p, ok := parent.(*fastValuesCtx); ok {
+		return p, p.Context
+	}
+	if v := parent.Value(fastValuesAncestorKey{}); v != nil {
+		return v.(*fastValuesCtx), parent
+	}
+	return nil, parent
+}
+
+const fastValuesContainerSize = 3
+
+// fastValuesContainer is used when we expect successive WithFastValue(s) calls.
+// We allocate a multi-object container for the first call and then use it as
+// needed.
+type fastValuesContainer struct {
+	used atomic.Int32
+	buf  [fastValuesContainerSize]fastValuesCtx
+}
+
+// get returns a new fastValueCtx from the container, and whether there are
+// still available slots in the container. If the container is used, returns nil
+// and false.
+func (c *fastValuesContainer) get() (_ *fastValuesCtx, more bool) {
+	for {
+		u := c.used.Load()
+		if u >= fastValuesContainerSize {
+			return nil, false
+		}
+		if c.used.CompareAndSwap(u, u+1) {
+			return &c.buf[u], u+1 < fastValuesContainerSize
+		}
 	}
-	return c.Context.Value(key)
 }
 
 func init() {

pkg/util/log/ambient_context.go

@@ -126,6 +126,17 @@ func (ac *AmbientContext) ResetAndAnnotateCtx(ctx context.Context) context.Conte
 	}
 }
 
+// ResetAndAnnotateCtxPrealloc is like ResetAndAnnotateCtx but allocates a
+// container to avoid allocations on future annotations on the context and its
+// descendants.
+func (ac *AmbientContext) ResetAndAnnotateCtxPrealloc(ctx context.Context) context.Context {
+	bld := ctxutil.WithFastValuesPrealloc(ctx)
+	// We set these unconditionally in case they are already set in the context.
+	bld.Set(ctxutil.LogTagsKey, ac.tags)
+	bld.Set(serverident.ServerIdentificationContextKey, ac.ServerIDs)
+	return bld.Finish()
+}
+
 func (ac *AmbientContext) annotateCtxInternal(ctx context.Context) context.Context {
 	bld := ctxutil.WithFastValues(ctx)
 	if ac.tags != nil {