142 lines
3 KiB
Go
142 lines
3 KiB
Go
package consul
|
|
|
|
import (
|
|
"github.com/ugorji/go/codec"
|
|
"net"
|
|
"net/rpc"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Conn is a pooled connection to a Consul server
|
|
type Conn struct {
|
|
addr net.Addr
|
|
conn *net.TCPConn
|
|
client *rpc.Client
|
|
}
|
|
|
|
// 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.
|
|
type ConnPool struct {
|
|
sync.Mutex
|
|
|
|
// The maximum connectsion to maintain per server
|
|
maxConns int
|
|
|
|
// Pool maps an address to a list of connections
|
|
pool map[string][]*Conn
|
|
|
|
// Used to indicate the pool is shutdown
|
|
shutdown bool
|
|
}
|
|
|
|
// NewPool is used to make a new connection pool
|
|
// Maintain at most maxConns per host
|
|
func NewPool(maxConns int) *ConnPool {
|
|
pool := &ConnPool{
|
|
maxConns: maxConns,
|
|
pool: make(map[string][]*Conn),
|
|
}
|
|
return pool
|
|
}
|
|
|
|
// Shutdown is used to close the connection pool
|
|
func (p *ConnPool) Shutdown() error {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
|
|
for _, conns := range p.pool {
|
|
for _, c := range conns {
|
|
c.conn.Close()
|
|
}
|
|
}
|
|
p.pool = make(map[string][]*Conn)
|
|
p.shutdown = true
|
|
|
|
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) {
|
|
// 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
|
|
conns := p.pool[addr.String()]
|
|
if len(conns) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Remove the last conn from the pool
|
|
conn := conns[len(conns)-1]
|
|
conns = conns[:len(conns)-1]
|
|
p.pool[addr.String()] = conns
|
|
return conn
|
|
}
|
|
|
|
// getNewConn is used to return a new connection
|
|
func (p *ConnPool) getNewConn(addr net.Addr) (*Conn, error) {
|
|
// Try to dial the conn
|
|
rawConn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Cast to TCPConn
|
|
conn := rawConn.(*net.TCPConn)
|
|
|
|
// Enable keep alives
|
|
conn.SetKeepAlive(true)
|
|
conn.SetNoDelay(true)
|
|
|
|
// Write the Consul RPC byte to set the mode
|
|
conn.Write([]byte{byte(rpcConsul)})
|
|
|
|
// Create the RPC client
|
|
cc := codec.GoRpc.ClientCodec(conn, &codec.MsgpackHandle{})
|
|
client := rpc.NewClientWithCodec(cc)
|
|
|
|
// Wrap the connection
|
|
c := &Conn{
|
|
addr: addr,
|
|
conn: conn,
|
|
client: client,
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
// Return is used to return a connection once done. Connections
|
|
// that are in an error state should not be returned
|
|
func (p *ConnPool) Return(conn *Conn) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
|
|
// Look for existing connections
|
|
conns := p.pool[conn.addr.String()]
|
|
|
|
// Check for limit on connections or shutdown
|
|
if p.shutdown || len(conns) >= p.maxConns {
|
|
conn.conn.Close()
|
|
return
|
|
}
|
|
|
|
// Retain the connection
|
|
conns = append(conns, conn)
|
|
p.pool[conn.addr.String()] = conns
|
|
}
|