From 82b1a34a360d41f46553e8c0d353c428711bc445 Mon Sep 17 00:00:00 2001
From: Alexander Simmerl <a.simmerl@gmail.com>
Date: Tue, 13 Feb 2018 19:08:21 +0100
Subject: [PATCH] Separate connect logic

* break out connect functionality out of OnStart
* introduce max retries
---
 types/priv_validator/socket.go | 93 +++++++++++++++++++++-------------
 1 file changed, 57 insertions(+), 36 deletions(-)

diff --git a/types/priv_validator/socket.go b/types/priv_validator/socket.go
index 4093ce9c..1ca3b7d1 100644
--- a/types/priv_validator/socket.go
+++ b/types/priv_validator/socket.go
@@ -17,6 +17,16 @@ import (
 	"github.com/tendermint/tendermint/types"
 )
 
+const (
+	dialRetryIntervalSeconds = 1
+	dialRetryMax             = 10
+)
+
+// Socket errors.
+var (
+	ErrDialRetryMax = errors.New("Error max client retries")
+)
+
 //-----------------------------------------------------------------
 
 var _ types.PrivValidator2 = (*PrivValidatorSocketClient)(nil)
@@ -33,10 +43,6 @@ type PrivValidatorSocketClient struct {
 	SocketAddress string
 }
 
-const (
-	dialRetryIntervalSeconds = 1
-)
-
 // NewPrivValidatorSocketClient returns an instance of
 // PrivValidatorSocketClient.
 func NewPrivValidatorSocketClient(
@@ -60,39 +66,14 @@ func (pvsc *PrivValidatorSocketClient) OnStart() error {
 		return err
 	}
 
-	var (
-		err  error
-		conn net.Conn
-	)
-
-RETRY_LOOP:
-	for {
-		conn, err = cmn.Connect(pvsc.SocketAddress)
-		if err != nil {
-			pvsc.Logger.Error(
-				"OnStart",
-				"addr", pvsc.SocketAddress,
-				"err", errors.Wrap(err, "connection failed"),
-			)
-
-			time.Sleep(time.Second * dialRetryIntervalSeconds)
-			continue RETRY_LOOP
-		}
-
-		if pvsc.privKey != nil {
-			conn, err = p2pconn.MakeSecretConnection(conn, pvsc.privKey.Wrap())
-			if err != nil {
-				pvsc.Logger.Error(
-					"OnStart",
-					"err", errors.Wrap(err, "encrypting connection failed"),
-				)
-				continue RETRY_LOOP
-			}
-		}
-
-		pvsc.conn = conn
-		return nil
+	conn, err := pvsc.connect()
+	if err != nil {
+		return err
 	}
+
+	pvsc.conn = conn
+
+	return nil
 }
 
 // OnStop implements cmn.Service.
@@ -175,6 +156,46 @@ func (pvsc *PrivValidatorSocketClient) SignHeartbeat(chainID string, heartbeat *
 	return nil
 }
 
+func (pvsc *PrivValidatorSocketClient) connect() (net.Conn, error) {
+	retries := dialRetryMax
+
+RETRY_LOOP:
+	for retries > 0 {
+		if retries != dialRetryMax {
+			time.Sleep(time.Second * dialRetryIntervalSeconds)
+		}
+
+		retries--
+
+		conn, err := cmn.Connect(pvsc.SocketAddress)
+		if err != nil {
+			pvsc.Logger.Error(
+				"OnStart",
+				"addr", pvsc.SocketAddress,
+				"err", errors.Wrap(err, "connection failed"),
+			)
+
+			continue RETRY_LOOP
+		}
+
+		if pvsc.privKey != nil {
+			conn, err = p2pconn.MakeSecretConnection(conn, pvsc.privKey.Wrap())
+			if err != nil {
+				pvsc.Logger.Error(
+					"OnStart",
+					"err", errors.Wrap(err, "encrypting connection failed"),
+				)
+
+				continue RETRY_LOOP
+			}
+		}
+
+		return conn, nil
+	}
+
+	return nil, ErrDialRetryMax
+}
+
 //---------------------------------------------------------
 
 // PrivValidatorSocketServer implements PrivValidator.