From 134dbb1f060646a1dc2693cbf0402d0024d84eba Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 15 Aug 2016 17:24:09 -0700
Subject: [PATCH] If the client detects that a heartbeat has failed because it
 is not registered, reregister

---
 client/client.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/client/client.go b/client/client.go
index 6412b357dea6..c83a49c17f88 100644
--- a/client/client.go
+++ b/client/client.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -737,7 +738,17 @@ func (c *Client) registerAndHeartbeat() {
 		select {
 		case <-heartbeat:
 			if err := c.updateNodeStatus(); err != nil {
-				heartbeat = time.After(c.retryIntv(registerRetryIntv))
+				// The servers have changed such that this node has not been
+				// registered before
+				if strings.Contains(err.Error(), "node not found") {
+					// Re-register the node
+					c.logger.Printf("[INFO] client: re-registering node")
+					c.retryRegisterNode()
+					heartbeat = time.After(lib.RandomStagger(initialHeartbeatStagger))
+				} else {
+					c.logger.Printf("[ERR] client: heartbeating failed: %v", err)
+					heartbeat = time.After(c.retryIntv(registerRetryIntv))
+				}
 			} else {
 				c.heartbeatLock.Lock()
 				heartbeat = time.After(c.heartbeatTTL)