lib/protocol: Alwasy return buffers to the pool (#7409)
This commit is contained in:
@@ -536,6 +536,7 @@ func (c *rawConnection) readMessageAfterHeader(hdr Header, fourByteBuf []byte) (
|
|||||||
|
|
||||||
buf := BufferPool.Get(int(msgLen))
|
buf := BufferPool.Get(int(msgLen))
|
||||||
if _, err := io.ReadFull(c.cr, buf); err != nil {
|
if _, err := io.ReadFull(c.cr, buf); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return nil, errors.Wrap(err, "reading message")
|
return nil, errors.Wrap(err, "reading message")
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -561,9 +562,11 @@ func (c *rawConnection) readMessageAfterHeader(hdr Header, fourByteBuf []byte) (
|
|||||||
|
|
||||||
msg, err := c.newMessage(hdr.Type)
|
msg, err := c.newMessage(hdr.Type)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if err := msg.Unmarshal(buf); err != nil {
|
if err := msg.Unmarshal(buf); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return nil, errors.Wrap(err, "unmarshalling message")
|
return nil, errors.Wrap(err, "unmarshalling message")
|
||||||
}
|
}
|
||||||
BufferPool.Put(buf)
|
BufferPool.Put(buf)
|
||||||
@@ -586,15 +589,17 @@ func (c *rawConnection) readHeader(fourByteBuf []byte) (Header, error) {
|
|||||||
|
|
||||||
buf := BufferPool.Get(int(hdrLen))
|
buf := BufferPool.Get(int(hdrLen))
|
||||||
if _, err := io.ReadFull(c.cr, buf); err != nil {
|
if _, err := io.ReadFull(c.cr, buf); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return Header{}, errors.Wrap(err, "reading header")
|
return Header{}, errors.Wrap(err, "reading header")
|
||||||
}
|
}
|
||||||
|
|
||||||
var hdr Header
|
var hdr Header
|
||||||
if err := hdr.Unmarshal(buf); err != nil {
|
err := hdr.Unmarshal(buf)
|
||||||
|
BufferPool.Put(buf)
|
||||||
|
if err != nil {
|
||||||
return Header{}, errors.Wrap(err, "unmarshalling header")
|
return Header{}, errors.Wrap(err, "unmarshalling header")
|
||||||
}
|
}
|
||||||
|
|
||||||
BufferPool.Put(buf)
|
|
||||||
return hdr, nil
|
return hdr, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -767,11 +772,13 @@ func (c *rawConnection) writeCompressedMessage(msg message) error {
|
|||||||
size := msg.ProtoSize()
|
size := msg.ProtoSize()
|
||||||
buf := BufferPool.Get(size)
|
buf := BufferPool.Get(size)
|
||||||
if _, err := msg.MarshalTo(buf); err != nil {
|
if _, err := msg.MarshalTo(buf); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return errors.Wrap(err, "marshalling message")
|
return errors.Wrap(err, "marshalling message")
|
||||||
}
|
}
|
||||||
|
|
||||||
compressed, err := c.lz4Compress(buf)
|
compressed, err := c.lz4Compress(buf)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return errors.Wrap(err, "compressing message")
|
return errors.Wrap(err, "compressing message")
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -784,17 +791,20 @@ func (c *rawConnection) writeCompressedMessage(msg message) error {
|
|||||||
panic("impossibly large header")
|
panic("impossibly large header")
|
||||||
}
|
}
|
||||||
|
|
||||||
totSize := 2 + hdrSize + 4 + len(compressed)
|
compressedSize := len(compressed)
|
||||||
|
totSize := 2 + hdrSize + 4 + compressedSize
|
||||||
buf = BufferPool.Upgrade(buf, totSize)
|
buf = BufferPool.Upgrade(buf, totSize)
|
||||||
|
|
||||||
// Header length
|
// Header length
|
||||||
binary.BigEndian.PutUint16(buf, uint16(hdrSize))
|
binary.BigEndian.PutUint16(buf, uint16(hdrSize))
|
||||||
// Header
|
// Header
|
||||||
if _, err := hdr.MarshalTo(buf[2:]); err != nil {
|
if _, err := hdr.MarshalTo(buf[2:]); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
|
BufferPool.Put(compressed)
|
||||||
return errors.Wrap(err, "marshalling header")
|
return errors.Wrap(err, "marshalling header")
|
||||||
}
|
}
|
||||||
// Message length
|
// Message length
|
||||||
binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(len(compressed)))
|
binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(compressedSize))
|
||||||
// Message
|
// Message
|
||||||
copy(buf[2+hdrSize+4:], compressed)
|
copy(buf[2+hdrSize+4:], compressed)
|
||||||
BufferPool.Put(compressed)
|
BufferPool.Put(compressed)
|
||||||
@@ -802,7 +812,7 @@ func (c *rawConnection) writeCompressedMessage(msg message) error {
|
|||||||
n, err := c.cw.Write(buf)
|
n, err := c.cw.Write(buf)
|
||||||
BufferPool.Put(buf)
|
BufferPool.Put(buf)
|
||||||
|
|
||||||
l.Debugf("wrote %d bytes on the wire (2 bytes length, %d bytes header, 4 bytes message length, %d bytes message (%d uncompressed)), err=%v", n, hdrSize, len(compressed), size, err)
|
l.Debugf("wrote %d bytes on the wire (2 bytes length, %d bytes header, 4 bytes message length, %d bytes message (%d uncompressed)), err=%v", n, hdrSize, compressedSize, size, err)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return errors.Wrap(err, "writing message")
|
return errors.Wrap(err, "writing message")
|
||||||
}
|
}
|
||||||
@@ -827,12 +837,14 @@ func (c *rawConnection) writeUncompressedMessage(msg message) error {
|
|||||||
binary.BigEndian.PutUint16(buf, uint16(hdrSize))
|
binary.BigEndian.PutUint16(buf, uint16(hdrSize))
|
||||||
// Header
|
// Header
|
||||||
if _, err := hdr.MarshalTo(buf[2:]); err != nil {
|
if _, err := hdr.MarshalTo(buf[2:]); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return errors.Wrap(err, "marshalling header")
|
return errors.Wrap(err, "marshalling header")
|
||||||
}
|
}
|
||||||
// Message length
|
// Message length
|
||||||
binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(size))
|
binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(size))
|
||||||
// Message
|
// Message
|
||||||
if _, err := msg.MarshalTo(buf[2+hdrSize+4:]); err != nil {
|
if _, err := msg.MarshalTo(buf[2+hdrSize+4:]); err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return errors.Wrap(err, "marshalling message")
|
return errors.Wrap(err, "marshalling message")
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1033,6 +1045,7 @@ func (c *rawConnection) lz4Compress(src []byte) ([]byte, error) {
|
|||||||
buf := BufferPool.Get(lz4.CompressBound(len(src)))
|
buf := BufferPool.Get(lz4.CompressBound(len(src)))
|
||||||
compressed, err := lz4.Encode(buf, src)
|
compressed, err := lz4.Encode(buf, src)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if &compressed[0] != &buf[0] {
|
if &compressed[0] != &buf[0] {
|
||||||
@@ -1050,6 +1063,7 @@ func (c *rawConnection) lz4Decompress(src []byte) ([]byte, error) {
|
|||||||
buf := BufferPool.Get(int(size))
|
buf := BufferPool.Get(int(size))
|
||||||
decoded, err := lz4.Decode(buf, src)
|
decoded, err := lz4.Decode(buf, src)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
BufferPool.Put(buf)
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if &decoded[0] != &buf[0] {
|
if &decoded[0] != &buf[0] {
|
||||||
|
|||||||
Reference in New Issue
Block a user