shutdowntest.go

// Program shutdowntest spins up lots of goroutines and times how long it takes
// to cancel them using various methods.
//
// Baseline measure:
// $ go run . -n 1000000
// 2026/02/12 13:53:42 1000001 goroutines ready in 2.36s
// 2026/02/12 13:53:43 cancelled rootCancel 0.52s after ready
// 2026/02/12 13:53:43 done 0.53s after ready
package main

import (
	"context"
	"flag"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"sync"
	"sync/atomic"
	"time"
)

func main() {
	n := flag.Int("n", 1000, "number of child goroutines to start then be cancelled")
	a := flag.Bool("a", false, "cancel child goroutines in a separate goroutine")
	p := flag.Int("p", -1, "parallelism with which to cancel derived goroutines (<1 disables)")
	t := flag.Bool("t", false, "trace parallel child goroutine context cancellation before root")
	C := flag.Bool("C", false, "have child goroutines use the parent goroutine directly (removes all looping)")
	d := flag.Bool("d", false, "detach parent context from children (doubles goroutine count)")

	cpuFile := flag.String("profile", "", "write a cpu profile to the `file`")
	traceFile := flag.String("memprofile", "", "write a trace profile to the `file`")

	flag.Parse()

	if *cpuFile != "" {
		f, err := os.Create(*cpuFile)
		if err != nil {
			log.Fatalf("error creating cpu profile: %s", err)
		}
		defer f.Close()
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatalf("error starting cpu profiling: %s", err)
		}
		defer pprof.StopCPUProfile()
	}

	if *traceFile != "" {
		f, err := os.Create(*traceFile)
		if err != nil {
			log.Fatalf("error creating cpu trace: %s", err)
		}
		defer f.Close()
		if err := trace.Start(f); err != nil {
			log.Fatalf("error starting cpu profiling: %s", err)
		}
		defer trace.Stop()
	}

	ctx, rootCancel := context.WithCancel(context.Background())

	if *a {
		// Disconnect rootCancel from the cancellation of derived ctxs'
		// cancellation. This causes rootCancel() to return immediately but
		// doesn't affect the performance of cancelling the children.
		//
		// $ go run . -n 1000000 -a
		// 2026/02/12 13:54:08 1000001 goroutines ready in 2.60s
		// 2026/02/12 13:54:08 cancelled rootCancel 0.00s after ready
		// 2026/02/12 13:54:08 done 0.53s after ready
		ctx, _ = withAsyncCancel(ctx)
	}

	if *d {
		// Wrap ctx in a new Context implementation whose cancellation follows
		// ctx but which the context package cannot traverse past. This allows
		// rootCancel to return quickly - because it didn't hear about
		// opacqueCanceller's children to loop over them. To counteract that,
		// however, the children each spin up a goroutine to watch for
		// cancellation, doubling our goroutine count.
		//
		// $ go run . -n 1000000 -d
		// 2026/02/12 13:54:18 1999883 goroutines ready in 3.26s
		// 2026/02/12 13:54:18 cancelled rootCancel 0.00s after ready
		// 2026/02/12 13:54:19 done 0.62s after ready
		doneCh := make(chan struct{})
		go func(ctx context.Context) {
			<-ctx.Done()
			close(doneCh)
		}(ctx)
		ctx = &opaqueCanceller{ctx, doneCh}
	}

	var (
		ctxfn    func() context.Context
		fnCancel func(func(int, context.CancelFunc))
	)
	if *p < 1 {
		ctxfn = func() context.Context { return ctx }
		fnCancel = func(f func(int, context.CancelFunc)) { f(0, rootCancel) }
		if *t {
			log.Fatal("-t can only be used with -p >= 1")
		}
	} else {
		// Build a cancel func which splits the p goroutines and cancels them
		// concurrently in that many goroutines concurrently.
		//
		// $ go run . -n 1000000 -p 1
		// 2026/02/12 13:54:46 1000001 goroutines ready in 2.56s
		// 2026/02/12 13:54:46 cancelled rootCancel 0.00s after ready
		// 2026/02/12 13:54:46 ... parallel cancel
		// 2026/02/12 13:54:47 done 0.53s after ready
		// $ go run . -n 1000000 -p 4
		// 2026/02/12 13:54:52 1000001 goroutines ready in 2.00s
		// 2026/02/12 13:54:52 cancelled rootCancel 0.00s after ready
		// 2026/02/12 13:54:52 ... parallel cancel
		// 2026/02/12 13:54:52 done 0.17s after ready
		ctxfn, fnCancel = withParallelCancellation(ctx, *p)
	}

	startSpinup := time.Now()

	// Start n child goroutines.
	var ready, done sync.WaitGroup
	ready.Add(*n)
	for range *n {
		done.Go(func() {
			ctx := ctxfn()
			if !*C {
				// Wrap the root context with another cancellation. Disabling
				// this sees the cancellation run very quickly, but is
				// unrealistic. In practice each *http.Request context is
				// derived from the BaseContext and wraps for its own
				// cancellation. So this is the default.
				//
				// $ go run . -n 1000000 -C
				// 2026/02/12 13:53:51 1000001 goroutines ready in 0.91s
				// 2026/02/12 13:53:51 cancelled rootCancel 0.21s after ready
				// 2026/02/12 13:53:51 done 0.21s after ready
				c, cancel := context.WithCancel(ctx)
				defer cancel()
				ctx = c
			}
			ready.Done()
			<-ctx.Done()
		})
	}

	log.Printf("%d goroutines ready in %.2fs", runtime.NumGoroutine(), time.Since(startSpinup).Seconds())
	startCancel := time.Now()

	if *t {
		// Cancel the children goroutines before cancelling the root context.
		// This doesn't impact the overall cancellation time, but confirms that
		// we don't accidentally do double work when both cancel funcs are called.
		//
		// $ go run . -n 1000000 -p 4 -t
		// 2026/02/12 13:55:01 1000001 goroutines ready in 2.06s
		// 2026/02/12 13:55:01 ... parallel cancel
		// 2026/02/12 13:55:01 cancelled fnCancel/2 0.16s after ready
		// 2026/02/12 13:55:01 cancelled fnCancel/1 0.16s after ready
		// 2026/02/12 13:55:01 cancelled fnCancel/0 0.16s after ready
		// 2026/02/12 13:55:01 cancelled fnCancel/3 0.16s after ready
		// 2026/02/12 13:55:01 cancelled fnCancel 0.16s after ready
		// 2026/02/12 13:55:01 cancelled rootCancel 0.16s after ready
		// 2026/02/12 13:55:01 done 0.16s after ready
		fnCancel(func(i int, cancel context.CancelFunc) {
			cancel()
			log.Printf("cancelled fnCancel/%01d %0.2fs after ready", i, time.Since(startCancel).Seconds())
		})
		log.Printf("cancelled fnCancel %.2fs after ready", time.Since(startCancel).Seconds())
	}
	rootCancel()
	log.Printf("cancelled rootCancel %.2fs after ready", time.Since(startCancel).Seconds())

	done.Wait()
	log.Printf("done %.2fs after ready", time.Since(startCancel).Seconds())
}

