go-redis 源码分析：连接池
笔者最近在项目中基于 go-redis 实现 Redis 缓存优化性能。go-redis 是一个 Go 语言实现的 Redis 客户端，既然是网络服务的客户端，为了高效利用有限资源，避免重复创建和销毁网络连接，就必需对其进行管理。而资源管理又是编程领域中的一个重点难点，抱着对是否能利用 Go 语言语法简洁的特点来优雅实现连接池的好奇，笔者决定阅读并分析 go-redis 连接池部分的源码，一探究竟。以下是对源码的分析，分为接口与结构体、连接池管理、建立与关闭连接、获取与放回连接、监控统计等5大部分，源码链接。
接口与结构体
连接结构体：
type Conn struct {
netConn net.Conn // 基于 tcp 的网络连接
rd *proto.Reader // 根据 Redis 通信协议实现的 Reader
wr *proto.Writer // 根据 Redis 通信协议实现的 Writer
Inited bool // 是否完成初始化
pooled bool // 是否放进连接池
createdAt time.Time // 创建时间
usedAt int64 // 使用时间，atomic
}
连接池接口：
type Pooler interface {
NewConn(context.Context) (*Conn, error) // 创建连接
CloseConn(*Conn) error // 关闭连接
Get(context.Context) (*Conn, error) // 获取连接
Put(*Conn) // 放回连接
Remove(*Conn, error) // 移除连接
Len() int // 连接池长度
IdleLen() int // 空闲连接数量
Stats() *Stats // 连接池统计
Close() error // 关闭连接池
}
连接池结构体：
type ConnPool struct {
opt *Options // 连接池配置
dialErrorsNum uint32 // 连接错误次数，atomic
lastDialErrorMu sync.RWMutex // 上一次连接错误锁，读写锁
lastDialError error // 上一次连接错误
queue chan struct{} // 工作连接队列
connsMu sync.Mutex // 连接队列锁
conns []*Conn // 连接队列
idleConns []*Conn // 空闲连接队列
poolSize int // 连接池大小
idleConnsLen int // 空闲连接队列长度
stats Stats // 连接池统计
_closed uint32 // 连接池关闭标志，atomic
closedCh chan struct{} // 通知连接池关闭通道
}
连接池管理
初始化
var _ Pooler = (*ConnPool)(nil)
func NewConnPool(opt *Options) *ConnPool {
p := &ConnPool{
opt: opt,
queue: make(chan struct{}, opt.PoolSize),
conns: make([]*Conn, 0, opt.PoolSize),
idleConns: make([]*Conn, 0, opt.PoolSize),
closedCh: make(chan struct{}),
}
p.checkMinIdleConns()
if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
go p.reaper(opt.IdleCheckFrequency)
}
return p
}
- 创建连接池，传入连接池配置选项参数 opt，工厂函数根据 opt 创建连接池实例。连接池主要依靠以下四个数据结构实现管理和通信：
- queue： 存储工作连接的缓冲通道
- conns：存储所有连接的切片
- idleConns：存储空闲连接的切片
- closed：用于通知所有协程连接池已经关闭的通道
- 检查连接池的空闲连接数量是否满足最小空闲连接数量要求，若不满足，则创建足够的空闲连接。
- 若连接池配置选项规定了空闲连接超时和检查空闲连接频率，则开启一个清理空闲连接的协程。
关闭
func (p *ConnPool) Close() error {
if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
return ErrClosed
}
close(p.closedCh)
var firstErr error
p.connsMu.Lock()
for _, cn := range p.conns {
if err := p.closeConn(cn); err != nil && firstErr == nil {
firstErr = err
}
}
p.conns = nil
p.poolSize = 0
p.idleConns = nil
p.idleConnsLen = 0
p.connsMu.Unlock()
return firstErr
}
- 原子性检查连接池是否已经关闭，若没关闭，则将关闭标志置为1
- 关闭 closedCh 通道，连接池中的所有协程都可以通过判断该通道是否关闭来确定连接池是否已经关闭。
- 连接队列锁上锁，关闭队列中的所有连接，并置空所有维护连接池状态的数据结构，解锁。
过滤
func (p *ConnPool) Filter(fn func(*Conn) bool) error {
var firstErr error
p.connsMu.Lock()
for _, cn := range p.conns {
if fn(cn) {
if err := p.closeConn(cn); err != nil && firstErr == nil {
firstErr = err
}
}
}
p.connsMu.Unlock()
return firstErr
}
实质上是遍历连接池中的所有连接，并调用传入的 fn 过滤函数作用在每个连接上，过滤出符合业务要求的连接。
清理
func (p *ConnPool) reaper(frequency time.Duration) {
ticker := time.NewTicker(frequency)
defer ticker.Stop()
for {
select {
case <-ticker.C:
// It is possible that ticker and closedCh arrive together,
// and select pseudo-randomly pick ticker case, we double
// check here to prevent being executed after closed.
if p.closed() {
return
}
_, err := p.ReapStaleConns()
if err != nil {
internal.Logger.Printf("ReapStaleConns failed: %s", err)
continue
}
case <-p.closedCh:
return
}
}
}
func (p *ConnPool) ReapStaleConns() (int, error) {
var n int
for {
p.getTurn()
p.connsMu.Lock()
cn := p.reapStaleConn()
p.connsMu.Unlock()
p.freeTurn()
if cn != nil {
_ = p.closeConn(cn)
n++
} else {
break
}
}
atomic.AddUint32(&p.stats.StaleConns, uint32(n))
return n, nil
}
func (p *ConnPool) reapStaleConn() *Conn {
if len(p.idleConns) == 0 {
return nil
}
cn := p.idleConns[0]
if !p.isStaleConn(cn) {
return nil
}
p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
p.idleConnsLen--
p.removeConn(cn)
return cn
}
- 开启一个用于检查并清理过期连接的 goroutine 每隔 frequency 时间遍历检查连接池中是否存在过期连接，并清理。
- 创建一个时间间隔为 frequency 的计时器，在连接池关闭时关闭该计时器
- 循环判断计时器是否到时和连接池是否关闭
- 移除空闲连接队列中的过期连接
建立与关闭连接
建立连接
func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
cn, err := p.dialConn(ctx, pooled)
if err != nil {
return nil, err
}
p.connsMu.Lock()
p.conns = append(p.conns, cn)
if pooled {
// If pool is full remove the cn on next Put.
if p.poolSize >= p.opt.PoolSize {
cn.pooled = false
} else {
p.poolSize++
}
}
p.connsMu.Unlock()
return cn, nil
}
func (p *ConnPool) dialConn(ctx context.Context, pooled bool) (*Conn, error) {
if p.closed() {
return nil, ErrClosed
}
if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
return nil, p.getLastDialError()
}
netConn, err := p.opt.Dialer(ctx)
if err != nil {
p.setLastDialError(err)
if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
go p.tryDial()
}
return nil, err
}
cn := NewConn(netConn)
cn.pooled = pooled
return cn, nil
}
func (p *ConnPool) tryDial() {
for {
if p.closed() {
return
}
conn, err := p.opt.Dialer(context.Background())
if err != nil {
p.setLastDialError(err)
time.Sleep(time.Second)
continue
}
atomic.StoreUint32(&p.dialErrorsNum, 0)
_ = conn.Close()
return
}
}
创建连接流程图：
newConn流程图.png
DialConn流程图.png
移除与关闭连接
func (p *ConnPool) Remove(cn *Conn, reason error) {
p.removeConnWithLock(cn)
p.freeTurn()
_ = p.closeConn(cn)
}
func (p *ConnPool) CloseConn(cn *Conn) error {
p.removeConnWithLock(cn)
return p.closeConn(cn)
}
func (p *ConnPool) removeConnWithLock(cn *Conn) {
p.connsMu.Lock()
p.removeConn(cn)
p.connsMu.Unlock()
}
func (p *ConnPool) removeConn(cn *Conn) {
for i, c := range p.conns {
if c == cn {
p.conns = append(p.conns[:i], p.conns[i+1:]...)
if cn.pooled {
p.poolSize--
p.checkMinIdleConns()
}
return
}
}
}
func (p *ConnPool) closeConn(cn *Conn) error {
if p.opt.OnClose != nil {
_ = p.opt.OnClose(cn)
}
return cn.Close()
}
连接池无论移除还是关闭连接，底层调用的都是 removeConnWithLock 函数。removeConnWithLock 函数的工作流程如下：
- 连接队列上锁
- 遍历连接队列找到要关闭的连接，并将其移除出连接队列
- 更新连接池统计数据
- 检查连接池最小空闲连接数量
- 连接队列解锁
- 关闭连接，先执行关闭连接时的回调函数（创建连接池时的配置选项传入），再关闭连接
获取与放回连接
获取
// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
if p.closed() {
return nil, ErrClosed
}
err := p.waitTurn(ctx)
if err != nil {
return nil, err
}
for {
p.connsMu.Lock()
cn := p.popIdle()
p.connsMu.Unlock()
if cn == nil {
break
}
if p.isStaleConn(cn) {
_ = p.CloseConn(cn)
continue
}
atomic.AddUint32(&p.stats.Hits, 1)
return cn, nil
}
atomic.AddUint32(&p.stats.Misses, 1)
newcn, err := p.newConn(ctx, true)
if err != nil {
p.freeTurn()
return nil, err
}
return newcn, nil
}
func (p *ConnPool) getTurn() {
p.queue <- struct{}{}
}
func (p *ConnPool) waitTurn(ctx context.Context) error {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
select {
case p.queue <- struct{}{}:
return nil
default:
}
timer := timers.Get().(*time.Timer)
timer.Reset(p.opt.PoolTimeout)
select {
case <-ctx.Done():
if !timer.Stop() {
<-timer.C
}
timers.Put(timer)
return ctx.Err()
case p.queue <- struct{}{}:
if !timer.Stop() {
<-timer.C