Skip to content

Instantly share code, notes, and snippets.

@fracasula
Last active May 19, 2022 20:49
Show Gist options
  • Save fracasula/b579d52daf15426e58aa133d0340ccb0 to your computer and use it in GitHub Desktop.
Save fracasula/b579d52daf15426e58aa133d0340ccb0 to your computer and use it in GitHub Desktop.
GoLang exiting from multiple go routines with context and wait group
package main
// Here's a simple example to show how to properly terminate multiple go routines by using a context.
// Thanks to the WaitGroup we'll be able to end all go routines gracefully before the main function ends.
import (
"context"
"fmt"
"math/rand"
"os"
"sync"
"time"
)
func init() {
rand.Seed(time.Now().UnixNano())
}
// it prints all values pushed into "ch" ("ch" here is read only)
func reader(ctx context.Context, wg *sync.WaitGroup, ch <-chan int) {
defer wg.Done() // decrements the WaitGroup counter by one when the function returns
for {
select {
case <-ctx.Done(): // Done returns a channel that's closed when work done on behalf of this context is canceled
fmt.Println("Exiting from reading go routine")
return
case v, ok := <-ch:
if !ok {
fmt.Println("Channel has been closed")
return
}
fmt.Println(v)
}
}
}
// it writes a random integer into "ch" every second ("ch" here is write only)
func writer(ctx context.Context, wg *sync.WaitGroup, ch chan<- int) {
defer wg.Done() // decrements the WaitGroup counter by one when the function returns
for {
select {
case <-ctx.Done(): // Done returns a channel that's closed when work done on behalf of this context is canceled
fmt.Println("Exiting from writing go routine")
return
case ch <- rand.Intn(100): // pushes a random integer from 0 to 100 into the channel
time.Sleep(1 * time.Second) // sleeps one second
}
}
}
func main() {
channel := make(chan int)
defer close(channel)
// a WaitGroup waits for a collection of goroutines to finish, pass this by address
waitGroup := sync.WaitGroup{}
// context.WithCancel returns a copy of parent with a new Done channel.
// The returned context's Done channel is closed when the returned cancel function is called or when the parent
// context's Done channel is closed, whichever happens first.
ctx, cancel := context.WithCancel(context.Background())
waitGroup.Add(2) // adds delta, if the counter becomes zero, all goroutines blocked on Wait are released
go reader(ctx, &waitGroup, channel) // go routine that prints all values pushed into "channel"
go writer(ctx, &waitGroup, channel) // go routine that writes a random integer into "channel" every second
// go routine that listens for an Enter keystroke to terminate the program
go func() {
os.Stdin.Read(make([]byte, 1)) // wait for Enter keystroke
cancel() // cancel the associated context
}()
waitGroup.Wait() // it blocks until the WaitGroup counter is zero
fmt.Println("Exiting from main")
}
@stigok
Copy link

stigok commented Sep 6, 2019

It's a big gotcha to pass the WaitGroup pointer reference. Thanks!

@xianghuzhao
Copy link

xianghuzhao commented Oct 15, 2019

It is unsafe to put wg.Add(1) inside the goroutine. There are chances that this runs after waitGroup.Wait(). So it should be safer to move wg.Add(1) before the goroutine like this:

	waitGroup.Add(1)
	go reader(&waitGroup, ctx, channel) // go routine that prints all values pushed into "channel"
	waitGroup.Add(1)
	go writer(&waitGroup, ctx, channel) // go routine that writes a random integer into "channel" every second

@fracasula
Copy link
Author

Good spot @xianghuzhao I've updated the gist 👍

@akshaybharambe14
Copy link

Thanks, this gist is really helpful.

@landrisek
Copy link

In the main function waitGroup variable should have name "wg" or "waitGroup" otherwise this will not compile.

@johnmackenzie91
Copy link

Amazing. Thanks for sharing this.

@landrisek
Copy link

Do you want me to open merge request for this? It is really nice example, pitty to not have it workable...

@fracasula
Copy link
Author

Do you want me to open merge request for this? It is really nice example, pitty to not have it workable...

@landrisek Busy times 😅 I updated the gist now it works 👍
@johnmackenzie91 glad it's been useful 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment