Fix parent vs child race (#1141)

Author: demirkayaender

internal/context.go

@@ -223,14 +223,19 @@ func propagateCancel(parent Context, child canceler) {
 		return // parent is never canceled
 	}
 	if p, ok := parentCancelCtx(parent); ok {
+		p.cancelLock.Lock()
 		if p.err != nil {
+			p.cancelLock.Unlock()
 			// parent has already been canceled
 			child.cancel(false, p.err)
 		} else {
+			p.childrenLock.Lock()
 			if p.children == nil {
 				p.children = make(map[canceler]bool)
 			}
 			p.children[child] = true
+			p.childrenLock.Unlock()
+			p.cancelLock.Unlock()
 		}
 	} else {
 		go func() {
@@ -269,6 +274,9 @@ func removeChild(parent Context, child canceler) {
 	if !ok {
 		return
 	}
+
+	p.childrenLock.Lock()
+	defer p.childrenLock.Unlock()
 	if p.children != nil {
 		delete(p.children, child)
 	}
@@ -288,11 +296,12 @@ type cancelCtx struct {
 
 	done Channel // closed by the first cancel call.
 
-	mu       sync.Mutex
-	canceled bool
+	cancelLock sync.Mutex
+	canceled   bool
 
-	children map[canceler]bool // set to nil by the first cancel call
-	err      error             // set to non-nil by the first cancel call
+	childrenLock sync.Mutex
+	children     map[canceler]bool // set to nil by the first cancel call
+	err          error             // set to non-nil by the first cancel call
 }
 
 func (c *cancelCtx) Done() Channel {
@@ -307,32 +316,48 @@ func (c *cancelCtx) String() string {
 	return fmt.Sprintf("%v.WithCancel", c.Context)
 }
 
+func (c *cancelCtx) getChildren() []canceler {
+	c.childrenLock.Lock()
+	defer c.childrenLock.Unlock()
+
+	out := []canceler{}
+	for key := range c.children {
+		out = append(out, key)
+	}
+	return out
+}
+
 // cancel closes c.done, cancels each of c's children, and, if
 // removeFromParent is true, removes c from its parent's children.
 func (c *cancelCtx) cancel(removeFromParent bool, err error) {
-	c.mu.Lock()
+	c.cancelLock.Lock()
 	if c.canceled {
-		c.mu.Unlock()
+		c.cancelLock.Unlock()
 		// calling cancel from multiple go routines isn't safe
 		// avoid a data race by only allowing the first call
 		return
 	}
 	c.canceled = true
-	c.mu.Unlock()
 
 	if err == nil {
 		panic("context: internal error: missing cancel error")
 	}
 	if c.err != nil {
+		c.cancelLock.Unlock()
 		return // already canceled
 	}
 	c.err = err
+	c.cancelLock.Unlock()
 	c.done.Close()
-	for child := range c.children {
+
+	children := c.getChildren()
+	for _, child := range children {
 		// NOTE: acquiring the child's lock while holding parent's lock.
 		child.cancel(false, err)
 	}
+	c.childrenLock.Lock()
 	c.children = nil
+	c.childrenLock.Unlock()
 
 	if removeFromParent {
 		removeChild(c.Context, c)

internal/context_test.go

@@ -28,7 +28,49 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestContext_RaceRegression(t *testing.T) {
+func TestContextChildParentCancelRace(t *testing.T) {
+	/*
+		Testing previous race happened while child and parent cancelling at the same time
+		While child is trying to remove itself from the parent, parent tries to iterate
+		its children and cancel them at the same time.
+	*/
+	env := newTestWorkflowEnv(t)
+
+	wf := func(ctx Context) error {
+		parentCtx, parentCancel := WithCancel(ctx)
+		defer parentCancel()
+
+		type cancelerContext struct {
+			ctx      Context
+			canceler func()
+		}
+
+		children := []cancelerContext{}
+		numChildren := 100
+
+		for i := 0; i < numChildren; i++ {
+			c, canceler := WithCancel(parentCtx)
+			children = append(children, cancelerContext{
+				ctx:      c,
+				canceler: canceler,
+			})
+		}
+
+		for i := 0; i < numChildren; i++ {
+			go children[i].canceler()
+			if i == numChildren/2 {
+				go parentCancel()
+			}
+		}
+
+		return nil
+	}
+	env.RegisterWorkflow(wf)
+	env.ExecuteWorkflow(wf)
+	assert.NoError(t, env.GetWorkflowError())
+}
+
+func TestContextConcurrentCancelRace(t *testing.T) {
 	/*
 		A race condition existed due to concurrently ending goroutines on shutdown (i.e. closing their chan without waiting
 		on them to finish shutdown), which executed... quite a lot of non-concurrency-safe code in a concurrent way.  All
@@ -57,3 +99,32 @@ func TestContext_RaceRegression(t *testing.T) {
 	env.ExecuteWorkflow(wf)
 	assert.NoError(t, env.GetWorkflowError())
 }
+
+func TestContextAddChildCancelParentRace(t *testing.T) {
+	/*
+		It's apparently also possible to race on adding children while propagating the cancel to children.
+	*/
+	env := newTestWorkflowEnv(t)
+	wf := func(ctx Context) error {
+		ctx, cancel := WithCancel(ctx)
+		racyCancel := func(ctx Context) {
+			defer cancel() // defer is necessary as Sleep will never return due to Goexit
+			defer func() {
+				_, ccancel := WithCancel(ctx)
+				cancel()
+				ccancel()
+			}()
+			_ = Sleep(ctx, time.Hour)
+		}
+		// start a handful to increase odds of a race being detected
+		for i := 0; i < 10; i++ {
+			Go(ctx, racyCancel)
+		}
+
+		_ = Sleep(ctx, time.Minute) // die early
+		return nil
+	}
+	env.RegisterWorkflow(wf)
+	env.ExecuteWorkflow(wf)
+	assert.NoError(t, env.GetWorkflowError())
+}