golang中如何保证main退出时，让goroutine处理逻辑执行全部完成？

Question

golang中一般使用context.Context来协调父子调用关系，但是如果ctx.Done() 的处理逻辑处理的比较慢，那么有没有办法保证在main结束以前让goroutine的逻辑全部执行完成？

Answer 1

通过 close channel + sync.WaitGroup 的方式通知。

close channel 的原理是当 channel 被 close 的时候，会向所有 channel 的接收方广播消息。

sync.WaitGroup 是为了保证在 goroutine 在接收到退出消息，主 goroutine 等待其他 goroutine 处理完成后再退出。

一个例子如下：

func sayHelloLoop(done <-chan struct{}) {  
   for {  
      select {  
      case <-done:  
         fmt.Println("done")  
         time.Sleep(3 * time.Second)  
         return  
 default:  
         // done someting  
 }  
   }  
}  
  
func main() {  
   done := make(chan struct{})  
   var wg sync.WaitGroup  
  
 wg.Add(1)  
   go func() {  
      defer wg.Done()  
      sayHelloLoop(done)  
   }()  
  
   wg.Add(1)  
   go func() {  
      defer wg.Done()  
      sayHelloLoop(done)  
   }()  
  
   close(done)  
   wg.Wait()  
}

其中的 close(done) 负责通知两个 sayHelloLoop 退出循环，而 wg.Wait() 负责等待 goroutine 退出。

注：代码纯手写，有错勿怪。?

Answer 2

如果你能保证不泄漏 goroutine，那通过在 main 中判断剩余 goroutine 个数，可作为退出参考。

package main

import (
    "runtime"
    "time"
)

func foo() {
    time.Sleep(time.Second * 10)
}

func main() {
    go foo()
    go foo()

    for {
        if runtime.NumGoroutine() == 1 {
            break
        }
        time.Sleep(time.Second)
    }
}

Answer 3

用一个（或几个） channel 让 goroutine 通知 main 自己已经结束。

Answer 4

sync.WaitGroup

Answer 5

总结一下，共有两种方案：
一种是sync.WaitGroup，另一种是手工判断goroutine的数量是否等于1。不知道还有没有更好的方案的？

以下为 sync.WaitGroup + context.Context来实现的

package main

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

func main() {
    var wg sync.WaitGroup

    ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    for i := 0; i < 15; i++ {
        wg.Add(1)
        go db(ctx, &wg, i)
    }

    // 放弃执行
    cancel()

    // 等待所有goroutine 都处理完成
    wg.Wait()

    fmt.Println("大家都干完了！")
}

func db(ctx context.Context, wg *sync.WaitGroup, id int) {
    defer wg.Done()
    for {
        select {
        case <-ctx.Done():
            fmt.Println(ctx.Err())
            fmt.Printf("收到exit信号, 清理工作开始:%d\n", id)
            time.Sleep(time.Second * 2)
            fmt.Printf("清理工作结束:%d\n", id)
            return
        default:
            fmt.Println("a")
        }
    }
}

每次都要传递两个参数wg和ctx

Answer 6

还可以接收外部信号被动的退出

func main() {  
   var wg sync.WaitGroup  
   wg.Add(1)  
   go func() {  
      defer wg.Done()  
      ch := make(chan os.Signal)  
      signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)  
      <-ch  
      fmt.Println("over")  
      return  
   }()  
   wg.Wait()  
}