iobufpool.go (1229B)
1 // Package iobufpool implements a global segregated-fit pool of buffers for IO. 2 // 3 // It uses *[]byte instead of []byte to avoid the sync.Pool allocation with Put. Unfortunately, using a pointer to avoid 4 // an allocation is purposely not documented. https://github.com/golang/go/issues/16323 5 package iobufpool 6 7 import "sync" 8 9 const minPoolExpOf2 = 8 10 11 var pools [18]*sync.Pool 12 13 func init() { 14 for i := range pools { 15 bufLen := 1 << (minPoolExpOf2 + i) 16 pools[i] = &sync.Pool{ 17 New: func() any { 18 buf := make([]byte, bufLen) 19 return &buf 20 }, 21 } 22 } 23 } 24 25 // Get gets a []byte of len size with cap <= size*2. 26 func Get(size int) *[]byte { 27 i := getPoolIdx(size) 28 if i >= len(pools) { 29 buf := make([]byte, size) 30 return &buf 31 } 32 33 ptrBuf := (pools[i].Get().(*[]byte)) 34 *ptrBuf = (*ptrBuf)[:size] 35 36 return ptrBuf 37 } 38 39 func getPoolIdx(size int) int { 40 size-- 41 size >>= minPoolExpOf2 42 i := 0 43 for size > 0 { 44 size >>= 1 45 i++ 46 } 47 48 return i 49 } 50 51 // Put returns buf to the pool. 52 func Put(buf *[]byte) { 53 i := putPoolIdx(cap(*buf)) 54 if i < 0 { 55 return 56 } 57 58 pools[i].Put(buf) 59 } 60 61 func putPoolIdx(size int) int { 62 minPoolSize := 1 << minPoolExpOf2 63 for i := range pools { 64 if size == minPoolSize<<i { 65 return i 66 } 67 } 68 69 return -1 70 }