Best practices on go context cancelation functions

I have been reading some articles on the use of the context package from golang. I recently came across the following article in a blog: http://p.agnihotry.com/post/understanding_the_context_package_in_golang/

The article states the following regarding context cancelation functions in go:

"You can pass around the cancel function if you wanted to, but, that is highly not recommended. This can lead to the invoker of cancel not realizing what the downstream impact of canceling the context may be. There may be other contexts that are derived from this which may cause the program to behave in an unexpected fashion. In short, NEVER pass around the cancel function."

However, passing the cancelation function as a parameter seems to be the only option in case I want the parent context.Done() channel to be activated (see the code snippet below). For instance, the code Done channel in the code snippet below is activated only when function2 is executed.

package main

import (
    "context"
    "fmt"
    "time"
)

func function1(ctx context.Context) {
    _, cancelFunction := context.WithCancel(ctx)
    fmt.Println("cancel called from function1")
    cancelFunction()
}

func function2(ctx context.Context, cancelFunction context.CancelFunc) {
    fmt.Println("cancel called from function2")
    cancelFunction()
}

func main() {
    //Make a background context
    ctx := context.Background()
    //Derive a context with cancel
    ctxWithCancel, cancelFunction := context.WithCancel(ctx)

    go function1(ctxWithCancel)
    time.Sleep(5 * time.Second)

    go function2(ctxWithCancel, cancelFunction)

    time.Sleep(5 * time.Second)

    // Done signal is only received when function2 is called
    <-ctxWithCancel.Done()
    fmt.Println("Done")
}

So, is passing this cancellation function actually an issue? Are there any best practices related to the use of the context package and their cancel function?

Solution

In your specific example, there's a small enough amount of code that there's probably no issue understanding how it all works. The problems start when you replace function1 and function2 with something more complicated. The article you link to gives a specific reason why passing around cancellation contexts can do things that are hard to reason about, but the more general principle is that you should try to separate coordination work (cancellation, spinning up goroutines) from the underlying work to be done (whatever function1 and function2 are doing) as much as possible. This just helps make it easier to reason about sub-sections of your code independently and can help make testing easier. "function2 does <something>" is a lot easier to understand than "function2 does <something> and also coordinates with function1".

Rather than pass the cancellation function into function2, you can just invoke it inside the goroutune you spawn to run function2:

func main() {
  //...
  go func() {
    function2(ctxWithCancel)
    cancelFunction()
  }()
  //...
}

This is niece because the coordination work of figuring out when to cancel is all contained within the calling function rather than be split across multiple functions.

If you want to have function2 cancel the context conditionally, have it explicitly return some kind of value indicating whether some cancellable condition happened:

func function2(ctx context.Context) bool {
  //...
  if workShouldBecanceled() {
    return true
  }
  //...
  return false
}

func main() {
  //...
  go func() {
    if function2(ctxWithCancel) {
      cancelFunction()
    }
  }()
  //...
}

Here I used a boolean, but commonly this pattern is used with errors - if function2 returns a non-nil error, cancel the rest of the work.

Depending on what you're doing, something like errgroup.WithContext might be useful to you. This can coordinate several concurrent actions all of which might fail and cancels the others as soon as the first one fails.

One other best practice that I try to follow with context cancellation: always make sure that the cancel function gets called at some point. From the docs, it is safe to call a cancel function twice, so frequently I end up doing something like this:

func main() {
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  //...
  if shouldCancel() {
    cancel()
  }
  //...
}

Edit in response to comment:

If you have a situation where you have multiple long-running operations (e.g., servers, connections, etc.) and you want to shut all of them down as soon as the first one stops, context cancellation is a reasonable way to do that. However, I'd still recommend that you handle all the context interaction in a single function. Something like this would work:

func operation1(ctx context.Context) {
   for {
     select {
     case <-ctx.Done():
       return
     default:
     }
     //...
   }
}

func operation2(ctx context.Context) {
  // Similar code to operatoin1()
}

func main() {
  ctx, cancel := context.WithCancel(context.Background())
  var wg sync.WaitGroup
  wg.Add(2)
  go func() {
    defer wg.Done()
    defer cancel()
    operation1(ctx)
  }()
  go func() {
    defer wg.Done()
    defer cancel()
    operation2(ctx)
  }()
  wg.Wait()
}

As soon as one of the operations terminates, the other is canceled, but main still waits for both to finish. Neither operation needs to worry about managing this at all.