From 9aae608650b0d7254a9bdd713ce0e8167e9749e9 Mon Sep 17 00:00:00 2001
From: vankli <vankli@qq.com>
Date: Thu, 12 Sep 2019 16:54:12 +0800
Subject: [PATCH] fix deadlock when batch canceled, some workUnit's done
 channel can't be closed

---
 limited_pool.go   | 2 ++
 unlimited_pool.go | 2 ++
 work_unit.go      | 8 +++++++-
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/limited_pool.go b/limited_pool.go
index cd4e317..4dbb5c6 100644
--- a/limited_pool.go
+++ b/limited_pool.go
@@ -88,6 +88,7 @@ func (p *limitedPool) newWorker(work chan *workUnit, cancel chan struct{}) {
 				if wu.cancelled.Load() == nil {
 					value, err = wu.fn(wu)
 
+					wu.Lock()
 					wu.writing.Store(struct{}{})
 
 					// need to check again in case the WorkFunc cancelled this unit of work
@@ -100,6 +101,7 @@ func (p *limitedPool) newWorker(work chan *workUnit, cancel chan struct{}) {
 						// of work to be done first so we use close
 						close(wu.done)
 					}
+					wu.Unlock()
 				}
 
 			case <-cancel:
diff --git a/unlimited_pool.go b/unlimited_pool.go
index d1f5beb..d60e74e 100644
--- a/unlimited_pool.go
+++ b/unlimited_pool.go
@@ -75,6 +75,7 @@ func (p *unlimitedPool) Queue(fn WorkFunc) WorkUnit {
 		if w.cancelled.Load() == nil {
 			val, err := w.fn(w)
 
+			w.Lock()
 			w.writing.Store(struct{}{})
 
 			// need to check again in case the WorkFunc cancelled this unit of work
@@ -88,6 +89,7 @@ func (p *unlimitedPool) Queue(fn WorkFunc) WorkUnit {
 				// of work to be done first so we use close
 				close(w.done)
 			}
+			w.Unlock()
 		}
 	}(w)
 
diff --git a/work_unit.go b/work_unit.go
index 9d0c75f..67533a2 100644
--- a/work_unit.go
+++ b/work_unit.go
@@ -1,6 +1,9 @@
 package pool
 
-import "sync/atomic"
+import (
+	"sync"
+	"sync/atomic"
+)
 
 // WorkUnit contains a single uint of works values
 type WorkUnit interface {
@@ -35,6 +38,7 @@ type workUnit struct {
 	cancelled  atomic.Value
 	cancelling atomic.Value
 	writing    atomic.Value
+	sync.Mutex
 }
 
 // Cancel cancels this specific unit of work, if not already committed to processing.
@@ -44,6 +48,8 @@ func (wu *workUnit) Cancel() {
 
 func (wu *workUnit) cancelWithError(err error) {
 
+	wu.Lock()
+	defer wu.Unlock()
 	wu.cancelling.Store(struct{}{})
 
 	if wu.writing.Load() == nil && wu.cancelled.Load() == nil {