实现带timeout的sync.Wait函数,在函数体内实现。
// WaitTimeout() 如果timeout到了超时时间返回true,如果WaitGroup自然结束返回false
func WaitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
ctx, cancel := context.WithCancel(context.Background())
quit := make(chan bool)
go func() {
wg.Wait()
quit <- false
}()
go func() {
for {
select {
case <-ctx.Done():
return
case <-time.After(timeout):
quit <- true
}
}
}()
ret := <-quit
cancel()
return ret
}
func main() {
wg := sync.WaitGroup{}
c := make(chan struct{})
for i := 0; i < 10; i++ {
wg.Add(1)
go func(num int, close <-chan struct{}) {
defer wg.Done()
ret := <-close
fmt.Println(num, ret)
}(i, c)
}
//for i := 0; i < 10; i++ {
// c<-struct{}{}
//}
if WaitTimeout(&wg, time.Second*2) {
close(c)
fmt.Println("timeout exit")
}
time.Sleep(time.Second * 5)
}