go实现连接池

1.什么是连接池？

连接池是一组连接组成的一个池子(集合)，成为连接池。

2.为什么需要连接池

因为TCP的三次握手等等原因，建立连接是一件成本比较高的行为。所以在一个需要多次与特定实体交互的程序中，就需要维持一个连接池，里面有服用的链接可供重复使用。

开源实例

silenceper/pool

连接池的get原理

连接池的put原理

以下是对silenceper的本人看法，在代码注释

1.接口抽象

var (
	// ErrClosed 连接池已经关闭Error
	ErrClosed = errors.New("pool is closed")
)

// Pool 基本方法
type Pool interface {
	Get() (interface{}, error)

	Put(interface{}) error

	Close(interface{}) error

	Release()

	Len() int
}

2.实际实现

var (
	// ErrMaxActiveConnReached 连接池超限
	ErrMaxActiveConnReached = errors.New("MaxActiveConnReached")
)

// Config 连接池相关配置,用于初始化连接池配置
type Config struct {
	// 连接池的初始连接数
	InitialCap int
	// 最大并发存活连接数
	MaxCap int
	// 最大空闲连接
	MaxIdle int
	// 生成连接的方法
	Factory func() (interface{}, error)
	// 关闭连接的方法
	Close func(interface{}) error
	// 检查连接是否有效的方法
	Ping func(interface{}) error
	// 连接最大空闲时间，超过该事件则将失效
	IdleTimeout time.Duration
}

// 空闲连接
type idleConn struct {
	// 连接
	conn interface{}
	// 上次使用的时间
	t time.Time
}

// 等待连接
type connReq struct {
	idleConn *idleConn
}

// channelPool 存放连接的信息
type channelPool struct {
	mu sync.RWMutex
	// 存放空闲连接的管道
	conns chan *idleConn
	// 存放等待连接的管道切片
	connReq []chan connReq
	// 生成连接的方法
	factory func() (interface{}, error)
	// 关闭连接的方法
	close func(interface{}) error
	// 测试连接是否可用的方法
	ping func(interface{}) error
	// 空闲和等待的超时时间
	idleTimeout, waitTimeout time.Duration
	// 最大并发存活连接数
	maxActive int
	// 当前的连接数
	openingConns int
}

// NewChannelPool 初始化连接
func NewChannelPool(poolConfig Config) (Pool, error) {
	// 先判断连接池的配置信息是否符合条件
	if !(poolConfig.InitialCap <= poolConfig.MaxIdle && poolConfig.MaxIdle <= poolConfig.MaxCap && poolConfig.InitialCap >= 0) {
		return nil, errors.New("invalid capacity settings")
	}
	if poolConfig.Factory == nil {
		return nil, errors.New("invalid factory func settings")
	}
	if poolConfig.Close == nil {
		return nil, errors.New("invalid close func settings")
	}

	// 创建存放连接池信息的结构体
	c := &channelPool{
		conns:        make(chan *idleConn, poolConfig.MaxIdle),
		factory:      poolConfig.Factory,
		close:        poolConfig.Close,
		idleTimeout:  poolConfig.IdleTimeout,
		maxActive:    poolConfig.MaxCap,
		openingConns: poolConfig.InitialCap,
	}

	// 检查ping方法是否可用
	if poolConfig.Ping != nil {
		c.ping = poolConfig.Ping
	}

	// 创建初始化连接
	for i := 0; i < poolConfig.InitialCap-1; i++ {
		conn, err := c.factory()
		if err != nil {
			// 创建初始连接失败，释放所有链接
			return nil, fmt.Errorf("factory is not able to fill the pool: %s", err)
		}
		c.conns <- &idleConn{conn: conn, t: time.Now()}
	}

	return c, nil
}

// getConns 获取所有空闲连接
func (c *channelPool) getConns() chan *idleConn {
	c.mu.Lock()
	conns := c.conns
	c.mu.Unlock()
	return conns
}

