golang 超时处理

60 阅读1分钟

实现带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)
}