Reduce marshaling allocations with MarshalAppend

[emcfarlane]

Use the MarshalAppend methods introduced in v1.31 of the protobuf runtime to more aggressively reuse buffers from Connect's internal pool (for both binary and JSON). This reduces allocations and substantially improves throughput:

BenchmarkConnect/unary-12                 8054           1318347 ns/op        14414394 B/op        251 allocs/op
BenchmarkConnect/unary-12                 10448          1153636 ns/op         6115680 B/op        236 allocs/op

[CLAassistant]

All committers have signed the CLA.

[akshayjshah]

Thanks for the PR, @emcfarlane! The performance win is definitely nice :)

I don't love adding more optional methods to Codec, but the only alternative is to have connect's internal codecs cheat and use the client and handler's buffer pools. On balance, the approach you have is probably the better option.

Before merging this, I'd like:

• To wait for a released version of the protobuf runtime to have all the necessary features, and
• To have tests (probably using testing/quick) verifying that a MarshalAppend followed by an Unmarshal leaves the original protobuf message unchanged.

[akshayjshah]

I think performance along this path has gotten worse, since we no longer put raw into the pool when we're done with it. Adding a bit of code to the else block should fix it:

if len(raw) > 0 {
  defer w.bufferPool.Put(bytes.NewBuffer(raw))
}

[akshayjshah]

Even with this change, I think we're failing to effectively pool here. At base, the buffers in the pool are pretty small. MarshalAppend is likely to grow the slice, in which case returning buffer to the pool no longer helps much. Give me a moment to play with this a bit locally and find a way to fix this that isn't too prone to bugs down the road.

[pkwarren]

I think it will do the right thing - doing buffer.Write(raw) will grow the buffer as needed, so it will be returned to the pool with a larger size. I think we can remove adding an additional bytes.Buffer to the pool.

[emcfarlane]

Protobuf v1.31.0 just released! Updated the library and added a test case.

[pkwarren]

Changes look great - just a couple of small things. Looking forward to the perf improvements.

[pkwarren]

Looks great to me - leaving for @akshayjshah to take another look.

[akshayjshah]

Even with this change, I think we're failing to effectively pool here. At base, the buffers in the pool are pretty small. MarshalAppend is likely to grow the slice, in which case returning buffer to the pool no longer helps much. Give me a moment to play with this a bit locally and find a way to fix this that isn't too prone to bugs down the road.

[akshayjshah]

Made a small change to allow the buffers in the pool to grow if they're not large enough; this change looks good to me now. @emcfarlane can you take a look and ensure that it's still ok with you? I'd like to get this into the next release if possible, since we're already planning to release some performance improvements.

[emcfarlane]

@akshayjshah I don't think this defer should be optional. We always want to defer the fetched buffer.

buffer := w.bufferPool.Get()
defer w.bufferPool.Put(buffer)

[akshayjshah]

Do we? For best performance, we want the buffers in the pool to be large enough to contain the marshaled bytes without resizing (at least most of the time). We know that the buffer we've pulled from the pool was too small for this message. I think we're better off letting the too-small buffer be GC'ed and putting the new buffer into the pool instead, under the assumption that we need the buffers in the pool to grow to accommodate the application workload.

None of this is perfect, but I think we're probably better off pooling over- rather than under-aggressively. The GC is still free to empty the pool on each cycle if it needs to reclaim the space. (Last I looked, I think it just categorically emptied pools...but I haven't checked in many years.)

[akshayjshah]

🤦🏽 @pkwarren pointed out to me that we're calling buffer.Write just a few lines below, which allocates the same amount of memory all over again. That makes returning the original buffer to the pool fine, but also means we're eating allocs that we definitely don't need.

[emcfarlane]

Calling bytes.NewBuffer(largerThanBuf) will cause an allocation like buffer.Write(largerThanBuf) but has the chance for reuse. Maybe it's better to split these two methods to keep the old behaviour

	var buffer *bytes.Buffer
	defer w.bufferPool.Put(buffer)

	switch codec := w.codec.(type) {
	case marshalAppender:
		buffer = w.bufferPool.Get()
		raw, err := codec.MarshalAppend(buffer.Bytes(), message)
		if err != nil {
			return errorf(CodeInternal, "marshal append message: %w", err)
		}
		buffer.Write(raw)
	default:
		raw, err := w.codec.Marshal(message)
		if err != nil {
			return errorf(CodeInternal, "marshal message: %w", err)
		}
		// We can't avoid allocating the byte slice, so we may as well reuse it once
		// we're done with it.
		buffer = bytes.NewBuffer(raw)
	}
	envelope := &envelope{Data: buffer}
	return w.Write(envelope)

[akshayjshah]

Did something similar-ish in the next commit to more clearly separate the two paths. I think I've gotten the best of both worlds: we avoid allocating if MarshalAppend didn't grow the slice capacity, and we eat a fixed-size allocation instead of an O(N) allocation if it did. Performance improved slightly:

[main] BenchmarkConnect/unary-12                   8054           1318347 ns/op        14414394 B/op        251 allocs/op
[original PR] BenchmarkConnect/unary-12            10192          1176274 ns/op         6122383 B/op        236 allocs/op
[latest] BenchmarkConnect/unary-12                 10448          1153636 ns/op         6115680 B/op        236 allocs/op

[emcfarlane]

Looks good to me!