Go-如何优雅的使用字节池

背景

在某些场景下,我们可能会大量的使用字节数组,比如IO操作、编解码,如果不进行优化,大量的申请和释放字节数组会造成一定的性能损耗,因此有必要复用字节数组。

最简单的方式

对于Go语言来说,我们第一个想到的就是使用sync.Pool来做字节数组的对象池,比如这样:

package bufferpool

import "sync"

type BytePool struct {
	p sync.Pool
}

func NewBytePool(size, cap int) *BytePool {
	if size > cap {
		panic("size must be less then cap")
	}
	p := &BytePool{}
	p.p.New = func() any {
		return make([]byte, size, cap)
	}
	return p
}

// 获取字节数组
func (p *BytePool) Get() []byte {
	return p.p.Get().([]byte)
}

// 归还字节数组
func (p *BytePool) Put(b []byte) {
	// 重置已用大小
	b = b[:0]
	p.p.Put(b)
}

我们简单的封装了sync.Poolsync.Pool.New根据指定的初始大小申请新的字节数组,在Put的时候重置字节数组的已用空间(这样下次才能从头开始使用)。

测试

我们进行一个简单性能测试,也就是不断的申请字节数组,然后写入长度为1024的字节数组块,共64块,也就是64KB,测试样例共3个:

不预先申请空间

这个样例我们不预先申请字节数组空间,因此在append的过程中会不断的申请新的更大的空间,然后转移字节数组内容。

func BenchmarkByte(b *testing.B) {
	for n := 0; n < b.N; n++ {
                // 从长度为0的字节数组开始
		var b []byte
		for i := 0; i < blocks; i++ {
			b = append(b, block...)
		}
	}
}

预先申请空间

由于这个测试的总大小的预先知道的,因此我们可以先提前申请空间,这样就不用在append过程中不断的申请新的更大空间,然后转移字节数组内容了。

func BenchmarkMake(b *testing.B) {
	for n := 0; n < b.N; n++ {
                // 预先保留需要的空间
		b := make([]byte, 0, blocks*blockSize)
		for i := 0; i < blocks; i++ {
			b = append(b, block...)
		}
	}
}

字节数组池

这里我们每次先从字节池拿一个字节数组Get(),使用完之后归还字节池Put()

func BenchmarkBytePool(b *testing.B) {
	pool := NewBytePool(0, blocks*blockSize)
	for n := 0; n < b.N; n++ {
                // 拿字节数组
		b := pool.Get()
		for i := 0; i < blocks; i++ {
			b = append(b, block...)
		}
                // 归还
		pool.Put(b)
	}
}

测试结果

可以看到我们简单的字节池就可以带来很大的性能提升!

BenchmarkByte-16                   32470             38136 ns/op
BenchmarkMake-16                  605449              1962 ns/op
BenchmarkBytePool-16             1000000              1162 ns/op

更优雅的方式

在实际的编程中,我们在使用字节数组时,很多时候都需要以一个的形式去读写,同时也可能很难提前计算出需要的大小,因此bytes.Buffer可能更加适合实际的编程。

package bufferpool

import (
	"bytes"
	"sync"
)

type BufferPool struct {
	p sync.Pool
}

func NewBufferPool(size, cap int) *BufferPool {
	if size > cap {
		panic("size must be less then cap")
	}
	p := &BufferPool{}
	p.p.New = func() any {
		var b []byte
		if cap > 0 {
			b = make([]byte, size, cap)
		}
		return bytes.NewBuffer(b)
	}
	return p
}

// 获取字节数组
func (p *BufferPool) Get() *bytes.Buffer {
	return p.p.Get().(*bytes.Buffer)
}

// 归还字节数组
func (p *BufferPool) Put(b *bytes.Buffer) {
	// 重置已用大小
	b.Reset()
	p.p.Put(b)
}

测试

测试条件与上面相同。

直接使用Buffer

作为对比实验我们直接使用Buffer。

