Release buffer for end-stream messages back to the pool

envelope.go

@@ -228,7 +228,12 @@ type envelopeReader struct {
 
 func (r *envelopeReader) Unmarshal(message any) *Error {
 	buffer := r.bufferPool.Get()
-	defer r.bufferPool.Put(buffer)
+	var dontRelease *bytes.Buffer
+	defer func() {
+		if buffer != dontRelease {
+			r.bufferPool.Put(buffer)
+		}
+	}()
 
 	env := &envelope{Data: buffer}
 	err := r.Read(env)
@@ -256,7 +261,11 @@ func (r *envelopeReader) Unmarshal(message any) *Error {
 			)
 		}
 		decompressed := r.bufferPool.Get()
-		defer r.bufferPool.Put(decompressed)
+		defer func() {
+			if decompressed != dontRelease {
+				r.bufferPool.Put(decompressed)
+			}
+		}()
 		if err := r.compressionPool.Decompress(decompressed, data, int64(r.readMaxBytes)); err != nil {
 			return err
 		}
@@ -276,14 +285,13 @@ func (r *envelopeReader) Unmarshal(message any) *Error {
 		}
 		// One of the protocol-specific flags are set, so this is the end of the
 		// stream. Save the message for protocol-specific code to process and
-		// return a sentinel error. Since we've deferred functions to return env's
-		// underlying buffer to a pool, we need to keep a copy.
-		copiedData := make([]byte, data.Len())
-		copy(copiedData, data.Bytes())
+		// return a sentinel error. We alias the buffer with dontRelease as a
+		// way of marking it so above defers don't release it to the pool.
 		r.last = envelope{
-			Data:  bytes.NewBuffer(copiedData),
+			Data:  data,
 			Flags: env.Flags,
 		}
+		dontRelease = data
 		return errSpecialEnvelope
 	}
 

protocol_connect.go

@@ -863,12 +863,15 @@ func (u *connectStreamingUnmarshaler) Unmarshal(message any) *Error {
 	if !errors.Is(err, errSpecialEnvelope) {
 		return err
 	}
-	env := u.envelopeReader.last
+	env := u.last
+	data := env.Data
+	u.last.Data = nil // don't keep a reference to it
+	defer u.bufferPool.Put(data)
 	if !env.IsSet(connectFlagEnvelopeEndStream) {
 		return errorf(CodeInternal, "protocol error: invalid envelope flags %d", env.Flags)
 	}
 	var end connectEndStreamMessage
-	if err := json.Unmarshal(env.Data.Bytes(), &end); err != nil {
+	if err := json.Unmarshal(data.Bytes(), &end); err != nil {
 		return errorf(CodeInternal, "unmarshal end stream message: %w", err)
 	}
 	for name, value := range end.Trailer {

protocol_grpc.go

@@ -602,9 +602,9 @@ func (m *grpcMarshaler) MarshalWebTrailers(trailer http.Header) *Error {
 }
 
 type grpcUnmarshaler struct {
-	envelopeReader envelopeReader
-	web            bool
-	webTrailer     http.Header
+	envelopeReader
+	web        bool
+	webTrailer http.Header
 }
 
 func (u *grpcUnmarshaler) Unmarshal(message any) *Error {
@@ -615,18 +615,21 @@ func (u *grpcUnmarshaler) Unmarshal(message any) *Error {
 	if !errors.Is(err, errSpecialEnvelope) {
 		return err
 	}
-	env := u.envelopeReader.last
+	env := u.last
+	data := env.Data
+	u.last.Data = nil // don't keep a reference to it
+	defer u.bufferPool.Put(data)
 	if !u.web || !env.IsSet(grpcFlagEnvelopeTrailer) {
 		return errorf(CodeInternal, "protocol error: invalid envelope flags %d", env.Flags)
 	}
 
 	// Per the gRPC-Web specification, trailers should be encoded as an HTTP/1
 	// headers block _without_ the terminating newline. To make the headers
 	// parseable by net/textproto, we need to add the newline.
-	if err := env.Data.WriteByte('\n'); err != nil {
+	if err := data.WriteByte('\n'); err != nil {
 		return errorf(CodeInternal, "unmarshal web trailers: %w", err)
 	}
-	bufferedReader := bufio.NewReader(env.Data)
+	bufferedReader := bufio.NewReader(data)
 	mimeReader := textproto.NewReader(bufferedReader)
 	mimeHeader, mimeErr := mimeReader.ReadMIMEHeader()
 	if mimeErr != nil {