// withAsyncCancel causes root context cancellation to trigger its child
// cancellation in a background goroutine.
func withAsyncCancel(parent context.Context) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(context.WithoutCancel(parent))
	context.AfterFunc(parent, cancel)
	return ctx, cancel
}

// withParalleCancellation returns a func which returns 1 of n contexts derived
// from ctx, and a cancel func which cancels al goroutines.
//
// As a special case if n == 1, withParallelCancellation is equivalent to
// context.WithCancel (but returns a func to get the ctx).
//
// When n is < 1, withParallelCancellation panics.
func withParallelCancellation(parent context.Context, n int) (
	func() context.Context,
	func(run func(i int, cancel context.CancelFunc)),
) {
	if n < 1 {
		panic("withParallelCancellation call with n < 1")
	}

	// We detach derived contexts' direct cancellations because we don't want
	// the default single-threaded cancellation. We rewire cancel propagation
	// with context.AfterFunc below.
	ctx := context.WithoutCancel(parent)

	// Create n child contexts.
	type ctxCancel struct {
		context.Context
		context.CancelFunc
	}
	ctxs := make([]ctxCancel, n)
	for i := range ctxs {
		ctx, cancel := context.WithCancel(ctx)
		ctxs[i] = ctxCancel{ctx, cancel}
	}

	// Cancel n child contexts with n separate goroutines.
	var cancelOnce sync.Once
	cancelWrap := func(run func(i int, cancel context.CancelFunc)) {
		cancelOnce.Do(func() {
			log.Print("... parallel cancel")
			var wg sync.WaitGroup
			defer wg.Wait()
			for i, cc := range ctxs {
				wg.Go(func() {
					defer cc.CancelFunc()
					if run != nil {
						run(i, cc.CancelFunc)
					}
				})
			}
		})
	}
	cancel := func() { cancelWrap(nil) }

	// Rewire the parent cancellation to cancel these children in parallel.
	context.AfterFunc(parent, cancel)

	// Return a func for getting one of n ctxs, and cancellation.
	var p atomic.Int64
	return func() context.Context { return ctxs[int(p.Add(1)%int64(n))] }, cancelWrap
}

// opaqueCanceller has no affect on context features:  cancellation, deadlines
// or values. It just prevents the context package from registering derived
// goroutines with the parent cancellation.
type opaqueCanceller struct {
	ctx  context.Context
	done chan struct{}
}

var _ context.Context = (*opaqueCanceller)(nil)

func (c *opaqueCanceller) Value(v any) any             { return c.ctx.Value(v) }
func (c *opaqueCanceller) Deadline() (time.Time, bool) { return c.ctx.Deadline() }
func (c *opaqueCanceller) Done() <-chan struct{}       { return c.done }
func (c *opaqueCanceller) Err() error {
	select {
	case <-c.done:
		return context.Canceled
	default:
		return nil
	}
}