From 66e914a8ab39567201e379eea4c4d78a64360ea2 Mon Sep 17 00:00:00 2001
From: Ed Robinson <ed.robinson@reevoo.com>
Date: Tue, 15 Nov 2016 19:14:11 +0000
Subject: [PATCH] Adds Cleanup method to safe.Pool

---
 safe/routine.go | 18 ++++++++++++++----
 server.go       |  2 +-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/safe/routine.go b/safe/routine.go
index 28014c5b8..70937fbb6 100644
--- a/safe/routine.go
+++ b/safe/routine.go
@@ -21,18 +21,20 @@ type Pool struct {
 	waitGroup   sync.WaitGroup
 	lock        sync.Mutex
 	baseCtx     context.Context
+	baseCancel  context.CancelFunc
 	ctx         context.Context
 	cancel      context.CancelFunc
 }
 
 // NewPool creates a Pool
 func NewPool(parentCtx context.Context) *Pool {
-	baseCtx, _ := context.WithCancel(parentCtx)
+	baseCtx, baseCancel := context.WithCancel(parentCtx)
 	ctx, cancel := context.WithCancel(baseCtx)
 	return &Pool{
-		baseCtx: baseCtx,
-		ctx:     ctx,
-		cancel:  cancel,
+		baseCtx:    baseCtx,
+		baseCancel: baseCancel,
+		ctx:        ctx,
+		cancel:     cancel,
 	}
 }
 
@@ -90,6 +92,14 @@ func (p *Pool) Stop() {
 	}
 }
 
+// Cleanup releases resources used by the pool, and should be called when the pool will no longer be used
+func (p *Pool) Cleanup() {
+	p.Stop()
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	p.baseCancel()
+}
+
 // Start starts all stopped routines
 func (p *Pool) Start() {
 	p.lock.Lock()
diff --git a/server.go b/server.go
index af338e963..83f97864b 100644
--- a/server.go
+++ b/server.go
@@ -133,7 +133,7 @@ func (server *Server) Close() {
 		}
 	}(ctx)
 	server.stopLeadership()
-	server.routinesPool.Stop()
+	server.routinesPool.Cleanup()
 	close(server.configurationChan)
 	close(server.configurationValidatedChan)
 	signal.Stop(server.signals)