net/textproto: reduce allocations in ReadMIMEHeader

[rs]

What version of Go are you using (go version)?

1.14

What did you do?

For our use-case (a HTTP proxy) textproto.ReadMIMEHeader is in the top 3 in term of allocations. In order to ease the GC, we could pre-allocate 512 bytes to be used for small header values. This can reduce by half the number of allocations for this function.

name                              old time/op    new time/op    delta
ReadMIMEHeader/client_headers-16    2.30µs ± 5%    2.11µs ± 2%   -8.21%  (p=0.000 n=10+10)
ReadMIMEHeader/server_headers-16    1.94µs ± 3%    1.85µs ± 3%   -4.56%  (p=0.000 n=10+10)

name                              old alloc/op   new alloc/op   delta
ReadMIMEHeader/client_headers-16    1.53kB ± 0%    1.69kB ± 0%  +10.48%  (p=0.000 n=10+10)
ReadMIMEHeader/server_headers-16    1.09kB ± 0%    1.44kB ± 0%  +32.32%  (p=0.000 n=10+10)

name                              old allocs/op  new allocs/op  delta
ReadMIMEHeader/client_headers-16      14.0 ± 0%       6.0 ± 0%  -57.14%  (p=0.000 n=10+10)
ReadMIMEHeader/server_headers-16      14.0 ± 0%       6.0 ± 0%  -57.14%  (p=0.000 n=10+10)

The patch is pretty small:

index d26e981ae4..6126a6685c 100644
--- a/src/net/textproto/reader.go
+++ b/src/net/textproto/reader.go
@@ -13,6 +13,7 @@ import (
        "strconv"
        "strings"
        "sync"
+       "unsafe"
 )

 // A Reader implements convenience methods for reading requests
@@ -502,6 +503,12 @@ func (r *Reader) ReadMIMEHeader() (MIMEHeader, error) {
                return m, ProtocolError("malformed MIME header initial line: " + string(line))
        }

+       // Create a pre-allocated byte slice for all header values to save
+       // allocations for the first 512 bytes of values, a size that fit typical
+       // small header values, larger ones will get their own allocation.
+       const valuesPreAllocSize = 1 << 9
+       valuesPreAlloc := make([]byte, 0, valuesPreAllocSize)
+
        for {
                kv, err := r.readContinuedLineSlice(mustHaveFieldNameColon)
                if len(kv) == 0 {
@@ -527,7 +534,17 @@ func (r *Reader) ReadMIMEHeader() (MIMEHeader, error) {
                for i < len(kv) && (kv[i] == ' ' || kv[i] == '\t') {
                        i++
                }
-               value := string(kv[i:])
+
+               // Try to fit the value in the pre-allocated buffer to save allocations.
+               var value string
+               if len(kv[i:]) <= valuesPreAllocSize-len(valuesPreAlloc) {
+                       off := len(valuesPreAlloc)
+                       valuesPreAlloc = append(valuesPreAlloc, kv[i:]...)
+                       v := valuesPreAlloc[off:]
+                       value = *(*string)(unsafe.Pointer(&v))
+               } else {
+                       value = string(kv[i:])
+               }

                vv := m[key]
                if vv == nil && len(strs) > 0 {

Please tell me if it's worth submitting a PR with this change.

[toothrot]

/cc @bradfitz @rsc

[smasher164]

I wonder if it's possible to make use of strings.Builder and pre-allocating with Grow(valuesPreAllocSize), instead of depending on unsafe and only using the buffer when the value fits inside it.

[bradfitz]

We're not adding unsafe usage to textproto.

[rs]

With or without strings.Builder, this trick won't work if you grow the slice, as a copy of the backing array will be performed, which would be a waste: the strings created before the grow would still point on the old backing array, copying data for nothing and wasting the first half of the new grown backing array.

I tried a slightly more complex implementation that would cover requests with arbitrary number header values, but it added no gain from my training set of real-life requests/responses.

We could still use strings.Builder to hide the use of unsafe, but I'm not sure to see the point.

[rs]

We're not adding unsafe usage to textproto.

@bradfitz would it be accepted if hidden by the use of strings.Builder?

EDIT: scratch that, this can't be implemented with strings.Buffer. So I guess it's a no-go.

[rs]

@bradfitz would you be open to adding the Slice(off, len int) string method to strings.Builder to allow implementation of such tricks without the need of unsafe?

[bcmills]

@rs, could you write all of the values into the strings.Builder and record the offsets, then convert all of the offsets to slices of the final string result?

[rs]

That could work yes, but it would make the code slightly more complex. I would personally prefer the Slice method on strings.Builder, but I guess it will require some more convincing :)

[josharian]

Another option (which would also complicate the code) is to have a fixed size, stack-allocated buffer for pending values to record and for the bytes involved. When a new kv pair is available, append the bytes to the bytes buffer and indices in the bytes buffer to the k/v pairs buffer. When either buffer is too full, convert the byte buffer to a string and add slices of that string to the map. This could all be wrapped up methods in a separate type, for sanity. (This is similar to the strings.Builder idea, but without the extra strings.Builder layer, since a plain ol' byte slice should suffice.)

Also, if the keys are commonly re-used, string interning might help.

[rs]

Good point, you would still pay for one extra copy, but it's better than nothing. String interning might work for some very common header values, I guess it could save a couple alloc per request with common content types and accept header values for instance.

Before spending more time on this, I'd love to hear from @bradfitz about the strings.Builder.Slice idea and if not, what is the probability to having more complex implementation as suggested above accepted.

[bcmills]

@rs, strings.Builder.Slice seems too prone to misuse: much too easy to accidentally create multiple underlying strings (due to growth) when the user really only intended to retain slices of a single big string.

[rs]

Agreed but it would still be safe and you would still save allocations by creating those substring from a common buffer.

[kokizzu]

Possibly related, this allocation grows to 200MB over a day or a week and cannot be reclaimed by GC: