How to Set a Timeout for WaitGroup.Wait() in Go?

Timeout for WaitGroup.Wait()

WaitGroup.Wait() can block indefinitely, waiting for all goroutines to complete. This can be problematic when you want to safeguard your system from errant workers potentially blocking execution indefinitely. While there is no idiomatic way to set a timeout for WaitGroup.Wait(), there are several approaches to achieve this functionality.

One common approach involves using a channel and a goroutine. When the goroutine finishes its work, it sends a signal to the channel. The main program can select on the channel and a timer to determine if the goroutine timed out. Here's an example:

import (
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    wg.Add(1)

    timeout := time.After(5 * time.Second)
    ch := make(chan struct{})

    go func() {
        defer wg.Done()
        defer close(ch)
        // Do some work
    }()

    select {
    case <-ch:
        // Goroutine finished before timeout
        wg.Wait() // Wait for all goroutines to finish
    case <-timeout:
        // Goroutine timed out
        wg.Wait() // Block until all goroutines have completed
        log.Println("Timeout occurred")
    }
}

This solution uses a select statement to wait on either the channel or a timer, providing a timeout mechanism for WaitGroup.Wait(). Other recommended optimizations include using defer statements to close channels and using untyped constants for timeout durations.

Alternatively, you can use a custom helper function that wraps the above functionality, simplifying the code and making it more reusable. Here's an example:

func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
    c := make(chan struct{})
    defer close(c)

    go func() {
        wg.Wait()
        c <- struct{}{}
    }()

    select {
    case <-c:
        return false // Completed normally
    case <-time.After(timeout):
        return true // Timed out
    }
}

This function can be used as follows:

if waitTimeout(&wg, 5*time.Second) {
    fmt.Println("Timed out waiting for wait group")
} else {
    fmt.Println("Wait group finished")
}

These approaches provide a way to set a timeout for WaitGroup.Wait(), allowing you to safeguard your system from errant workers blocking execution indefinitely.

The above is the detailed content of How to Set a Timeout for WaitGroup.Wait() in Go?. For more information, please follow other related articles on the PHP Chinese website!