From d250ac736fa5b05797640896dd66c0321e811cb2 Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Fri, 2 Aug 2019 15:17:48 +0100
Subject: [PATCH] In the event watchRegistry or watchTable exit then close the
 go routines

---
 network/router/default.go | 39 +++++++++++++++++++++++++++++++++------
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/network/router/default.go b/network/router/default.go
index bee3d017..498539a1 100644
--- a/network/router/default.go
+++ b/network/router/default.go
@@ -177,10 +177,24 @@ func (r *router) watchRegistry(w registry.Watcher) error {
 	// wait in the background for the router to stop
 	// when the router stops, stop the watcher and exit
 	r.wg.Add(1)
+
+	exit := make(chan bool)
+
+	defer func() {
+		// close the exit channel when the go routine finishes
+		close(exit)
+		r.wg.Done()
+	}()
+
 	go func() {
-		defer r.wg.Done()
-		<-r.exit
-		w.Stop()
+		defer w.Stop()
+
+		select {
+		case <-r.exit:
+			return
+		case <-exit:
+			return
+		}
 	}()
 
 	var watchErr error
@@ -208,10 +222,23 @@ func (r *router) watchTable(w Watcher) error {
 	// wait in the background for the router to stop
 	// when the router stops, stop the watcher and exit
 	r.wg.Add(1)
+	exit := make(chan bool)
+
+	defer func() {
+		// close the exit channel when the go routine finishes
+		close(exit)
+		r.wg.Done()
+	}()
+
 	go func() {
-		defer r.wg.Done()
-		<-r.exit
-		w.Stop()
+		defer w.Stop()
+
+		select {
+		case <-r.exit:
+			return
+		case <-exit:
+			return
+		}
 	}()
 
 	var watchErr error