1 files changed, 23 insertions, 5 deletions
diff --git a/vendor/google.golang.org/grpc/preloader.go b/vendor/google.golang.org/grpc/preloader.go
index 73bd63364..e87a17f36 100644
--- a/vendor/google.golang.org/grpc/preloader.go
+++ b/vendor/google.golang.org/grpc/preloader.go
@@ -20,6 +20,7 @@ package grpc
 
 import (
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/mem"
 	"google.golang.org/grpc/status"
 )
 
@@ -31,9 +32,10 @@ import (
 // later release.
 type PreparedMsg struct {
 	// Struct for preparing msg before sending them
-	encodedData []byte
+	encodedData mem.BufferSlice
 	hdr         []byte
-	payload     []byte
+	payload     mem.BufferSlice
+	pf          payloadFormat
 }
 
 // Encode marshalls and compresses the message using the codec and compressor for the stream.
@@ -57,11 +59,27 @@ func (p *PreparedMsg) Encode(s Stream, msg any) error {
 	if err != nil {
 		return err
 	}
-	p.encodedData = data
-	compData, err := compress(data, rpcInfo.preloaderInfo.cp, rpcInfo.preloaderInfo.comp)
+
+	materializedData := data.Materialize()
+	data.Free()
+	p.encodedData = mem.BufferSlice{mem.NewBuffer(&materializedData, nil)}
+
+	// TODO: it should be possible to grab the bufferPool from the underlying
+	//  stream implementation with a type cast to its actual type (such as
+	//  addrConnStream) and accessing the buffer pool directly.
+	var compData mem.BufferSlice
+	compData, p.pf, err = compress(p.encodedData, rpcInfo.preloaderInfo.cp, rpcInfo.preloaderInfo.comp, mem.DefaultBufferPool())
 	if err != nil {
 		return err
 	}
-	p.hdr, p.payload = msgHeader(data, compData)
+
+	if p.pf.isCompressed() {
+		materializedCompData := compData.Materialize()
+		compData.Free()
+		compData = mem.BufferSlice{mem.NewBuffer(&materializedCompData, nil)}
+	}
+
+	p.hdr, p.payload = msgHeader(p.encodedData, compData, p.pf)
+
 	return nil
 }