redis/internal/pool/pool_sticky.go

124 lines
1.8 KiB
Go

package pool
import "sync"
type StickyConnPool struct {
pool *ConnPool
reusable bool
cn *Conn
closed bool
mu sync.Mutex
}
var _ Pooler = (*StickyConnPool)(nil)
func NewStickyConnPool(pool *ConnPool, reusable bool) *StickyConnPool {
return &StickyConnPool{
pool: pool,
reusable: reusable,
}
}
func (p *StickyConnPool) NewConn() (*Conn, error) {
panic("not implemented")
}
func (p *StickyConnPool) CloseConn(*Conn) error {
panic("not implemented")
}
func (p *StickyConnPool) Get() (*Conn, bool, error) {
p.mu.Lock()
defer p.mu.Unlock()
if p.closed {
return nil, false, ErrClosed
}
if p.cn != nil {
return p.cn, false, nil
}
cn, _, err := p.pool.Get()
if err != nil {
return nil, false, err
}
p.cn = cn
return cn, true, nil
}
func (p *StickyConnPool) putUpstream() (err error) {
err = p.pool.Put(p.cn)
p.cn = nil
return err
}
func (p *StickyConnPool) Put(cn *Conn) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.closed {
return ErrClosed
}
return nil
}
func (p *StickyConnPool) removeUpstream() error {
err := p.pool.Remove(p.cn)
p.cn = nil
return err
}
func (p *StickyConnPool) Remove(cn *Conn) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.closed {
return nil
}
return p.removeUpstream()
}
func (p *StickyConnPool) Len() int {
p.mu.Lock()
defer p.mu.Unlock()
if p.cn == nil {
return 0
}
return 1
}
func (p *StickyConnPool) FreeLen() int {
p.mu.Lock()
defer p.mu.Unlock()
if p.cn == nil {
return 1
}
return 0
}
func (p *StickyConnPool) Stats() *Stats {
return nil
}
func (p *StickyConnPool) Close() error {
p.mu.Lock()
defer p.mu.Unlock()
if p.closed {
return ErrClosed
}
p.closed = true
var err error
if p.cn != nil {
if p.reusable {
err = p.putUpstream()
} else {
err = p.removeUpstream()
}
}
return err
}