bufferpool.go 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. // Copyright (C) 2016 The Protocol Authors.
  2. package protocol
  3. import "sync"
  4. // Global pool to get buffers from. Requires Blocksizes to be initialised,
  5. // therefore it is initialized in the same init() as BlockSizes
  6. var BufferPool bufferPool
  7. type bufferPool struct {
  8. pools []sync.Pool
  9. }
  10. func newBufferPool() bufferPool {
  11. return bufferPool{make([]sync.Pool, len(BlockSizes))}
  12. }
  13. func (p *bufferPool) Get(size int) []byte {
  14. // Too big, isn't pooled
  15. if size > MaxBlockSize {
  16. return make([]byte, size)
  17. }
  18. var i int
  19. for i = range BlockSizes {
  20. if size <= BlockSizes[i] {
  21. break
  22. }
  23. }
  24. var bs []byte
  25. // Try the fitting and all bigger pools
  26. for j := i; j < len(BlockSizes); j++ {
  27. if intf := p.pools[j].Get(); intf != nil {
  28. bs = *intf.(*[]byte)
  29. return bs[:size]
  30. }
  31. }
  32. // All pools are empty, must allocate.
  33. return make([]byte, BlockSizes[i])[:size]
  34. }
  35. // Put makes the given byte slice availabe again in the global pool
  36. func (p *bufferPool) Put(bs []byte) {
  37. c := cap(bs)
  38. // Don't buffer huge byte slices
  39. if c > 2*MaxBlockSize {
  40. return
  41. }
  42. for i := range BlockSizes {
  43. if c >= BlockSizes[i] {
  44. p.pools[i].Put(&bs)
  45. return
  46. }
  47. }
  48. }
  49. // Upgrade grows the buffer to the requested size, while attempting to reuse
  50. // it if possible.
  51. func (p *bufferPool) Upgrade(bs []byte, size int) []byte {
  52. if cap(bs) >= size {
  53. // Reslicing is enough, lets go!
  54. return bs[:size]
  55. }
  56. // It was too small. But it pack into the pool and try to get another
  57. // buffer.
  58. p.Put(bs)
  59. return p.Get(size)
  60. }