fix deadlock when batch canceled, some workUnit's done channel can't be closed

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:

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)
 

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 {