From e55535a53c9cc7e9bceff1114957a71b7165116a Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Thu, 10 Sep 2020 09:52:09 +0300
Subject: [PATCH] fix registry check

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 tcp.go | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/tcp.go b/tcp.go
index 08b14e6..893fa4d 100644
--- a/tcp.go
+++ b/tcp.go
@@ -280,6 +280,8 @@ func (h *tcpServer) Start() error {
 		return err
 	}
 
+	config := h.Options()
+
 	// register
 	if err = h.Register(); err != nil {
 		return err
@@ -308,10 +310,32 @@ func (h *tcpServer) Start() error {
 			select {
 			// register self on interval
 			case <-t.C:
-				if err := h.Register(); err != nil {
-					logger.Error("Server register error: ", err)
+				h.RLock()
+				registered := h.registered
+				h.RUnlock()
+				rerr := h.opts.RegisterCheck(h.opts.Context)
+				if rerr != nil && registered {
+					if logger.V(logger.ErrorLevel) {
+						logger.Errorf("Server %s-%s register check error: %s, deregister it", config.Name, config.Id, rerr)
+					}
+					// deregister self in case of error
+					if err := h.Deregister(); err != nil {
+						if logger.V(logger.ErrorLevel) {
+							logger.Errorf("Server %s-%s deregister error: %s", config.Name, config.Id, err)
+						}
+					}
+				} else if rerr != nil && !registered {
+					if logger.V(logger.ErrorLevel) {
+						logger.Errorf("Server %s-%s register check error: %s", config.Name, config.Id, rerr)
+					}
+					continue
 				}
-			// wait for exit
+				if err := h.Register(); err != nil {
+					if logger.V(logger.ErrorLevel) {
+						logger.Errorf("Server %s-%s register error: %s", config.Name, config.Id, err)
+					}
+				}
+				// wait for exit
 			case ch = <-h.exit:
 				break Loop
 			}