github.com/quic-go/quic-go@v0.44.0/buffer_pool.go (about)

     1  package quic
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/quic-go/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  }