open-consul/consul/pool.go

268 lines
5.7 KiB
Go
Raw Normal View History

2013-12-09 20:09:57 +00:00
package consul
import (
"crypto/tls"
2013-12-24 20:22:42 +00:00
"fmt"
"github.com/inconshreveable/muxado"
2013-12-09 22:52:22 +00:00
"github.com/ugorji/go/codec"
2013-12-09 20:09:57 +00:00
"net"
2013-12-09 22:52:22 +00:00
"net/rpc"
2013-12-09 20:09:57 +00:00
"sync"
"sync/atomic"
2013-12-09 20:09:57 +00:00
"time"
)
// Conn is a pooled connection to a Consul server
type Conn struct {
refCount int32
2013-12-19 23:42:17 +00:00
addr net.Addr
session muxado.Session
2013-12-19 23:42:17 +00:00
lastUsed time.Time
2013-12-09 20:09:57 +00:00
}
2013-12-09 22:58:49 +00:00
func (c *Conn) Close() error {
return c.session.Close()
2013-12-09 22:58:49 +00:00
}
2013-12-09 20:09:57 +00:00
// ConnPool is used to maintain a connection pool to other
// Consul servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcConsul mode. Raft connections
// are pooled seperately.
2013-12-09 20:09:57 +00:00
type ConnPool struct {
sync.Mutex
2013-12-19 23:42:17 +00:00
// The maximum time to keep a connection open
maxTime time.Duration
// Pool maps an address to a open connection
pool map[string]*Conn
2013-12-09 20:09:57 +00:00
// TLS settings
tlsConfig *tls.Config
2013-12-09 20:09:57 +00:00
// Used to indicate the pool is shutdown
shutdown bool
shutdownCh chan struct{}
2013-12-09 20:09:57 +00:00
}
// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. If TLS settings are provided
// outgoing connections use TLS.
func NewPool(maxTime time.Duration, tlsConfig *tls.Config) *ConnPool {
2013-12-09 20:09:57 +00:00
pool := &ConnPool{
maxTime: maxTime,
pool: make(map[string]*Conn),
tlsConfig: tlsConfig,
shutdownCh: make(chan struct{}),
2013-12-09 20:09:57 +00:00
}
2013-12-19 23:42:17 +00:00
if maxTime > 0 {
go pool.reap()
}
2013-12-09 20:09:57 +00:00
return pool
}
// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
p.Lock()
defer p.Unlock()
for _, conn := range p.pool {
conn.Close()
2013-12-09 20:09:57 +00:00
}
p.pool = make(map[string]*Conn)
2013-12-09 20:09:57 +00:00
if p.shutdown {
return nil
}
p.shutdown = true
close(p.shutdownCh)
2013-12-09 20:09:57 +00:00
return nil
}
// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(addr net.Addr) (*Conn, error) {
2013-12-09 20:09:57 +00:00
// Check for a pooled ocnn
if conn := p.getPooled(addr); conn != nil {
return conn, nil
}
// Create a new connection
return p.getNewConn(addr)
}
// getPooled is used to return a pooled connection
func (p *ConnPool) getPooled(addr net.Addr) *Conn {
p.Lock()
defer p.Unlock()
// Look for an existing connection
c := p.pool[addr.String()]
if c != nil {
c.lastUsed = time.Now()
atomic.AddInt32(&c.refCount, 1)
2013-12-09 20:09:57 +00:00
}
return c
2013-12-09 20:09:57 +00:00
}
// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(addr net.Addr) (*Conn, error) {
// Try to dial the conn
conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
2013-12-09 20:09:57 +00:00
if err != nil {
return nil, err
}
// Cast to TCPConn
if tcp, ok := conn.(*net.TCPConn); ok {
tcp.SetKeepAlive(true)
tcp.SetNoDelay(true)
}
// Check if TLS is enabled
if p.tlsConfig != nil {
// Switch the connection into TLS mode
if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
conn.Close()
return nil, err
}
2013-12-09 20:09:57 +00:00
// Wrap the connection in a TLS client
conn = tls.Client(conn, p.tlsConfig)
}
2013-12-09 20:09:57 +00:00
// Write the Consul multiplex byte to set the mode
if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
conn.Close()
return nil, err
}
// Create a multiplexed session
session := muxado.Client(conn)
2013-12-09 22:52:22 +00:00
2013-12-09 20:09:57 +00:00
// Wrap the connection
c := &Conn{
refCount: 1,
addr: addr,
session: session,
lastUsed: time.Now(),
2013-12-09 20:09:57 +00:00
}
// Monitor the session
go func() {
session.Wait()
p.Lock()
defer p.Unlock()
if conn, ok := p.pool[addr.String()]; ok && conn.session == session {
delete(p.pool, addr.String())
}
}()
// Track this connection, handle potential race condition
2013-12-09 20:09:57 +00:00
p.Lock()
defer p.Unlock()
if existing := p.pool[addr.String()]; existing != nil {
session.Close()
return existing, nil
} else {
p.pool[addr.String()] = c
return c, nil
2013-12-09 20:09:57 +00:00
}
}
// clearConn is used to clear any cached connection, potentially in response to an erro
func (p *ConnPool) clearConn(addr net.Addr) {
p.Lock()
defer p.Unlock()
delete(p.pool, addr.String())
}
2013-12-09 20:09:57 +00:00
// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
atomic.AddInt32(&conn.refCount, -1)
2013-12-09 20:09:57 +00:00
}
2013-12-09 22:58:49 +00:00
// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(addr net.Addr, method string, args interface{}, reply interface{}) error {
retries := 0
START:
2013-12-09 22:58:49 +00:00
// Try to get a conn first
conn, err := p.acquire(addr)
2013-12-09 22:58:49 +00:00
if err != nil {
2013-12-24 20:22:42 +00:00
return fmt.Errorf("failed to get conn: %v", err)
2013-12-09 22:58:49 +00:00
}
defer p.releaseConn(conn)
// Create a new stream
stream, err := conn.session.Open()
if err != nil {
p.clearConn(addr)
// Try to redial, possible that the TCP session closed due to timeout
if retries == 0 {
retries++
goto START
}
return fmt.Errorf("failed to start stream: %v", err)
}
defer stream.Close()
// Create the RPC client
cc := codec.GoRpc.ClientCodec(stream, &codec.MsgpackHandle{})
client := rpc.NewClientWithCodec(cc)
2013-12-09 22:58:49 +00:00
// Make the RPC call
err = client.Call(method, args, reply)
2013-12-09 22:58:49 +00:00
// Fast path the non-error case
if err == nil {
return nil
}
// If its a network error, nuke the connection
if _, ok := err.(net.Error); ok {
p.clearConn(addr)
2013-12-09 22:58:49 +00:00
}
2013-12-24 20:22:42 +00:00
return fmt.Errorf("rpc error: %v", err)
2013-12-09 22:58:49 +00:00
}
2013-12-19 23:42:17 +00:00
// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
for !p.shutdown {
// Sleep for a while
select {
case <-time.After(time.Second):
case <-p.shutdownCh:
return
}
2013-12-19 23:42:17 +00:00
// Reap all old conns
p.Lock()
var removed []string
2013-12-19 23:42:17 +00:00
now := time.Now()
for host, conn := range p.pool {
// Skip recently used connections
if now.Sub(conn.lastUsed) < p.maxTime {
continue
2013-12-19 23:42:17 +00:00
}
// Skip connections with active streams
if atomic.LoadInt32(&conn.refCount) > 0 {
continue
}
// Close the conn
conn.Close()
// Remove from pool
removed = append(removed, host)
}
for _, host := range removed {
delete(p.pool, host)
2013-12-19 23:42:17 +00:00
}
p.Unlock()
}
}