Playing with multiple Contexts in Go

5 min readFeb 22, 2019

I’ve decided to move my blog to Hashnode: https://dlahoza.dev/

Context package has been introduced quite a while but it is still a topic of discussions. Should we use Values or not, how should we pass it to the functions, etc. But this is not the topic of this article.

Let’s imagine that we have several different contexts, which are coming from different sources and we need to behave with them like with a single one.

For example, the first context comes from main() function where we want to control the general execution of the program, the second one, on the other hand, comes from requests or some event calls. Let’s assume that we’re working on some sort of web server with long running background tasks. And we’re starting these tasks on request to the web server. What should we cover here are graceful shutdown and cancellation of those background tasks.

Let’s write some simple task logic:

So we have pretty logger with Zerolog, which is added to requests in middleware. We have a task manager which does “a very complex job”, it adds timeout. Let’s suppose that task workload is blackbox which we should not change, but we can work around it.

Our test cases will be such:

We’re requesting http://localhost:8080/?time=duration. Where duration is a sleep time in seconds.

Happy path:

2019-02-22T10:52:24+02:00 INF Listening HTTP on :8080
2019-02-22T10:52:36+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T10:52:36+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T10:53:06+02:00 INF Task 30 second(s): FINISHED ip=::1 url=/?time=30 user_agent="..."

Task timeout

2019-02-21T21:02:05+02:00 INF Listening HTTP on :8080
2019-02-21T21:02:11+02:00 INF GET /?time=120 ip=::1 url=/?time=120 user_agent="..."
2019-02-21T21:02:11+02:00 INF Task 120 second(s): STARTED ip=::1 url=/?time=120 user_agent="..."
2019-02-21T21:03:11+02:00 INF Task 120 second(s): CANCELED ip=::1 url=/?time=120 user_agent="..."

SIGINT from console

2019-02-21T21:00:45+02:00 INF Listening HTTP on :8080
2019-02-21T21:00:53+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-21T21:00:53+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
^C

Ok, looks fine. We’re using context.Backgound() instead of r.Context() because context which comes from request will be canceled when handler will have finished the request processing and our task will be canceled too. That is why we use the new context with a timeout.

Let’s add graceful shutdown with the timeout of 10 seconds.

Cool… Now web server will be waiting for unfinished requests and only then allow us to stop the application. But what have we missed here?

2019-02-21T21:04:46+02:00 INF Listening HTTP on :8080
2019-02-21T21:04:53+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-21T21:04:53+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
^C
2019-02-21T21:04:57+02:00 INF Shutting down...
2019-02-21T21:04:59+02:00 INF Server has been stopped

All our background jobs will be rudely interrupted without any warning. This is bad. They should be gracefully finished. Moreover, we already have a classic interface to cancel all our jobs ctx context.Context. But this context comes from the handler and task manager, and we cannot seamlessly combine it with some global cancellation mechanism.

We should keep in mind that context.Context is an interface and we can replace it with something more functional which will help us with our issue. So we can just merge existing contexts in one piece. How to do that?

In merge.go you can see all the needed logic for merging two contexts. We combine only cancellation and deadline logic from both of them, but values are taken from the second context only. It’s easy to implement any behavior other on that. Now calling cancel or raising timeout in one of the contexts will cancel both.

In main.go we have some additional code to control cancellation.

mainCtx global variable (which should be moved from globals to task manager object in a real program) and mainCancel function. We will use them to control cancellation
We have RWMutex mu which will be used as a three-position switch. Tasks will be switching it to RLock simultaneously until mutex switched to the full Lock state
Channels timeoutCh and tasksCancelled are used for waiting for tasks to be finished. If timeout is reached we cancel all tasks with mainCancel function and wait for tasksCancelled channel to be closed.

That is it. Let’s try how it works.

2019-02-22T11:20:02+02:00 INF Listening HTTP on :8080
2019-02-22T11:20:07+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:20:07+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
^C
2019-02-22T11:20:10+02:00 INF Shutting down...
2019-02-22T11:20:20+02:00 INF Forcing shutdown
2019-02-22T11:20:20+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:20:20+02:00 INF Server has been stopped

Ok, what about simultaneous tasks?

2019-02-22T11:22:03+02:00 INF Listening HTTP on :8080
2019-02-22T11:22:04+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:04+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF GET /favicon.ico ip=::1 referer=http://localhost:8080/?time=30 url=/favicon.ico user_agent="..."
2019-02-22T11:22:05+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF GET /?time=30 ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:05+02:00 INF Task 30 second(s): STARTED ip=::1 url=/?time=30 user_agent="..."
^C
2019-02-22T11:22:07+02:00 INF Shutting down...
2019-02-22T11:22:17+02:00 INF Forcing shutdown
2019-02-22T11:22:17+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:17+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:17+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:17+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:17+02:00 INF Task 30 second(s): CANCELED ip=::1 url=/?time=30 user_agent="..."
2019-02-22T11:22:17+02:00 INF Server has been stopped

Works like a charm.

As we see here context.Context is a very powerful paradigm. We can easily change its behavior as we need by reimplementing several parts. And by doing that we can use a standard and convenient way to control or cancel our background tasks.

Cheers!