Leapcell

Posted on May 28

How to Wait for Multiple Goroutines in Go

#webdev #programming #backend #go

In Go, the main goroutine often needs to wait for other goroutines to finish their tasks before continuing execution or exiting the program. This is a common requirement for concurrent synchronization. Go provides several mechanisms to achieve this, depending on the scenario and requirements.

Method 1: Using sync.WaitGroup

sync.WaitGroup is the most commonly used synchronization tool in Go, designed to wait for a group of goroutines to finish their tasks. It works through a counter mechanism and is especially suitable when the main goroutine needs to wait for multiple sub-goroutines.

Example Code

package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup

    // Start 3 goroutines
    for i := 1; i <= 3; i++ {
        wg.Add(1) // Increment the counter by 1
        go func(id int) {
            defer wg.Done() // Decrement the counter by 1 when the task is done
            fmt.Printf("Goroutine %d is running\n", id)
        }(i)
    }

    wg.Wait() // Main goroutine waits for all goroutines to finish
    fmt.Println("All goroutines finished")
}

Output (order may vary):

Goroutine 1 is running
Goroutine 2 is running
Goroutine 3 is running
All goroutines finished

How it works:

wg.Add(n): Increases the counter to indicate the number of goroutines to wait for.
wg.Done(): Called by each goroutine upon completion, decreases the counter by 1.
wg.Wait(): Blocks the main goroutine until the counter reaches zero.

Advantages:

Simple and easy to use, suitable for a fixed number of goroutines.
No need for additional channels, low performance overhead.

Method 2: Using Channel

By passing signals through channels, the main goroutine can wait until all other goroutines have sent completion signals. This method is more flexible, but usually a bit more complex than WaitGroup.

Example Code

package main

import "fmt"

func main() {
    done := make(chan struct{}) // A signal channel to notify completion
    numGoroutines := 3

    for i := 1; i <= numGoroutines; i++ {
        go func(id int) {
            fmt.Printf("Goroutine %d is running\n", id)
            done <- struct{}{} // Send a signal when the task is done
        }(i)
    }

    // Wait for all goroutines to finish
    for i := 0; i < numGoroutines; i++ {
        <-done // Receive signals
    }
    fmt.Println("All goroutines finished")
}

Output (order may vary):

Goroutine 1 is running
Goroutine 2 is running
Goroutine 3 is running
All goroutines finished

How it works:

Each goroutine sends a signal to the done channel upon completion.
The main goroutine confirms that all tasks are done by receiving the specified number of signals.

Advantages:

High flexibility, can carry data (such as task results).
Suitable for a dynamic number of goroutines.

Disadvantages:

Need to manually manage the number of receives, which can make the code a bit cumbersome.

Method 3: Controlling Exit with Context

Using context.Context allows you to gracefully control goroutine exits and have the main goroutine wait until all tasks are done. This method is especially useful in scenarios requiring cancellation or timeouts.

Example Code

package main

import (
    "context"
    "fmt"
    "sync"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    var wg sync.WaitGroup

    for i := 1; i <= 3; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            select {
            case <-ctx.Done():
                fmt.Printf("Goroutine %d cancelled\n", id)
                return
            default:
                fmt.Printf("Goroutine %d is running\n", id)
            }
        }(i)
    }

    // Simulate task completion
    cancel()       // Send cancel signal
    wg.Wait()      // Wait for all goroutines to exit
    fmt.Println("All goroutines finished")
}

Output (may vary depending on when cancellation occurs):

Goroutine 1 is running
Goroutine 2 is running
Goroutine 3 is running
All goroutines finished

How it works:

The context is used to notify goroutines to exit.
WaitGroup ensures that the main goroutine waits for all goroutines to complete.

Advantages:

Supports cancellation and timeout, suitable for complex concurrent scenarios.

Disadvantages:

Slightly more complex code.

Method 4: Using errgroup (Recommended)

golang.org/x/sync/errgroup is an advanced tool that combines the waiting functionality of WaitGroup with error handling, making it especially suitable for waiting for a group of tasks and handling errors.

Example Code

package main

import (
    "fmt"
    "golang.org/x/sync/errgroup"
)

func main() {
    var g errgroup.Group

    for i := 1; i <= 3; i++ {
        id := i
        g.Go(func() error {
            fmt.Printf("Goroutine %d is running\n", id)
            return nil // No error
        })
    }

    if err := g.Wait(); err != nil {
        fmt.Println("Error:", err)
    } else {
        fmt.Println("All goroutines finished")
    }
}

Output:

Goroutine 1 is running
Goroutine 2 is running
Goroutine 3 is running
All goroutines finished

How it works:

g.Go() starts a goroutine and adds it to the group.
g.Wait() waits for all goroutines to finish and returns the first non-nil error (if any).

Advantages:

Simple and elegant, supports error propagation.
Built-in context support (can use errgroup.WithContext).

Installation:

Requires go get golang.org/x/sync/errgroup.

Which Method to Choose?

sync.WaitGroup

Applicable scenarios: Simple tasks with a fixed number.
Advantages: Simple and efficient.
Disadvantages: Does not support error handling or cancellation.

Channel

Applicable scenarios: Dynamic tasks or when results need to be passed.
Advantages: Highly flexible.
Disadvantages: Manual management is more complex.

context

Applicable scenarios: Complex situations where cancellation or timeout is required.
Advantages: Supports cancellation and timeout.
Disadvantages: Code is slightly more complex.

errgroup

Applicable scenarios: Modern applications that require error handling and waiting.
Advantages: Elegant and powerful.
Disadvantages: Requires extra dependency.

Others: Why Doesn’t the Main Goroutine Just Sleep?

time.Sleep only introduces a fixed delay and cannot accurately wait for tasks to finish. This may cause the program to exit prematurely or lead to unnecessary waiting. Synchronization tools are more reliable.

Summary

The most commonly used method for the main goroutine to wait for other goroutines is sync.WaitGroup, which is simple and efficient. If you need error handling or cancellation capabilities, errgroup or a combination with context is recommended. Choose the appropriate tool according to your specific requirements to ensure clear program logic and prevent resource leaks.