buffer_pool.go 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. package quic
  2. import (
  3. "sync"
  4. "github.com/lucas-clemente/quic-go/internal/protocol"
  5. "v2ray.com/core/common/bytespool"
  6. )
  7. type packetBuffer struct {
  8. Slice []byte
  9. // refCount counts how many packets the Slice is used in.
  10. // It doesn't support concurrent use.
  11. // It is > 1 when used for coalesced packet.
  12. refCount int
  13. }
  14. // Split increases the refCount.
  15. // It must be called when a packet buffer is used for more than one packet,
  16. // e.g. when splitting coalesced packets.
  17. func (b *packetBuffer) Split() {
  18. b.refCount++
  19. }
  20. // Release decreases the refCount.
  21. // It should be called when processing the packet is finished.
  22. // When the refCount reaches 0, the packet buffer is put back into the pool.
  23. func (b *packetBuffer) Release() {
  24. if cap(b.Slice) < 2048 {
  25. return
  26. }
  27. b.refCount--
  28. if b.refCount < 0 {
  29. panic("negative packetBuffer refCount")
  30. }
  31. // only put the packetBuffer back if it's not used any more
  32. if b.refCount == 0 {
  33. buffer := b.Slice[0:cap(b.Slice)]
  34. bufferPool.Put(buffer)
  35. }
  36. }
  37. var bufferPool *sync.Pool
  38. func getPacketBuffer() *packetBuffer {
  39. buffer := bufferPool.Get().([]byte)
  40. return &packetBuffer{
  41. refCount: 1,
  42. Slice: buffer[:protocol.MaxReceivePacketSize],
  43. }
  44. }
  45. func init() {
  46. bufferPool = bytespool.GetPool(int32(protocol.MaxReceivePacketSize))
  47. }