buffer_pool.go 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. package quic
  2. import (
  3. "sync"
  4. "github.com/lucas-clemente/quic-go/internal/protocol"
  5. )
  6. type packetBuffer struct {
  7. Slice []byte
  8. // refCount counts how many packets the Slice is used in.
  9. // It doesn't support concurrent use.
  10. // It is > 1 when used for coalesced packet.
  11. refCount int
  12. }
  13. // Split increases the refCount.
  14. // It must be called when a packet buffer is used for more than one packet,
  15. // e.g. when splitting coalesced packets.
  16. func (b *packetBuffer) Split() {
  17. b.refCount++
  18. }
  19. // Decrement decrements the reference counter.
  20. // It doesn't put the buffer back into the pool.
  21. func (b *packetBuffer) Decrement() {
  22. b.refCount--
  23. if b.refCount < 0 {
  24. panic("negative packetBuffer refCount")
  25. }
  26. }
  27. // MaybeRelease puts the packet buffer back into the pool,
  28. // if the reference counter already reached 0.
  29. func (b *packetBuffer) MaybeRelease() {
  30. // only put the packetBuffer back if it's not used any more
  31. if b.refCount == 0 {
  32. b.putBack()
  33. }
  34. }
  35. // Release puts back the packet buffer into the pool.
  36. // It should be called when processing is definitely finished.
  37. func (b *packetBuffer) Release() {
  38. b.Decrement()
  39. if b.refCount != 0 {
  40. panic("packetBuffer refCount not zero")
  41. }
  42. b.putBack()
  43. }
  44. func (b *packetBuffer) putBack() {
  45. if cap(b.Slice) != int(protocol.MaxReceivePacketSize) {
  46. panic("putPacketBuffer called with packet of wrong size!")
  47. }
  48. bufferPool.Put(b)
  49. }
  50. var bufferPool sync.Pool
  51. func getPacketBuffer() *packetBuffer {
  52. buf := bufferPool.Get().(*packetBuffer)
  53. buf.refCount = 1
  54. buf.Slice = buf.Slice[:protocol.MaxReceivePacketSize]
  55. return buf
  56. }
  57. func init() {
  58. bufferPool.New = func() interface{} {
  59. return &packetBuffer{
  60. Slice: make([]byte, 0, protocol.MaxReceivePacketSize),
  61. }
  62. }
  63. }