github.com/daeuniverse/quic-go@v0.0.0-20240413031024-943f218e0810/buffer_pool.go (about) 1 package quic 2 3 import ( 4 "sync" 5 6 "github.com/daeuniverse/quic-go/internal/protocol" 7 ) 8 9 type packetBuffer struct { 10 Data []byte 11 12 // refCount counts how many packets Data is used in. 13 // It doesn't support concurrent use. 14 // It is > 1 when used for coalesced packet. 15 refCount int 16 } 17 18 // Split increases the refCount. 19 // It must be called when a packet buffer is used for more than one packet, 20 // e.g. when splitting coalesced packets. 21 func (b *packetBuffer) Split() { 22 b.refCount++ 23 } 24 25 // Decrement decrements the reference counter. 26 // It doesn't put the buffer back into the pool. 27 func (b *packetBuffer) Decrement() { 28 b.refCount-- 29 if b.refCount < 0 { 30 panic("negative packetBuffer refCount") 31 } 32 } 33 34 // MaybeRelease puts the packet buffer back into the pool, 35 // if the reference counter already reached 0. 36 func (b *packetBuffer) MaybeRelease() { 37 // only put the packetBuffer back if it's not used any more 38 if b.refCount == 0 { 39 b.putBack() 40 } 41 } 42 43 // Release puts back the packet buffer into the pool. 44 // It should be called when processing is definitely finished. 45 func (b *packetBuffer) Release() { 46 b.Decrement() 47 if b.refCount != 0 { 48 panic("packetBuffer refCount not zero") 49 } 50 b.putBack() 51 } 52 53 // Len returns the length of Data 54 func (b *packetBuffer) Len() protocol.ByteCount { return protocol.ByteCount(len(b.Data)) } 55 func (b *packetBuffer) Cap() protocol.ByteCount { return protocol.ByteCount(cap(b.Data)) } 56 57 func (b *packetBuffer) putBack() { 58 if cap(b.Data) == protocol.MaxPacketBufferSize { 59 bufferPool.Put(b) 60 return 61 } 62 if cap(b.Data) == protocol.MaxLargePacketBufferSize { 63 largeBufferPool.Put(b) 64 return 65 } 66 panic("putPacketBuffer called with packet of wrong size!") 67 } 68 69 var bufferPool, largeBufferPool sync.Pool 70 71 func getPacketBuffer() *packetBuffer { 72 buf := bufferPool.Get().(*packetBuffer) 73 buf.refCount = 1 74 buf.Data = buf.Data[:0] 75 return buf 76 } 77 78 func getLargePacketBuffer() *packetBuffer { 79 buf := largeBufferPool.Get().(*packetBuffer) 80 buf.refCount = 1 81 buf.Data = buf.Data[:0] 82 return buf 83 } 84 85 func init() { 86 bufferPool.New = func() any { 87 return &packetBuffer{Data: make([]byte, 0, protocol.MaxPacketBufferSize)} 88 } 89 largeBufferPool.New = func() any { 90 return &packetBuffer{Data: make([]byte, 0, protocol.MaxLargePacketBufferSize)} 91 } 92 }