我们如何确定Go中“最后一个”工作进程/线程何时完成？

Question

I'll use a hacky inefficient prime number finder to make this question a little more concrete.

Let's say our main function fires off a bunch of "worker" goroutines. They will report their results to a single channnel which prints them. But not every worker will report something so we can't use a counter to know when the last job is finished. Or is there a way?

For the concrete example, here, main fires off goroutines to check whether the values 2...1000 are prime (yeah I know it is inefficient).

package main

import (
    "fmt"
    "time"
)

func main() {
    c := make(chan int)
    go func () {
        for {
            fmt.Print(" ", <- c)
        }
    }()
    for n := 2; n < 1000; n++ {
        go printIfPrime(n, c)
    }
    time.Sleep(2 * time.Second)   // <---- THIS FEELS WRONG
}

func printIfPrime(n int, channel chan int) {
    for d := 2; d * d <= n; d++ {
        if n % d == 0 {
            return
        }
    }
    channel <- n
}

My problem is that I don't know how to reliably stop it at the right time. I tried adding a sleep at the end of main and it works (but it might take too long, and this is no way to write concurrent code!). I would like to know if there was a way to send a stop signal through a channel or something so main can stop at the right time.

The trick here is that I don't know how many worker responses there will be.

Is this impossible or is there a cool trick?

(If there's an answer for this prime example, great. I can probably generalize. Or maybe not. Maybe this is app specific?)

Answer 1

Use a WaitGroup.

The following code uses two WaitGroups. The main function uses wgTest to wait for print_if_prime functions to complete. Once they are done, it closes the channel to break the for loop in the printing goroutine. The main function uses wgPrint to wait for printing to complete.

package main

import (
  "fmt"
  "sync"
) 

func main() {
  c := make(chan int)
  var wgPrint, wgTest sync.WaitGroup

  wgPrint.Add(1)
  go func(wg *sync.WaitGroup) {
    defer wg.Done()
    for n := range c {
        fmt.Print(" ", n)
    }
  }(&wgPrint)

  for n := 2; n < 1000; n++ {
    wgTest.Add(1)
    go print_if_prime(&wgTest, n, c)
  }

  wgTest.Wait()
  close(c)
  wgPrint.Wait()
}

func print_if_prime(wg *sync.WaitGroup, n int, channel chan int) {
  defer wg.Done()
  for d := 2; d*d <= n; d++ {
    if n%d == 0 {
        return
    }
  }
  channel <- n
}

playground example