github.com/ooni/psiphon/tunnel-core@v0.0.0-20230105123940-fe12a24c96ee/oovendor/quic-go/buffer_pool.go (about)

     1  package quic
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/ooni/psiphon/tunnel-core/oovendor/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 {
    55  	return protocol.ByteCount(len(b.Data))
    56  }
    57  
    58  func (b *packetBuffer) putBack() {
    59  	if cap(b.Data) != int(protocol.MaxPacketBufferSize) {
    60  		panic("putPacketBuffer called with packet of wrong size!")
    61  	}
    62  	bufferPool.Put(b)
    63  }
    64  
    65  var bufferPool sync.Pool
    66  
    67  func getPacketBuffer() *packetBuffer {
    68  	buf := bufferPool.Get().(*packetBuffer)
    69  	buf.refCount = 1
    70  	buf.Data = buf.Data[:0]
    71  	return buf
    72  }
    73  
    74  func init() {
    75  	bufferPool.New = func() interface{} {
    76  		return &packetBuffer{
    77  			Data: make([]byte, 0, protocol.MaxPacketBufferSize),
    78  		}
    79  	}
    80  }