log/slog: optimize appendKey to reduce allocations

This change introduces a new method, `appendTwoStrings`, which
optimizes the `appendKey` function by avoiding the allocation of a
temporary string (string concatenation of prefix and key). Instead, it
directly appends the prefix and key to the buffer.

Additionally, added `BenchmarkAppendKey` benchmark tests to validate performance improvements.

This change improves performance in cases where large prefixes are used,
as verified by the following benchmarks:

goos: darwin
goarch: arm64
pkg: log/slog
cpu: Apple M1 Max
                             │   old.out   │               new.out               │
                             │   sec/op    │   sec/op     vs base                │
AppendKey/prefix_size_5-10     44.41n ± 0%   35.62n ± 0%  -19.80% (p=0.000 n=10)
AppendKey/prefix_size_10-10    48.17n ± 0%   39.12n ± 0%  -18.80% (p=0.000 n=10)
AppendKey/prefix_size_30-10    84.50n ± 0%   62.30n ± 0%  -26.28% (p=0.000 n=10)
AppendKey/prefix_size_50-10    124.9n ± 0%   102.3n ± 0%  -18.09% (p=0.000 n=10)
AppendKey/prefix_size_100-10   203.6n ± 1%   168.7n ± 0%  -17.14% (p=0.000 n=10)
geomean                        85.61n        68.41n       -20.09%

                             │    old.out    │                 new.out                 │
                             │     B/op      │    B/op     vs base                     │
AppendKey/prefix_size_5-10      0.000 ± 0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
AppendKey/prefix_size_10-10     0.000 ± 0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
AppendKey/prefix_size_30-10     48.00 ± 0%      0.00 ± 0%  -100.00% (p=0.000 n=10)
AppendKey/prefix_size_50-10    128.00 ± 0%     64.00 ± 0%   -50.00% (p=0.000 n=10)
AppendKey/prefix_size_100-10    224.0 ± 0%     112.0 ± 0%   -50.00% (p=0.000 n=10)
geomean                                    ²               ?                       ² ³
¹ all samples are equal
² summaries must be >0 to compute geomean
³ ratios must be >0 to compute geomean

                             │   old.out    │                 new.out                 │
                             │  allocs/op   │ allocs/op   vs base                     │
AppendKey/prefix_size_5-10     0.000 ± 0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
AppendKey/prefix_size_10-10    0.000 ± 0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
AppendKey/prefix_size_30-10    1.000 ± 0%     0.000 ± 0%  -100.00% (p=0.000 n=10)
AppendKey/prefix_size_50-10    2.000 ± 0%     1.000 ± 0%   -50.00% (p=0.000 n=10)
AppendKey/prefix_size_100-10   2.000 ± 0%     1.000 ± 0%   -50.00% (p=0.000 n=10)
geomean                                   ²               ?                       ² ³
¹ all samples are equal
² summaries must be >0 to compute geomean
³ ratios must be >0 to compute geomean

This patch improves performance without altering the external behavior of the `slog` package.

Change-Id: I8b47718de522196f06e0ddac48af73e352d2e5cd
Reviewed-on: https://go-review.googlesource.com/c/go/+/631415
Reviewed-by: Alan Donovan <[email protected]>
Reviewed-by: Jonathan Amsterdam <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: Aoang

src/log/slog/handler.go

@@ -525,8 +525,7 @@ func (s *handleState) appendError(err error) {
 func (s *handleState) appendKey(key string) {
 	s.buf.WriteString(s.sep)
 	if s.prefix != nil && len(*s.prefix) > 0 {
-		// TODO: optimize by avoiding allocation.
-		s.appendString(string(*s.prefix) + key)
+		s.appendTwoStrings(string(*s.prefix), key)
 	} else {
 		s.appendString(key)
 	}
@@ -538,6 +537,24 @@ func (s *handleState) appendKey(key string) {
 	s.sep = s.h.attrSep()
 }
 
+// appendTwoStrings implements appendString(prefix + key), but faster.
+func (s *handleState) appendTwoStrings(x, y string) {
+	buf := *s.buf
+	switch {
+	case s.h.json:
+		buf.WriteByte('"')
+		buf = appendEscapedJSONString(buf, x)
+		buf = appendEscapedJSONString(buf, y)
+		buf.WriteByte('"')
+	case !needsQuoting(x) && !needsQuoting(y):
+		buf.WriteString(x)
+		buf.WriteString(y)
+	default:
+		buf = strconv.AppendQuote(buf, x+y)
+	}
+	*s.buf = buf
+}
+
 func (s *handleState) appendString(str string) {
 	if s.h.json {
 		s.buf.WriteByte('"')

src/log/slog/handler_test.go

@@ -10,7 +10,9 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
+	"log/slog/internal/buffer"
 	"os"
 	"path/filepath"
 	"slices"
@@ -529,6 +531,20 @@ func TestJSONAndTextHandlers(t *testing.T) {
 			wantText: "name.first=Perry name.last=Platypus",
 			wantJSON: `{"name":{"first":"Perry","last":"Platypus"}}`,
 		},
+		{
+			name:    "group and key (or both) needs quoting",
+			replace: removeKeys(TimeKey, LevelKey),
+			attrs: []Attr{
+				Group("prefix",
+					String(" needs quoting ", "v"), String("NotNeedsQuoting", "v"),
+				),
+				Group("prefix needs quoting",
+					String(" needs quoting ", "v"), String("NotNeedsQuoting", "v"),
+				),
+			},
+			wantText: `msg=message "prefix. needs quoting "=v prefix.NotNeedsQuoting=v "prefix needs quoting. needs quoting "=v "prefix needs quoting.NotNeedsQuoting"=v`,
+			wantJSON: `{"msg":"message","prefix":{" needs quoting ":"v","NotNeedsQuoting":"v"},"prefix needs quoting":{" needs quoting ":"v","NotNeedsQuoting":"v"}}`,
+		},
 	} {
 		r := NewRecord(testTime, LevelInfo, "message", callerPC(2))
 		line := strconv.Itoa(r.source().Line)
@@ -732,3 +748,31 @@ func TestDiscardHandler(t *testing.T) {
 	l.Info("info", "a", []Attr{Int("i", 1)})
 	l.Info("info", "a", GroupValue(Int("i", 1)))
 }
+
+func BenchmarkAppendKey(b *testing.B) {
+	for _, size := range []int{5, 10, 30, 50, 100} {
+		for _, quoting := range []string{"no_quoting", "pre_quoting", "key_quoting", "both_quoting"} {
+			b.Run(fmt.Sprintf("%s_prefix_size_%d", quoting, size), func(b *testing.B) {
+				var (
+					hs     = NewJSONHandler(io.Discard, nil).newHandleState(buffer.New(), false, "")
+					prefix = bytes.Repeat([]byte("x"), size)
+					key    = "key"
+				)
+
+				if quoting == "pre_quoting" || quoting == "both_quoting" {
+					prefix[0] = '"'
+				}
+				if quoting == "key_quoting" || quoting == "both_quoting" {
+					key = "ke\""
+				}
+
+				hs.prefix = (*buffer.Buffer)(&prefix)
+
+				for b.Loop() {
+					hs.appendKey(key)
+					hs.buf.Reset()
+				}
+			})
+		}
+	}
+}