How to properly let goroutine finish gracefully when main is interrupted?

Question

In the following example, when Ctrl+C is pressed, the work gorountine finishes gracefully before main exits:

package main

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

func work(ctx context.Context, done chan struct{}) {
    defer close(done)
    fmt.Println("work starting its loop")
    out:
    for {
        select {
            case <-ctx.Done():
                fmt.Println("ctx.Done() inside work")
                break out
            default:
                time.Sleep(time.Second)
                fmt.Println("work heartbeat")
        }
    }
    fmt.Println("work finished")
}

func main() {
    interrupt := make(chan os.Signal, 1)
    signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    done := make(chan struct{})
    go func() {
        <-interrupt
        cancel()
    }()
    go work(ctx, done)
    <-ctx.Done()
    fmt.Println("ctx.Done() inside main")
    <-done
    fmt.Println("Main finished")
}

However, I suspect the code is suboptimal. Is the done channel necessary? Shouldn't graceful termination of goroutines be achieved with context alone? What if there were many worker goroutines, not just one?

Answer 1

Actually it has gotten much simpler with signal.NotifyContext: We simply pass a context and get a "ContextWithCancel". Now, if that context is "Done" we received a signal and if we pass the context to all of our goroutines, all goroutines will take note of that and finish.

package main

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

func work(ctx context.Context, number int, wg *sync.WaitGroup) {
    // We use a wait group to ensure that main blocks until the work has finished.
    defer wg.Done()
    fmt.Printf("worker %d starting its loop\n", number)
out:
    for {
        select {
        // We simply reuse the context to check if the work should stop.
        case <-ctx.Done():
            fmt.Println("ctx.Done() inside work")
            break out
        default:
            time.Sleep(time.Second)
            fmt.Printf("work heartbeat from worker %d\n", number)
        }
    }
    fmt.Printf("worker %d finished\n", number)
}

var wg sync.WaitGroup

func main() {
    // No need to use channels here, we can use the context directly.
    ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
    defer stop()

    for i := 0; i < 3; i++ {
        // Add 1 for each goroutine we start.
        wg.Add(1)
        go work(ctx, i, &wg)
    }

    // We received a signal, cancel the context.
    sig := <-ctx.Done()
    fmt.Printf("Got signal: %v\n", sig)
    fmt.Println("ctx.Done() inside main")

    // We ensure that the work has finished before exiting.
    wg.Wait()
    fmt.Println("Main finished")
}