Repository: grpc/grpc-go
File: mem/buffer_pool.go (Go)

Code before:
type simpleBufferPool struct {
	pool sync.Pool
}

func (p *simpleBufferPool) Get(size int) *[]byte {
	bs, ok := p.pool.Get().(*[]byte)
	if ok && cap(*bs) >= size {
		clear((*bs)[:cap(*bs)])
		*bs = (*bs)[:size]
		return bs
	}

	// A buffer was pulled from the pool, but it is too small. Put it back in
	// the pool and create one large enough.
	if ok {
		p.pool.Put(bs)
	}

	// If we're going to allocate, round up to the nearest page. This way if
	// requests frequently arrive with small variation we don't allocate
	// repeatedly if we get unlucky and they increase over time. By default we
	// only allocate here if size > 1MiB.
	allocSize := size
	if allocSize%goPageSize != 0 {
		allocSize += goPageSize - (allocSize % goPageSize)
	}

	b := make([]byte, allocSize)
	b = b[:size]
	return &b
}

func (p *simpleBufferPool) Put(buf *[]byte) {
	p.pool.Put(buf)
}

var _ BufferPool = NopBufferPool{}

// NopBufferPool is a buffer pool that returns new buffers without pooling.
type NopBufferPool struct{}

// Get returns a buffer with specified length from the pool.
func (NopBufferPool) Get(length int) *[]byte {
	b := make([]byte, length)
	return &b
}

// Put returns a buffer to the pool.
func (NopBufferPool) Put(*[]byte) {
}


Review comment (performance):
As a micro-optimization, to avoid the more expensive modulo operations, we can use bitwise operations since the page size is a power of 2.
```go
allocSize = (allocSize + goPageSize - 1) & ^(goPageSize - 1)
```

Here `goPageSize - 1` is a bitmask with the trailing `N` bits as 1s, where `2^N = goPageSize`. `^(goPageSize - 1)` inverts the mask, making the leading bits 1s and the trailing `N` bits 0s. The entire expression effectively performs:

$$\text{allocSize} = \text{ceil}(\text{allocSize} / \text{goPageSize}) \times \text{goPageSize}$$

The operation `(allocSize + goPageSize - 1)` ensures that the sum is rounded up to or past the next multiple of `goPageSize`. The final bitwise AND, `& ^(goPageSize - 1)`, then zeroes out the trailing `N` bits, guaranteeing the result is an exact multiple of `goPageSize`.

Code after:
type simpleBufferPool struct {
	pool sync.Pool
}

func (p *simpleBufferPool) Get(size int) *[]byte {
	bs, ok := p.pool.Get().(*[]byte)
	if ok && cap(*bs) >= size {
		clear((*bs)[:cap(*bs)])
		*bs = (*bs)[:size]
		return bs
	}

	// A buffer was pulled from the pool, but it is too small. Put it back in
	// the pool and create one large enough.
	if ok {
		p.pool.Put(bs)
	}

	// If we're going to allocate, round up to the nearest page. This way if
	// requests frequently arrive with small variation we don't allocate
	// repeatedly if we get unlucky and they increase over time. By default we
	// only allocate here if size > 1MiB. Because goPageSize is a power of 2, we
	// can round up efficiently.
	allocSize := (size + goPageSize - 1) & ^(goPageSize - 1)

	b := make([]byte, size, allocSize)
	return &b
}

func (p *simpleBufferPool) Put(buf *[]byte) {
	p.pool.Put(buf)
}

var _ BufferPool = NopBufferPool{}

// NopBufferPool is a buffer pool that returns new buffers without pooling.
type NopBufferPool struct{}

// Get returns a buffer with specified length from the pool.
func (NopBufferPool) Get(length int) *[]byte {
	b := make([]byte, length)
	return &b
}

// Put returns a buffer to the pool.
func (NopBufferPool) Put(*[]byte) {
}