adding timeouts and deadlines

cluster/listener_test.go

@@ -2,6 +2,7 @@ package cluster_test
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -112,7 +113,7 @@ func TestServer_WriteShardRequestSuccess(t *testing.T) {
 	// Close the server
 	defer s.Close()
 
-	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()})
+	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()}, time.Minute)
 
 	now := time.Now()
 
@@ -176,7 +177,7 @@ func TestServer_WriteShardRequestMultipleSuccess(t *testing.T) {
 	// Close the server
 	defer s.Close()
 
-	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()})
+	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()}, time.Minute)
 
 	now := time.Now()
 
@@ -250,7 +251,7 @@ func TestServer_WriteShardRequestFail(t *testing.T) {
 	// Close the server
 	defer s.Close()
 
-	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()})
+	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()}, time.Minute)
 	now := time.Now()
 
 	shardID := uint64(1)
@@ -264,3 +265,30 @@ func TestServer_WriteShardRequestFail(t *testing.T) {
 		t.Fatalf("expected error %s, got %v", exp, err)
 	}
 }
+
+func TestServer_DialTimeout(t *testing.T) {
+	var (
+		ts = newTestServer(writeShardSuccess)
+		s  = cluster.NewServer(ts, "127.0.0.1:0")
+	)
+	// Start on a random port
+	if e := s.Open(); e != nil {
+		t.Fatalf("err does not match.  expected %v, got %v", nil, e)
+	}
+	// Close the server
+	defer s.Close()
+
+	writer := cluster.NewWriter(&metaStore{host: s.Addr().String()}, time.Nanosecond)
+	now := time.Now()
+
+	shardID := uint64(1)
+	ownerID := uint64(2)
+	var points []tsdb.Point
+	points = append(points, tsdb.NewPoint(
+		"cpu", tsdb.Tags{"host": "server01"}, map[string]interface{}{"value": int64(100)}, now,
+	))
+
+	if err, exp := writer.Write(shardID, ownerID, points), "i/o timeout"; err == nil || !strings.Contains(err.Error(), exp) {
+		t.Fatalf("expected error %v, to contain %s", err, exp)
+	}
+}

cluster/writer.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"time"
 
 	"github.com/fatih/pool"
 	"github.com/influxdb/influxdb/meta"
@@ -16,7 +17,10 @@ const (
 	writeShardResponseMessage
 )
 
-const maxConnections = 500
+const (
+	maxConnections = 500
+	maxRetries     = 3
+)
 
 var errMaxConnectionsExceeded = fmt.Errorf("can not exceed max connections of %d", maxConnections)
 
@@ -27,6 +31,7 @@ type metaStore interface {
 type connFactory struct {
 	metaStore  metaStore
 	nodeID     uint64
+	timeout    time.Duration
 	clientPool interface {
 		size() int
 	}
@@ -42,22 +47,29 @@ func (c *connFactory) dial() (net.Conn, error) {
 		return nil, err
 	}
 
-	conn, err := net.Dial("tcp", nodeInfo.Host)
-	if err != nil {
-		return nil, err
+	var retries int
+	for {
+		conn, err := net.DialTimeout("tcp", nodeInfo.Host, c.timeout)
+		if err != nil && retries == maxRetries {
+			return nil, err
+		} else if err == nil {
+			return conn, nil
+		}
+		retries++
 	}
-	return conn, nil
 }
 
 type Writer struct {
 	pool      *clientPool
 	metaStore metaStore
+	timeout   time.Duration
 }
 
-func NewWriter(m metaStore) *Writer {
+func NewWriter(m metaStore, timeout time.Duration) *Writer {
 	return &Writer{
 		pool:      newClientPool(),
 		metaStore: m,
+		timeout:   timeout,
 	}
 }
 
@@ -65,7 +77,7 @@ func (c *Writer) dial(nodeID uint64) (net.Conn, error) {
 	// if we don't have a connection pool for that addr yet, create one
 	_, ok := c.pool.getPool(nodeID)
 	if !ok {
-		factory := &connFactory{nodeID: nodeID, metaStore: c.metaStore, clientPool: c.pool}
+		factory := &connFactory{nodeID: nodeID, metaStore: c.metaStore, clientPool: c.pool, timeout: c.timeout}
 		p, err := pool.NewChannelPool(1, 3, factory.dial)
 		if err != nil {
 			return nil, err
@@ -84,6 +96,7 @@ func (w *Writer) Write(shardID, ownerID uint64, points []tsdb.Point) error {
 	// This will return the connection to the data pool
 	defer conn.Close()
 
+	conn.SetWriteDeadline(time.Now().Add(w.timeout))
 	var mt byte = writeShardRequestMessage
 	if err := binary.Write(conn, binary.LittleEndian, &mt); err != nil {
 		return err
@@ -100,19 +113,23 @@ func (w *Writer) Write(shardID, ownerID uint64, points []tsdb.Point) error {
 
 	size := int64(len(b))
 
+	conn.SetWriteDeadline(time.Now().Add(w.timeout))
 	if err := binary.Write(conn, binary.LittleEndian, &size); err != nil {
 		return err
 	}
 
+	conn.SetWriteDeadline(time.Now().Add(w.timeout))
 	if _, err := conn.Write(b); err != nil {
 		return err
 	}
 
+	conn.SetReadDeadline(time.Now().Add(w.timeout))
 	// read back our response
 	if err := binary.Read(conn, binary.LittleEndian, &mt); err != nil {
 		return err
 	}
 
+	conn.SetReadDeadline(time.Now().Add(w.timeout))
 	if err := binary.Read(conn, binary.LittleEndian, &size); err != nil {
 		return err
 	}
@@ -120,6 +137,7 @@ func (w *Writer) Write(shardID, ownerID uint64, points []tsdb.Point) error {
 	message := make([]byte, size)
 
 	reader := io.LimitReader(conn, size)
+	conn.SetReadDeadline(time.Now().Add(w.timeout))
 	_, err = reader.Read(message)
 	if err != nil {
 		return err
@@ -136,6 +154,7 @@ func (w *Writer) Write(shardID, ownerID uint64, points []tsdb.Point) error {
 
 	return nil
 }
+
 func (w *Writer) Close() error {
 	if w.pool == nil {
 		return fmt.Errorf("client already closed")