有没有更简洁的方法来创建在通道上接收后被取消的上下文？

Question

I need to invoke a function that takes a Context as a parameter. This block of code has access to a channel that is used for signalling that the operation should be cancelled.

Here is what I am currently using to cancel the Context when a value is received:

func doSomething(stop <-chan bool) {
    ctx, cancel := context.WithCancel(context.Background())
    go func() {
        select {
        case <-ctx.Done():
        case <-stop:
            cancel()
        }
    }()
    longRunningFunction(ctx)
}

The intended control flow is as follows:

• If the task runs to completion, it will cancel the context, the <-ctx.Done() will fire, and the goroutine will terminate.

• If a value is received on stop, the context is cancelled, notifying the task that it should quit. Once again, the goroutine will terminate when this happens.

This seems overly complex. Is there a simpler way to accomplish the intended behavior?

Answer 1

As @ain mentionned, your code currently leaks the goroutine if the longRunningFunction runs to the end and nothing is sent on stop (or it is not closed): the select statement will never be fulfilled (the only way for the context to be done is when something comes out of stop to call cancel).

Here is a way to fix it (mainly an implementation of @ain's comment):

func doSomething(stop <-chan bool) {
    ctx := context.TODO() // because in the future, you might pass a ctx arg to this function, from which you could then "inherit"
    ctx, cancel := context.WithCancel(ctx)
    defer cancel() // to be sure to release the associated resources whatever happens (and prevent the following goroutine from leaking)
    go func() {
        select {
        case <-ctx.Done():
        case <-stop:
            cancel()
        }
    }()
    longRunningFunction(ctx)
}