Add Connection Pool for GetNodeAPIAddr
@@ -5,16 +5,27 @@ import (
"errors"
"fmt"
"io/ioutil"
+ "net"
+ "sync"
"time"
"github.com/golang/protobuf/proto"
"github.com/rqlite/rqlite/command"
+ "github.com/rqlite/rqlite/tcp/pool"
+)
+
+const (
+ initialPoolSize = 4
+ maxPoolCapacity = 64
)
// Client allows communicating with a remote node.
type Client struct {
dialer Dialer
timeout time.Duration
+
+ mu sync.RWMutex
+ pools map[string]pool.Pool
}
// NewClient returns a client instance for talking to a remote node.
@@ -22,17 +33,56 @@ func NewClient(dl Dialer) *Client {
return &Client{
dialer: dl,
timeout: 30 * time.Second,
+ pools: make(map[string]pool.Pool),
}
}
// GetNodeAPIAddr retrieves the API Address for the node at nodeAddr
func (c *Client) GetNodeAPIAddr(nodeAddr string) (string, error) {
- conn, err := c.dialer.Dial(nodeAddr, c.timeout)
+ var pl pool.Pool
+ var ok bool
+
+ c.mu.RLock()
+ pl, ok = c.pools[nodeAddr]
+ c.mu.RUnlock()
+
+ // Do we need a new pool for the given address?
+ if !ok {
+ if err := func() error {
+ c.mu.Lock()
+ defer c.mu.Unlock()
+ pl, ok = c.pools[nodeAddr]
+ if ok {
+ return nil // Pool was inserted just after we checked.
+ }
+
+ // New pool is needed for given address.
+ factory := func() (net.Conn, error) { return c.dialer.Dial(nodeAddr, c.timeout) }
+ p, err := pool.NewChannelPool(initialPoolSize, maxPoolCapacity, factory)
if err != nil {
- return "", fmt.Errorf("dial connection: %s", err)
+ return err
+ }
+ c.pools[nodeAddr] = p
+ pl = p
+ return nil
+ }(); err != nil {
+ return "", err
+ }
+ }
+
+ // Got pool, now get a connection.
+ conn, err := pl.Get()
+ if err != nil {
+ return "", fmt.Errorf("pool get: %s", err)
}
defer conn.Close()
+ handleConnError := func(c net.Conn) {
+ if pc, ok := conn.(*pool.PoolConn); ok {
+ pc.MarkUnusable()
+ }
+ }
+
// Send the request
command := &Command{
Type: Command_COMMAND_TYPE_GET_NODE_API_URL,
@@ -48,15 +98,18 @@ func (c *Client) GetNodeAPIAddr(nodeAddr string) (string, error) {
_, err = conn.Write(b)
if err != nil {
+ handleConnError(conn)
return "", fmt.Errorf("write protobuf length: %s", err)
}
_, err = conn.Write(p)
if err != nil {
+ handleConnError(conn)
return "", fmt.Errorf("write protobuf: %s", err)
}
b, err = ioutil.ReadAll(conn)
if err != nil {
+ handleConnError(conn)
return "", fmt.Errorf("read protobuf bytes: %s", err)
}
