Introduction

Lately, I’ve been working on a hobby project which is a distributed real-time chat app. The objective of this project was to learn the Go programming language.

Having worked with Python for over the past 6 years, I found Go extremely fascinating! Go has this amazing cheap and lightweight thread called goroutine that helps create concurrent applications so much easier.

The Problem

While writing the chat app, I came across a problem of cleaning up the resources when a user disconnects from the server. A general rule of thumb is you should figure out a way your goroutine is going to exit. If it is not handled correctly, your application’s memory and CPU usage will keep on increasing, and eventually end up in panic.

The following function creates a goroutine and listens for any message on a Redis channel

func onUserMessage(r <-chan *redis.Message, w chan string) {
    go func() {
		for m := range r {
			w <- m.Payload
		}
	}()
}

This will work as expected. It gets any message on channel r and sends it to another channel w which I’m using to send the message to a connected user. This goroutine will also exit as expected when the connection to Redis is lost. But what happens when the user gets disconnected? There should be a way to tell this goroutine that we don’t need it now.

The Solution

This is where the context comes into the picture. The same goroutine can be written as follows

func startListener(ctx context.Context) {
	ctx, cancelFn = context.WithCancel(ctx)
	defer cancelFn()
	go onUserMessage(ctx, r, w)
    // Blocking loop
	for {
		// ...do usual stuff
	}
}

func onUserMessage(
    ctx context.Context,
    r <-chan *redis.Message,
    w chan string,
) {
    for {
        select{
            case m, ok := <-r:
            if !ok{
                return
            }else{
                w <- m.Payload
            }
            case <-ctx.Done():
                // ctx.Done returns a channel which is closed
                // when cancel function of that context is called
                return
        }	
    }
}

This is a very clean and correct way to handle the cleanup of a goroutine.