cockroach: context: add PropagateCancel

Author: RaduBerinde

api/go1.txt

@@ -30869,3 +30869,4 @@ pkg unicode/utf8, func RuneLen(int32) int
 pkg unicode/utf8, func RuneStart(uint8) bool
 pkg unicode/utf8, func Valid([]uint8) bool
 pkg unicode/utf8, func ValidString(string) bool
+pkg context, func PropagateCancel(Context, func()) bool

src/context/context.go

@@ -62,6 +62,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	_ "unsafe"  // for go:linkname
 )
 
 // A Context carries a deadline, a cancellation signal, and other values across
@@ -371,6 +372,7 @@ type stopCtx struct {
 var goroutines atomic.Int32
 
 // &cancelCtxKey is the key that a cancelCtx returns itself for.
+//go:linkname cancelCtxKey
 var cancelCtxKey int
 
 // parentCancelCtx returns the underlying *cancelCtx for parent.
@@ -490,17 +492,7 @@ func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
 
 	if p, ok := parentCancelCtx(parent); ok {
 		// parent is a *cancelCtx, or derives from one.
-		p.mu.Lock()
-		if err := p.err.Load(); err != nil {
-			// parent has already been canceled
-			child.cancel(false, err.(error), p.cause)
-		} else {
-			if p.children == nil {
-				p.children = make(map[canceler]struct{})
-			}
-			p.children[child] = struct{}{}
-		}
-		p.mu.Unlock()
+		p.addChild(child)
 		return
 	}
 
@@ -528,6 +520,22 @@ func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
 	}()
 }
 
+// addChild adds child to the list of children.
+// NB: CockroachDB runtime patch.
+func (c *cancelCtx) addChild(child canceler) {
+	c.mu.Lock()
+	if err := c.err.Load(); err != nil {
+		// parent has already been canceled
+		child.cancel(false, err.(error), c.cause)
+	} else {
+		if c.children == nil {
+			c.children = make(map[canceler]struct{})
+		}
+		c.children[child] = struct{}{}
+	}
+	c.mu.Unlock()
+}
+
 type stringer interface {
 	String() string
 }
@@ -804,3 +812,33 @@ func value(c Context, key any) any {
 		}
 	}
 }
+
+// CockroachDB runtime patch.
+// cancelerAdapter invokes f when cancel context completes.
+type cancelerAdapter struct {
+	*cancelCtx
+	f func()
+}
+
+func (c *cancelerAdapter) cancel(removeFromParent bool, err, cause error) {
+	if removeFromParent {
+		removeChild(c.cancelCtx, c)
+	}
+	c.f()
+}
+
+// PropagateCancel arranges for f to be invoked when parent is done.
+// Parent must be one of the cancelable contexts.
+// Returns true if cancellation will be propagated, false if the parent
+// is not cancelable.
+// This is similar to AfterFunc(), but does not spin up goroutine, and instead
+// invokes f on whatever goroutine completed parent context.
+func PropagateCancel(parent Context, f func()) bool {
+	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
+	if !ok {
+		return false
+	}
+	a := cancelerAdapter{cancelCtx: p, f: f}
+	p.addChild(&a)
+	return true
+}