Golang Redis连接池实现原理及示例探究

type Pool struct {
	Config
	// 创建对象
	newObject func() (any, error)
	// 释放对象
	freeObject func(x any)
	// 空闲对象池
	idles chan any
	mu          sync.Mutex
	activeCount int// 已经创建的对象个数
	waiting     []chan any // 阻塞等待
	closed bool// 是否已关闭
}
func NewPool(new func() (any, error), free func(x any), conf Config) *Pool {
	ifnew == nil {
		logger.Error("NewPool argument new func is nil")
		returnnil
	}
	if free == nil {
		free = func(x any) {}
	}
	p := Pool{
		Config:      conf,
		newObject:   new,
		freeObject:  free,
		activeCount: 0,
		closed:      false,
	}
	p.idles = make(chan any, p.MaxIdles)
	return &p
}

func (p *Pool) Get() (any, error) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		returnnil, ErrClosed
	}
	select {
	case x := <-p.idles: // 从空闲中获取
		p.mu.Unlock() // 解锁
		return x, nil
	default:
		return p.getOne() // 获取一个新的
	}
}
func (p *Pool) getOne() (any, error) {
	// 说明已经创建了太多对象
	if p.activeCount >= p.Config.MaxActive {
		wait := make(chan any, 1)
		p.waiting = append(p.waiting, wait)
		p.mu.Unlock()
		// 阻塞等待
		x, ok := <-wait
		if !ok {
			returnnil, ErrClosed
		}
		return x, nil
	}
	p.activeCount++
	p.mu.Unlock()
	// 创建新对象
	x, err := p.newObject()
	if err != nil {
		p.mu.Lock()
		p.activeCount--
		p.mu.Unlock()
		returnnil, err
	}
	return x, nil
}

func (p *Pool) Put(x any) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		p.freeObject(x) // 直接释放
		return
	}

	//1.先判断等待中
	iflen(p.waiting) > 0 {
		// 弹出一个（从头部）
		wait := p.waiting[0]
		temp := make([]chan any, len(p.waiting)-1)
		copy(temp, p.waiting[1:])
		p.waiting = temp
		wait <- x // 取消阻塞
		p.mu.Unlock()
		return

	}
	// 2.直接放回空闲缓冲
	select {
	case p.idles <- x:
		p.mu.Unlock()
	default: // 说明空闲已满
		p.activeCount-- // 对象个数-1
		p.mu.Unlock()
		p.freeObject(x) // 释放
	}

}

type RedisConnPool struct {
	connDict *dict.ConcurrentDict // addr -> *pool.Pool
}
func NewRedisConnPool() *RedisConnPool {
	return &RedisConnPool{
		connDict: dict.NewConcurrentDict(16),
	}
}
func (r *RedisConnPool) GetConn(addr string) (*client.RedisClent, error) {
	var connectionPool *pool.Pool // 对象池
	// 通过不同的地址addr，获取不同的对象池
	raw, ok := r.connDict.Get(addr)
	if ok {
		connectionPool = raw.(*pool.Pool)
	} else {
		// 创建对象函数
		newClient := func() (any, error) {
			// redis的客户端连接
			cli, err := client.NewRedisClient(addr)
			if err != nil {
				returnnil, err
			}
			// 启动
			cli.Start()
			if conf.GlobalConfig.RequirePass != "" { // 说明服务需要密码
				reply, err := cli.Send(aof.Auth([]byte(conf.GlobalConfig.RequirePass)))
				if err != nil {
					returnnil, err
				}
				if !protocol.IsOKReply(reply) {
					returnnil, errors.New("auth failed:" + string(reply.ToBytes()))
				}
				return cli, nil
			}
			return cli, nil
		}
		// 释放对象函数
		freeClient := func(x any) {
			cli, ok := x.(*client.RedisClent)
			if ok {
				cli.Stop() // 释放
			}
		}
		// 针对addr地址，创建一个新的对象池
		connectionPool = pool.NewPool(newClient, freeClient, pool.Config{
			MaxIdles:  1,
			MaxActive: 20,
		})
		// addr -> *pool.Pool
		r.connDict.Put(addr, connectionPool)
	}
	// 从对象池中获取一个对象
	raw, err := connectionPool.Get()
	if err != nil {
		returnnil, err
	}
	conn, ok := raw.(*client.RedisClent)
	if !ok {
		returnnil, errors.New("connection pool make wrong type")
	}
	return conn, nil
}
func (r *RedisConnPool) ReturnConn(peer string, cli *client.RedisClent) error {
	raw, ok := r.connDict.Get(peer)
	if !ok {
		return errors.New("connection pool not found")
	}
	raw.(*pool.Pool).Put(cli)
	returnnil
}