// Get 从pool中取一个连接
func (c *channelPool) Get() (interface{}, error) {
	// 获取空闲连接
	conns := c.getConns()
	if conns == nil {
		return nil, ErrClosed
	}
	for {
		select {
		case wrapConn := <-conns:
			{
				// 有空闲连接，判断连接是否超时或者为空
				if wrapConn == nil {
					return nil, ErrClosed
				}
				if timeout := c.idleTimeout; timeout > 0 {
					if wrapConn.t.Add(c.idleTimeout).Before(time.Now()) {
						// 过期就关闭丢弃该连接
						c.Close(wrapConn.conn)
						continue
					}
				}

				// 如果用户编写有ping方法，调用ping方法检查一下连接是否失效，失效则丢弃，如果没编写，就不检查
				if c.ping != nil {
					if err := c.ping(wrapConn.conn); err != nil {
						c.Close(wrapConn.conn)
						continue
					}
				}
				return wrapConn.conn, nil
			}
		default:
			// 没有空闲连接
			// 上锁操作，放在并发操作问题
			c.mu.Lock()
			logrus.Debugf("openConn %v %v", c.openingConns, c.maxActive)
			// 判断当前连接数是否大于最大并发连接数，如果大于就阻塞等待别人给连接唤醒，如果小于，就新建一个连接
			if c.openingConns >= c.maxActive {
				// 大于，加入阻塞连接队列等待唤醒
				req := make(chan connReq, 1)
				c.connReq = append(c.connReq, req)
				c.mu.Unlock()
				// 等待别人给连接唤醒
				ret, ok := <-req
				if !ok {
					// 一直没人唤醒
					return nil, errors.New(fmt.Sprintf("%v no one wakes up", req))
				}
				// 等待到别人给连接被唤醒，检查连接是否超时 (其实这里也可以不检测，因为毕竟是刚从别人那里拿到的连接，应该是能用的)
				if timeout := c.idleTimeout; timeout > 0 {
					if ret.idleConn.t.Add(timeout).Before(time.Now()) {
						// 连接超时，丢弃并关闭该连接
						c.Close(ret.idleConn.conn)
						continue
					}
				}
				// 上面没有错误的情况下，直接返回拿到的连接
				return ret.idleConn.conn, nil
			}

			// 小于最大并发连接数，创建连接返回
			// 创建之前检查一下用户编写的创建连接函数是否能用 （其实这里也可以不检查,因为创建的时候已经检查过了,保险而已）
			if c.factory == nil {
				c.mu.Unlock()
				return nil, errors.New("your factory is nil")
			}
			//创建连接
			conn, err := c.factory()
			if err != nil {
				c.mu.Unlock()
				return nil, err
			}
			// 创建连接成功，当前连接数加1
			c.openingConns++
			return conn, nil
		}
	}
}

// Put 将连接放回pool中
func (c *channelPool) Put(conn interface{}) error {
	// 先判断传入的连接是否可用
	if conn == nil {
		return errors.New("connection is nil. rejecting")
	}

	//上锁，防止并发重复操作
	c.mu.Lock()

	// 检查一下连接有没有释放，释放连接的标志是c.conns == nil，因为释放的时候把channelPool所有赋值为nil
	if c.conns == nil {
		c.mu.Unlock()
		return c.Close(conn)
	}

	// 先看看有没有在阻塞队列等待连接，有就直接给连接，没有就看看空闲队列能不能放，不能放就直接关闭连接
	if connReqLen := len(c.connReq); connReqLen > 0 {
		// 有等待，把等待连接的从等待队列取出来，直接给连接
		// 从等待队列的头部取 （其实这里也可以从尾部取，这样可以保证连接不超时的概率）
		req := c.connReq[0]
		// 把取出的头部去掉, 其实这里感觉可以直接这样写 c.connReq = c.connReq[1:]
		copy(c.connReq, c.connReq[1:])
		c.connReq = c.connReq[:connReqLen-1]
		// 给连接
		req <- connReq{idleConn: &idleConn{conn: conn, t: time.Now()}}
		// 给完连接记得解锁退出
		c.mu.Unlock()
		return nil
	}

	// 等待队列没有等待，就看空闲队列是否放满了
	select {
	case c.conns <- &idleConn{conn: conn, t: time.Now()}:
		// 空闲队列没满直接放就解锁退出
		c.mu.Unlock()
		return nil
	default:
		// 空闲队列满了,关闭连接，释放锁，退出
		c.mu.Unlock()
		return c.Close(conn)
	}
}

// Release 释放所有链接
func (c *channelPool) Release() {
	c.mu.Lock()
	conns := c.conns
	c.conns = nil
	c.connReq = nil
	c.factory = nil
	c.ping = nil
	closeFunc := c.close
	c.close = nil
	c.mu.Unlock()

	if conns == nil {
		return
	}

	// 关闭空闲连接的管道
	close(conns)
	// 关闭空闲管道里面的连接
	for wrapConn := range conns {
		closeFunc(wrapConn)
	}
}

// Close 关闭连接
func (c *channelPool) Close(conn interface{}) error {
	// 判断连接是否未空
	if conn == nil {
		return errors.New("connection is nil. rejecting")
	}
	// 上锁，准备调用用户编写的close关闭连接，减少当前的连接数
	c.mu.Lock()
	defer c.mu.Unlock()

	// 判断用户的close方法是否为空
	if c.close == nil {
		return errors.New("your close method is nil,rejecting")
	}

	// 减少当前连接数
	c.openingConns--
	return c.close(conn)
}

// Ping 检查单挑连接是否有效
func (c *channelPool) Ping(conn interface{}) error {
	// 检查连接是否可用
	if conn == nil {
		return errors.New("connection is nil. rejecting")
	}

	return c.ping(conn)
}

// Len 连接池中已有的连接
func (c *channelPool) Len() int {
	return len(c.getConns())
}

// MyLen 连接池中已有的连接 (感觉这个返回当前的连接比较正确)
func (c *channelPool) MyLen() int {
	return c.openingConns
}