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?
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")
}