Shutdown after shutdown

kc596 · Nov 3, 2020 · 7d0d73a · 7d0d73a
1 parent 0640e39
commit 7d0d73a
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 2 deletions.
diff --git a/go.mod b/go.mod
@@ -3,6 +3,6 @@ module github.com/kc596/priorityworkerpool
 go 1.15
 
 require (
-	github.com/kc596/UGCPriorityQueue v1.2.3
+	github.com/kc596/UGCPriorityQueue v1.2.5
 	github.com/stretchr/testify v1.6.1
 )
diff --git a/go.sum b/go.sum
@@ -2,6 +2,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/kc596/UGCPriorityQueue v1.2.3 h1:jVJ0nMT8W7Y9KYfBFhz3/1guLiGzyIpFmz8ABavh65c=
 github.com/kc596/UGCPriorityQueue v1.2.3/go.mod h1:gxnMPbwjCHqnHngEq3VcE8Cx3MQXOwK3U6dmMGyNoDU=
+github.com/kc596/UGCPriorityQueue v1.2.5 h1:/rcTVRZS59XrGx+sZI474t0skQqO/hYj4PngQ2ssl0A=
+github.com/kc596/UGCPriorityQueue v1.2.5/go.mod h1:gxnMPbwjCHqnHngEq3VcE8Cx3MQXOwK3U6dmMGyNoDU=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

diff --git a/pool.go b/pool.go
@@ -72,6 +72,9 @@ func (pool *Pool) WaitGroup() *sync.WaitGroup {
 // ShutDown prevents pickup of next job from the queue
 // For stopping the already picked up work, use context
 func (pool *Pool) ShutDown() {
+	if atomic.LoadUint32(&pool.active) == uint32(0) {
+		return
+	}
 	pool.shutDownCh <- true
 }
 
@@ -85,7 +88,8 @@ func (pool *Pool) start() {
 			select {
 			case <-pool.shutDownCh:
 				atomic.StoreUint32(&pool.active, uint32(0))
-				pool.jobQueue = nil // for garbage collection
+				pool.jobQueue.Clear()
+				close(pool.shutDownCh)
 				return
 			default:
 				if pool.jobQueue.Size() > 0 {

diff --git a/pool_test.go b/pool_test.go
@@ -124,4 +124,10 @@ func TestPoolShutDown(t *testing.T) {
 	pool.Submit(func() { atomic.AddUint32(&executed, 1) }, 1+rand.Float64())
 	assert.Equal(uint32(1), atomic.LoadUint32(&executed))
 	assert.Equal(uint32(1), panicCount)
+
+	// shutdown after shutdown
+	pool.ShutDown()
+	pool.ShutDown()
+	assert.Zero(panicCount)
+	assert.Equal(uint32(1), atomic.LoadUint32(&executed))
 }