Gracefully stopping goroutines

Recently I developed StrangerBot, a bot built using the Telegram Bot API. It’s pretty cool. It got more than 10,000 users in a week.

This bot was quickly hacked together, so I keep improving it and deploying fixes on a daily basis. I make use of channels that holds the received messages, and redeploying causes a loss of these messages.

My application contains one goroutine, that continuously tries to receive the newest messages sent to the bot. Once received it passes them into a channel and they are handled by several workers.

So I went to figure out how to gracefully stop all running goroutines before quitting the application, in order to prevent losing the messages.

I ended up using a sync.WaitGroup and a channel that listens to signals. See the example:

package main

import (
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

func worker(jobs <-chan int) {
    for job := range jobs {
        fmt.Println(job)
        time.Sleep(500 * time.Millisecond)
    }
}

func main() {

        jobs := make(chan int, 1000)

        var wg sync.WaitGroup

        // Spawn example workers
        for i := 0; i < 3; i++ {
            wg.Add(1)
            go func(jobs <-chan int) {
                defer wg.Done()
                worker(jobs)
                fmt.Println("Worker is done.")
            }(jobs)
        }

        // Create example messages
        for i := 0; i < 100; i++ {
            jobs <- i
        }

        waitForSignal()

        close(jobs)
        fmt.Println("Closed jobs!")
        wg.Wait()
}

func waitForSignal() {
    sigs := make(chan os.Signal, 1)
    signal.Notify(sigs, os.Interrupt)
    signal.Notify(sigs, syscall.SIGTERM)
    <-sigs
}

In this example the application will block on waitForSignal(). When a signal is received the jobs channel will be closed. To prevent the loss of messages we use a WaitGroup, this way we can make sure all goroutines finished running, otherwise the application will still exit immediately and we will still lose messages.

Be wary though, once you’re using multiple channels that communicate with each other the solution might not be as simple as this, and you might require several WaitGroups in order to achieve a graceful shutdown without losing any messages.

Hope it was helpful :).