Attention when deferring sync.WaitGroup Wait()

In Go it's very common to see the pattern that after creating a sync.WaitGroup() the Wait() call is immediately defer-red, like this:

func f() {
  var wg sync.WaitGroup
  defer wg.Wait()

  // rest of the logic
}

We need to be careful with this pattern due to how and when Go captures return values. For example, consider the following function that does some concurrent processing and using a sync.WaitGroup before returning the processed data to the caller:

func f(xs []int) map[int]string {
  var mu sync.Mutex
  var wg sync.WaitGroup
  defer wg.Wait()

  results := make(map[int]string, len(xs))

  for _, x := range xs {
    wg.Add(1)
    go func() {
      defer wg.Done()
      mu.Lock()
      defer mu.Unlock

      // some calculations
      results[x] = calc(x)
    }()
  }

  return results
}

This works fine, but consider the same using a string slice as return value, like this:

func f(xs []int) []string {
  var mu sync.Mutex
  var wg sync.WaitGroup
  defer wg.Wait()

  results := make([]string, 0, len(xs))

  for _, x := range xs {
    wg.Add(1)
    go func() {
      defer wg.Done()
      mu.Lock()
      defer mu.Unlock

      // some calculations
      results = append(results, calc(x))
    }()
  }

  return results
}

This will lead to a race condition. The problem is that the function f will (more or less) return immediately and capture the pointer to results, while the goroutines make an assignment to results.

You can avoid this problem in multiple ways:

1. do not defer the Wait() call, but call it manually right before the return.
2. do not assign to the return variable in the goroutines. Instead use the index assignment, like results[i] = calc(x).
3. use a named return value.

(1): manual calling Wait

Manually deferring the Wait() is very easy, to code example above would change to this:

func f(xs []int) []string {
  var mu sync.Mutex
  var wg sync.WaitGroup

  results := make([]string, 0, len(xs))

  for _, x := range xs {
    wg.Add(1)
    go func() {
      defer wg.Done()
      mu.Lock()
      defer mu.Unlock

      // some calculations
      results = append(results, calc(x))
    }()
  }

  wg.Wait()

  return results
}

(2): do not assign to the return variable

Since the underlying problem illustrated here is about the assignment to the captured return variable, we may just not re-assign it in goroutines:

func f(xs []int) []string {
  var mu sync.Mutex
  var wg sync.WaitGroup
  defer wg.Wait()

  results := make([]string, 0, len(xs))

  for i, x := range xs {
    wg.Add(1)
    go func() {
      defer wg.Done()
      mu.Lock()
      defer mu.Unlock

      // some calculations
      results[i] = calc(x)
    }()
  }

  return results
}

In this example, this may work because we know the length of the returned slice ahead of time, but that may not always be possible.

(3): use a named return

We can also use a named return which changes when Go captures the return value. The following works just fine:

func f(xs []int) (results []string) {
  var mu sync.Mutex
  var wg sync.WaitGroup
  defer wg.Wait()

  for _, x := range xs {
    wg.Add(1)
    go func() {
      defer wg.Done()
      mu.Lock()
      defer mu.Unlock

      // some calculations
      results = append(results, calc(x))
    }()
  }

  return results
}

Conclusion

Pay attention when using sync.WaitGroup and deferring Wait() calls. Make sure your return values or whatever is awaited with the wait group doesn't depend on simple variable assignments.