@@ -496,6 +496,16 @@ const (
496496 serverSide
497497)
498498
499+ // maxWriteBufSize is the maximum length (number of elements) the cached
500+ // writeBuf can grow to. The length depends on the number of buffers
501+ // contained within the BufferSlice produced by the codec, which is
502+ // generally small.
503+ //
504+ // If a writeBuf larger than this limit is required, it will be allocated
505+ // and freed after use, rather than being cached. This avoids holding
506+ // on to large amounts of memory.
507+ const maxWriteBufSize = 64
508+
499509// Loopy receives frames from the control buffer.
500510// Each frame is handled individually; most of the work done by loopy goes
501511// into handling data frames. Loopy maintains a queue of active streams, and each
@@ -530,6 +540,8 @@ type loopyWriter struct {
530540
531541 // Side-specific handlers
532542 ssGoAwayHandler func (* goAway ) (bool , error )
543+
544+ writeBuf [][]byte // cached slice to avoid heap allocations for calls to mem.Reader.Peek.
533545}
534546
535547func newLoopyWriter (s side , fr * framer , cbuf * controlBuffer , bdpEst * bdpEstimator , conn net.Conn , logger * grpclog.PrefixLogger , goAwayHandler func (* goAway ) (bool , error ), bufferPool mem.BufferPool ) * loopyWriter {
@@ -962,11 +974,11 @@ func (l *loopyWriter) processData() (bool, error) {
962974
963975 if len (dataItem .h ) == 0 && reader .Remaining () == 0 { // Empty data frame
964976 // Client sends out empty data frame with endStream = true
965- if err := l .framer .fr . WriteData (dataItem .streamID , dataItem .endStream , nil ); err != nil {
977+ if err := l .framer .writeData (dataItem .streamID , dataItem .endStream , nil ); err != nil {
966978 return false , err
967979 }
968980 str .itl .dequeue () // remove the empty data item from stream
969- _ = reader .Close ()
981+ reader .Close ()
970982 if str .itl .isEmpty () {
971983 str .state = empty
972984 } else if trailer , ok := str .itl .peek ().(* headerFrame ); ok { // the next item is trailers.
@@ -999,25 +1011,20 @@ func (l *loopyWriter) processData() (bool, error) {
9991011 remainingBytes := len (dataItem .h ) + reader .Remaining () - hSize - dSize
10001012 size := hSize + dSize
10011013
1002- var buf * [] byte
1003-
1004- if hSize != 0 && dSize == 0 {
1005- buf = & dataItem . h
1006- } else {
1007- // Note: this is only necessary because the http2.Framer does not support
1008- // partially writing a frame, so the sequence must be materialized into a buffer.
1009- // TODO: Revisit once https://github.com/golang/go/issues/66655 is addressed.
1010- pool := l . bufferPool
1011- if pool == nil {
1012- // Note that this is only supposed to be nil in tests. Otherwise, stream is
1013- // always initialized with a BufferPool.
1014- pool = mem . DefaultBufferPool ()
1014+ l . writeBuf = l . writeBuf [: 0 ]
1015+ if hSize > 0 {
1016+ l . writeBuf = append ( l . writeBuf , dataItem . h [: hSize ])
1017+ }
1018+ if dSize > 0 {
1019+ var err error
1020+ l . writeBuf , err = reader . Peek ( dSize , l . writeBuf )
1021+ if err != nil {
1022+ // This must never happen since the reader must have at least dSize
1023+ // bytes.
1024+ // Log an error to fail tests.
1025+ l . logger . Errorf ( "unexpected error while reading Data frame payload: %v" , err )
1026+ return false , err
10151027 }
1016- buf = pool .Get (size )
1017- defer pool .Put (buf )
1018-
1019- copy ((* buf )[:hSize ], dataItem .h )
1020- _ , _ = reader .Read ((* buf )[hSize :])
10211028 }
10221029
10231030 // Now that outgoing flow controls are checked we can replenish str's write quota
@@ -1030,15 +1037,22 @@ func (l *loopyWriter) processData() (bool, error) {
10301037 if dataItem .onEachWrite != nil {
10311038 dataItem .onEachWrite ()
10321039 }
1033- if err := l .framer .fr .WriteData (dataItem .streamID , endStream , (* buf )[:size ]); err != nil {
1040+ err := l .framer .writeData (dataItem .streamID , endStream , l .writeBuf )
1041+ reader .Discard (dSize )
1042+ if cap (l .writeBuf ) > maxWriteBufSize {
1043+ l .writeBuf = nil
1044+ } else {
1045+ clear (l .writeBuf )
1046+ }
1047+ if err != nil {
10341048 return false , err
10351049 }
10361050 str .bytesOutStanding += size
10371051 l .sendQuota -= uint32 (size )
10381052 dataItem .h = dataItem .h [hSize :]
10391053
10401054 if remainingBytes == 0 { // All the data from that message was written out.
1041- _ = reader .Close ()
1055+ reader .Close ()
10421056 str .itl .dequeue ()
10431057 }
10441058 if str .itl .isEmpty () {
0 commit comments