implement a buffer pool for STREAM frames

internal/wire/stream_frame.go

@@ -17,6 +17,8 @@ type StreamFrame struct {
 	DataLenPresent bool
 	Offset         protocol.ByteCount
 	Data           []byte
+
+	fromPool bool
 }
 
 func parseStreamFrame(r *bytes.Reader, version protocol.VersionNumber) (*StreamFrame, error) {
@@ -26,45 +28,53 @@ func parseStreamFrame(r *bytes.Reader, version protocol.VersionNumber) (*StreamF
 	}
 
 	hasOffset := typeByte&0x4 > 0
-	frame := &StreamFrame{
-		FinBit:         typeByte&0x1 > 0,
-		DataLenPresent: typeByte&0x2 > 0,
-	}
+	fin := typeByte&0x1 > 0
+	hasDataLen := typeByte&0x2 > 0
 
 	streamID, err := utils.ReadVarInt(r)
 	if err != nil {
 		return nil, err
 	}
-	frame.StreamID = protocol.StreamID(streamID)
+	var offset uint64
 	if hasOffset {
-		offset, err := utils.ReadVarInt(r)
+		offset, err = utils.ReadVarInt(r)
 		if err != nil {
 			return nil, err
 		}
-		frame.Offset = protocol.ByteCount(offset)
 	}
 
 	var dataLen uint64
-	if frame.DataLenPresent {
+	if hasDataLen {
 		var err error
 		dataLen, err = utils.ReadVarInt(r)
 		if err != nil {
 			return nil, err
 		}
-		// shortcut to prevent the unnecessary allocation of dataLen bytes
-		// if the dataLen is larger than the remaining length of the packet
-		// reading the packet contents would result in EOF when attempting to READ
-		if dataLen > uint64(r.Len()) {
-			return nil, io.EOF
-		}
 	} else {
 		// The rest of the packet is data
 		dataLen = uint64(r.Len())
 	}
+
+	var frame *StreamFrame
+	if dataLen < protocol.MinStreamFrameBufferSize {
+		frame = &StreamFrame{Data: make([]byte, dataLen)}
+	} else {
+		frame = getStreamFrame()
+		// The STREAM frame can't be larger than the StreamFrame we obtained from the buffer,
+		// since those StreamFrames have a buffer length of the maximum packet size.
+		if dataLen > uint64(cap(frame.Data)) {
+			return nil, io.EOF
+		}
+		frame.Data = frame.Data[:dataLen]
+	}
+
+	frame.StreamID = protocol.StreamID(streamID)
+	frame.Offset = protocol.ByteCount(offset)
+	frame.FinBit = fin
+	frame.DataLenPresent = hasDataLen
+
 	if dataLen != 0 {
-		frame.Data = make([]byte, dataLen)
 		if _, err := io.ReadFull(r, frame.Data); err != nil {
-			// this should never happen, since we already checked the dataLen earlier
 			return nil, err
 		}
 	}
@@ -156,16 +166,25 @@ func (f *StreamFrame) MaybeSplitOffFrame(maxSize protocol.ByteCount, version pro
 	if n == 0 {
 		return nil, true
 	}
-	newFrame := &StreamFrame{
-		FinBit:         false,
-		StreamID:       f.StreamID,
-		Offset:         f.Offset,
-		Data:           f.Data[:n],
-		DataLenPresent: f.DataLenPresent,
-	}
 
-	f.Data = f.Data[n:]
+	new := getStreamFrame()
+	new.StreamID = f.StreamID
+	new.Offset = f.Offset
+	new.FinBit = false
+	new.DataLenPresent = f.DataLenPresent
+
+	// swap the data slices
+	new.Data, f.Data = f.Data, new.Data
+	new.fromPool, f.fromPool = f.fromPool, new.fromPool
+
+	f.Data = f.Data[:protocol.ByteCount(len(new.Data))-n]
+	copy(f.Data, new.Data[n:])
+	new.Data = new.Data[:n]
 	f.Offset += n
 
-	return newFrame, true
+	return new, true
+}
+
+func (f *StreamFrame) PutBack() {
+	putStreamFrame(f)
 }