Writing cancellable functions

Using existing functions that accept cancellable context is easy.

Writing a function that can be cancelled via context is much harder.

When a time experies or you call cancel function returned by context.WithCancel() or context.WithTimeout(), a channel in the context is signalled.

When writing a cancellable function, you have to periodically check channel returned by context.Done() and return immediately if it has been signalled.

It does make for an awkward code:

func longMathOp(ctx context.Context, n int) (int, error) {
	res := n
	for i := 0; i < 100; i++ {
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		default:
			res += i
			// simulate long operation by sleeping
			time.Sleep(time.Millisecond)
		}
	}
	return res, nil
}

func main() {
	ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*200)
	res, err := longMathOp(ctx, 5)
	fmt.Printf("Called longMathOp() with 200ms timeout. res; %d, err: %v\n", res, err)

	ctx, _ = context.WithTimeout(context.Background(), time.Millisecond*10)
	res, err = longMathOp(ctx, 5)
	fmt.Printf("Called longMathOp() with 10ms timeout. res: %d, err: %v\n", res, err)
}

Called longMathOp() with 200ms timeout. res; 4955, err: <nil>
Called longMathOp() with 10ms timeout. res: 0, err: context deadline exceeded

For clarity, this is an artificial task.

Our longMathOp function performs simple operation 100 times and simulates slowness by sleeping for 1 ms on every iteration.

We can expect it to take ~100 ms.

A select with default clause is non-blocking. If there’s nothing in the ctx.Done() channel, we don’t wait for values and immediately execute default part, which is where the logic of the program lives.

We can see in our test that if timeout is greater than 100 ms, the function finishes.

If timeout is smaller than 100 ms, ctx.Done() channel is signalled, we detect it in longMathOp and return ctx.Err().