func BenchmarkBuffer(b *testing.B) {
	for n := 0; n < b.N; n++ {
		b := bytes.NewBuffer(make([]byte, 0, blocks*blockSize))
		for i := 0; i < blocks; i++ {
			b.Write(block)
		}
	}
}

bytes.Buffer池

func BenchmarkBufferPool(b *testing.B) {
	pool := NewBufferPool(0, blocks*blockSize)
	for n := 0; n < b.N; n++ {
		b := pool.Get()
		for i := 0; i < blocks; i++ {
			b.Write(block)
		}
		pool.Put(b)
	}
}

测试结果

可以看到使用bytes.Buffer池字节数组池性能差了一点,主要是因为bytes.Buffer比较复杂,但是bytes.Buffer的功能比字节数组强大很多。

BenchmarkByte-16                   31748             38131 ns/op
BenchmarkMake-16                  605847              1964 ns/op
BenchmarkBytePool-16             1000000              1162 ns/op
BenchmarkBuffer-16                589336              2030 ns/op
BenchmarkBufferPool-16            962132              1235 ns/op

限制池大小

有时候我们不想对象池无限大,因此我们需要限制对象池的大小,对于Go语言来说,我们可以使用channel+select,也就是申请一个固定长度缓冲区的channel,配合select的default分支。

  • Put:channel不满则put,否则default分支丢弃这个对象。
  • Get:channel不空则get,否则default分支申请新对象。

这里我们直接使用minio的实现: github.com/minio/minio…

package bufferpool

type ByteFixPool struct {
	cache chan []byte
	size  int
	cap   int
}

// cacheSize: 字节池缓存长度
// size: 字节数组长度
// cap: 字节数组容量
func NewByteFixPool(cacheSize, size, cap int) *ByteFixPool {
	if size > cap {
		panic("size must be less then cap")
	}
	return &ByteFixPool{
		cache: make(chan []byte, cacheSize),
		size:  size,
		cap:   cap,
	}
}

func (p *ByteFixPool) Get() []byte {
	select {
	// 从channel读
	case b := <-p.cache:
		return b
		// 如果channel空则申请一个新的字节数组
	default:
		return make([]byte, p.size, p.cap)
	}
}

func (p *ByteFixPool) Put(b []byte) {
	// 重置已用大小
	b = b[:0]
	select {
	// 放入channel
	case p.cache <- b:
	// channel满了则丢弃字节数组
	default:
	}
}


测试

固定大小字节池

这里使用固定大小字节池,同时预先分配空间。

func BenchmarkByteFixPool(b *testing.B) {
	pool := NewByteFixPool(16, 0, blocks*blockSize)
	for n := 0; n < b.N; n++ {
		b := pool.Get()
		for i := 0; i < blocks; i++ {
			b = append(b, block...)
		}
		pool.Put(b)
	}
}

测试结果

可以看到使用channel+select的性能甚至更好一点,而且还能限制字节池大小,当然相比于sync.Pool的实现,它在字节池channel里面的空间是没办法自动回收的。

BenchmarkByte-16                   31748             38131 ns/op
BenchmarkMake-16                  605847              1964 ns/op
BenchmarkBytePool-16             1000000              1162 ns/op
BenchmarkBuffer-16                589336              2030 ns/op
BenchmarkBufferPool-16            962132              1235 ns/op
BenchmarkByteFixPool-16          1000000              1130 ns/op

总结

对于字节池来说。

字节对象可以是:

  • []byte:字节数组
  • bytes.Buffer:功能更加强大的字节数组
  • 其他:比如一组bytes.Buffer

实现方式可以是:

  • sync.Pool:根据GC期间对象是否使用回收对象
  • channel+select:限制字节池长度
  • 其他:比如限制对象池使用空间

当然,最通用的实现是sync.Pool+bytes.Buffer,因为sync.Pool能够自动回收字节对象,bytes.Buffer又能提供强大的功能。

上面介绍的几种都是比较常用的,而且实现也非常简单的字节池,如果在业务中有更加复杂的需求,也可以根据需求实现一个字节池。