github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/membuf/buffer.go (about)

     1  // Copyright 2021 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package membuf
    15  
    16  const bigValueSize = 1 << 16 // 64K
    17  
    18  var allocBufLen = 1 << 20 // 1M
    19  
    20  // Allocator is the abstract interface for allocating and freeing memory.
    21  type Allocator interface {
    22  	Alloc(n int) []byte
    23  	Free([]byte)
    24  }
    25  
    26  type stdAllocator struct{}
    27  
    28  func (stdAllocator) Alloc(n int) []byte {
    29  	return make([]byte, n)
    30  }
    31  
    32  func (stdAllocator) Free(_ []byte) {}
    33  
    34  // Pool is like `sync.Pool`, which manages memory for all bytes buffers.
    35  //
    36  // NOTE: we don't used a `sync.Pool` because when will sync.Pool release is depending on the
    37  // garbage collector which always release the memory so late. Use a fixed size chan to reuse
    38  // can decrease the memory usage to 1/3 compare with sync.Pool.
    39  type Pool struct {
    40  	allocator Allocator
    41  	recycleCh chan []byte
    42  }
    43  
    44  // NewPool creates a new pool.
    45  func NewPool(size int, allocator Allocator) *Pool {
    46  	return &Pool{
    47  		allocator: allocator,
    48  		recycleCh: make(chan []byte, size),
    49  	}
    50  }
    51  
    52  func (p *Pool) acquire() []byte {
    53  	select {
    54  	case b := <-p.recycleCh:
    55  		return b
    56  	default:
    57  		return p.allocator.Alloc(allocBufLen)
    58  	}
    59  }
    60  
    61  func (p *Pool) release(b []byte) {
    62  	select {
    63  	case p.recycleCh <- b:
    64  	default:
    65  		p.allocator.Free(b)
    66  	}
    67  }
    68  
    69  // NewBuffer creates a new buffer in current pool.
    70  func (p *Pool) NewBuffer() *Buffer {
    71  	return &Buffer{pool: p, bufs: make([][]byte, 0, 128), curBufIdx: -1}
    72  }
    73  
    74  var globalPool = NewPool(1024, stdAllocator{})
    75  
    76  // NewBuffer creates a new buffer in global pool.
    77  func NewBuffer() *Buffer { return globalPool.NewBuffer() }
    78  
    79  // Buffer represents the reuse buffer.
    80  type Buffer struct {
    81  	pool      *Pool
    82  	bufs      [][]byte
    83  	curBuf    []byte
    84  	curIdx    int
    85  	curBufIdx int
    86  	curBufLen int
    87  }
    88  
    89  // addBuf adds buffer to Buffer.
    90  func (b *Buffer) addBuf() {
    91  	if b.curBufIdx < len(b.bufs)-1 {
    92  		b.curBufIdx++
    93  		b.curBuf = b.bufs[b.curBufIdx]
    94  	} else {
    95  		buf := b.pool.acquire()
    96  		b.bufs = append(b.bufs, buf)
    97  		b.curBuf = buf
    98  		b.curBufIdx = len(b.bufs) - 1
    99  	}
   100  
   101  	b.curBufLen = len(b.curBuf)
   102  	b.curIdx = 0
   103  }
   104  
   105  // Reset resets the buffer.
   106  func (b *Buffer) Reset() {
   107  	if len(b.bufs) > 0 {
   108  		b.curBuf = b.bufs[0]
   109  		b.curBufLen = len(b.bufs[0])
   110  		b.curBufIdx = 0
   111  		b.curIdx = 0
   112  	}
   113  }
   114  
   115  // Destroy frees all buffers.
   116  func (b *Buffer) Destroy() {
   117  	for _, buf := range b.bufs {
   118  		b.pool.release(buf)
   119  	}
   120  	b.bufs = nil
   121  }
   122  
   123  // TotalSize represents the total memory size of this Buffer.
   124  func (b *Buffer) TotalSize() int64 {
   125  	return int64(len(b.bufs) * allocBufLen)
   126  }
   127  
   128  // AllocBytes allocates bytes with the given length.
   129  func (b *Buffer) AllocBytes(n int) []byte {
   130  	if n > bigValueSize {
   131  		return make([]byte, n)
   132  	}
   133  	if b.curIdx+n > b.curBufLen {
   134  		b.addBuf()
   135  	}
   136  	idx := b.curIdx
   137  	b.curIdx += n
   138  	return b.curBuf[idx:b.curIdx:b.curIdx]
   139  }
   140  
   141  // AddBytes adds the bytes into this Buffer.
   142  func (b *Buffer) AddBytes(bytes []byte) []byte {
   143  	buf := b.AllocBytes(len(bytes))
   144  	copy(buf, bytes)
   145  	return buf
   146  }