bufferpool.go 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. // Copyright (C) 2016 The Protocol Authors.
  2. package protocol
  3. import "sync"
  4. type bufferPool struct {
  5. minSize int
  6. pool sync.Pool
  7. }
  8. // get returns a new buffer of the requested size
  9. func (p *bufferPool) get(size int) []byte {
  10. intf := p.pool.Get()
  11. if intf == nil {
  12. // Pool is empty, must allocate.
  13. return p.new(size)
  14. }
  15. bs := intf.([]byte)
  16. if cap(bs) < size {
  17. // Buffer was too small, leave it for someone else and allocate.
  18. p.put(bs)
  19. return p.new(size)
  20. }
  21. return bs[:size]
  22. }
  23. // upgrade grows the buffer to the requested size, while attempting to reuse
  24. // it if possible.
  25. func (p *bufferPool) upgrade(bs []byte, size int) []byte {
  26. if cap(bs) >= size {
  27. // Reslicing is enough, lets go!
  28. return bs[:size]
  29. }
  30. // It was too small. But it pack into the pool and try to get another
  31. // buffer.
  32. p.put(bs)
  33. return p.get(size)
  34. }
  35. // put returns the buffer to the pool
  36. func (p *bufferPool) put(bs []byte) {
  37. p.pool.Put(bs)
  38. }
  39. // new creates a new buffer of the requested size, taking the minimum
  40. // allocation count into account. For internal use only.
  41. func (p *bufferPool) new(size int) []byte {
  42. allocSize := size
  43. if allocSize < p.minSize {
  44. // Avoid allocating tiny buffers that we won't be able to reuse for
  45. // anything useful.
  46. allocSize = p.minSize
  47. }
  48. return make([]byte, allocSize)[:size]
  49. }