Cancellable Functions

Summary: A combination of channel and Context is a very powerful tool for cancellable functions.

There are multiple use cases that we might want to create a cancellable function, especially when developing concurrent programs. We might perform some tasks and if they take too long to complete, we just go ahead to cancel them. Or we might want to cancel them based on other conditions. This post is about common ways to create cancellable functions in Go.

Let's start with a simple use case where we want to abort a function if it takes too long to complete.

Timeout Functions #

Let's assume that we have a function func f() and we want to limit the execution time of f within duration d. This kind of function will stop when either f finishes its operation or d timeout is reached. This reminds us of the select statement in Go, which can listen on multiple channels at the same time.

To solve the above problem, we can use select in combination with channel and a timer like below:

func Do(d time.Duration, f func()) {
	c := make(chan struct{})
	go func() {
		defer close(c)
		f()
	}()
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-c:
		return
	case <-t.C:
		return
	}
}

Playground

The above solution is good, but what if you want to cancel the function not only by timeout but also for other reasons, such as on shutdown or when its caller aborts the mission? We need a more powerful tool for this, and I believe nothing is more suitable than context.Context.

Cancellable Context #

Context was designed for carrying contextual values during function executions but I feel it's more for cancellation. Context provides multiple ways for sending cancellation signals to multiple levels of functions calls either by timeout or on-demand:

// cancellation signal is sent when cancel() is called.
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

// cancellation signal is sent when the timeout is reached.
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

// cancellation signal is sent when the deadline is reached.
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(5*time.Minute))
defer cancel()

// disable cancellation if any in parent context.
ctx := context.WithoutCancel(parentContext)

The same version of the above function can be done with Context as below:

func Do(ctx context.Context, f func()) {
	c := make(chan struct{})
	go func() {
		defer close(c)
		f()
	}()
	select {
	case <-c:
		return
	case <-ctx.Done():
		return
	}
}

And here is how you use it:

// Cancel on demand
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
Do(ctx, func() { /*do something*/ })

// Cancel when timeout is reached.
ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
Do(ctx, func() { /*do something*/ })

// Don't want to cancel it at all
ctx = context.WithoutCancel(context.Background())
Do(ctx, func() { /*do something*/ })

Playground

Generics #

If the function has return values, a similar thing can be done by changing the channel value, and if it's timeout, we want to return error to indicate that the function was failed because of timeout. In this case we can use it with context.Cause as below:

func Do[T any](ctx context.Context, f func() (T, error)) (t T, err error) {
	type v struct {
		t T
		e error
	}
	c := make(chan v)
	go func() {
		defer close(c)
		t, err := f()
		c <- v{t: t, e: err}
	}()
	select {
	case v := <-c:
		return v.t, v.e
	case <-ctx.Done():
		return t, context.Cause(ctx)
	}
}

With this generic version of the function, we can use it with many different functions with different return types:

Do(ctx, func() (int, error) { return 0, nil })
Do(ctx, func() ([]int, error) { return []int{}, nil })
Do(ctx, func() (string, error) { return "", nil })

Playground

Conclusion #

With just a simple combination of channel and context, but we could make a very powerful tool for cancellable functions. You can based on the above suggestion implementations to make your own versions. And I believe that you will find that it's even more powerful than what I have shown you so far.