open-consul/agent/consul/wanfed/pool.go

256 lines
4.2 KiB
Go

package wanfed
import (
"fmt"
"net"
"sync"
"time"
)
// connPool pools idle negotiated ALPN_WANGossipPacket flavored connections to
// remote servers. Idle connections only remain pooled for up to maxTime after
// they were last acquired.
type connPool struct {
// maxTime is the maximum time to keep a connection open.
maxTime time.Duration
// mu protects pool and shutdown
mu sync.Mutex
pool map[string][]*conn
shutdown bool
shutdownCh chan struct{}
reapWg sync.WaitGroup
}
func newConnPool(maxTime time.Duration) (*connPool, error) {
if maxTime == 0 {
return nil, fmt.Errorf("wanfed: conn pool needs a max time configured")
}
p := &connPool{
maxTime: maxTime,
pool: make(map[string][]*conn),
shutdownCh: make(chan struct{}),
}
p.reapWg.Add(1)
go p.reap()
return p, nil
}
func (p *connPool) Close() error {
p.mu.Lock()
defer p.mu.Unlock()
if p.shutdown {
return nil
}
for _, conns := range p.pool {
for _, conn := range conns {
conn.Close()
}
}
p.pool = nil
p.shutdown = true
close(p.shutdownCh)
p.reapWg.Wait()
return nil
}
// AcquireOrDial either removes an idle connection from the pool or
// estabilishes a new one using the provided dialer function.
func (p *connPool) AcquireOrDial(key string, dialer func() (net.Conn, error)) (*conn, error) {
c, err := p.maybeAcquire(key)
if err != nil {
return nil, err
}
if c != nil {
c.markForUse()
return c, nil
}
nc, err := dialer()
if err != nil {
return nil, err
}
c = &conn{
key: key,
pool: p,
Conn: nc,
}
c.markForUse()
return c, nil
}
var errPoolClosed = fmt.Errorf("wanfed: connection pool is closed")
// maybeAcquire removes an idle connection from the pool if possible otherwise
// returns nil indicating there were no idle connections ready. It is the
// caller's responsibility to open a new connection if that is desired.
func (p *connPool) maybeAcquire(key string) (*conn, error) {
p.mu.Lock()
defer p.mu.Unlock()
if p.shutdown {
return nil, errPoolClosed
}
conns, ok := p.pool[key]
if !ok {
return nil, nil
}
switch len(conns) {
case 0:
delete(p.pool, key) // stray cleanup
return nil, nil
case 1:
c := conns[0]
delete(p.pool, key)
return c, nil
default:
sz := len(conns)
remaining, last := conns[0:sz-1], conns[sz-1]
p.pool[key] = remaining
return last, nil
}
}
// returnConn puts the connection back into the idle pool for reuse.
func (p *connPool) returnConn(c *conn) error {
p.mu.Lock()
defer p.mu.Unlock()
if p.shutdown {
return c.Conn.Close() // actual shutdown
}
p.pool[c.key] = append(p.pool[c.key], c)
return nil
}
// reap periodically scans the idle pool for connections that have not been
// used recently and closes them.
func (p *connPool) reap() {
defer p.reapWg.Done()
for {
select {
case <-p.shutdownCh:
return
case <-time.After(time.Second):
}
p.reapOnce()
}
}
func (p *connPool) reapOnce() {
p.mu.Lock()
defer p.mu.Unlock()
if p.shutdown {
return
}
now := time.Now()
var removedKeys []string
for key, conns := range p.pool {
if len(conns) == 0 {
removedKeys = append(removedKeys, key) // cleanup
continue
}
var retain []*conn
for _, c := range conns {
// Skip recently used connections
if now.Sub(c.lastUsed) < p.maxTime {
retain = append(retain, c)
} else {
c.Conn.Close()
}
}
if len(retain) == len(conns) {
continue // no change
} else if len(retain) == 0 {
removedKeys = append(removedKeys, key)
continue
}
p.pool[key] = retain
}
for _, key := range removedKeys {
delete(p.pool, key)
}
}
type conn struct {
key string
mu sync.Mutex
lastUsed time.Time
failed bool
closed bool
pool *connPool
net.Conn
}
func (c *conn) ReturnOrClose() error {
c.mu.Lock()
closed := c.closed
failed := c.failed
if failed {
c.closed = true
}
c.mu.Unlock()
if closed {
return nil
}
if failed {
return c.Conn.Close()
}
return c.pool.returnConn(c)
}
func (c *conn) Close() error {
c.mu.Lock()
closed := c.closed
c.closed = true
c.mu.Unlock()
if closed {
return nil
}
return c.Conn.Close()
}
func (c *conn) markForUse() {
c.mu.Lock()
c.lastUsed = time.Now()
c.failed = false
c.mu.Unlock()
}
func (c *conn) MarkFailed() {
c.mu.Lock()
c.failed = true
c.mu.Unlock()
}