sync.WaitGroup增加timeout

要等待多个goroutine完成，我们可以使用sync.WaitGroup，但是它会一直等待，这里简单封装下增加超时机制。

直接上代码：

type WaitGroupTimeout struct {
	wg sync.WaitGroup
	done chan struct{}
	timeout time.Duration
}

func NewWaitGroupTimeout(timeout time.Duration)*WaitGroupTimeout {
	return &WaitGroupTimeout{
		done: make(chan struct{}),
		timeout: timeout,
	}
}

func(tw *WaitGroupTimeout)Add(delta int) {
	tw.wg.Add(delta)
}

func(tw *WaitGroupTimeout)Done() {
	tw.wg.Done()
}

func(tw *WaitGroupTimeout)Wait() {
	go func() {
		tw.wg.Wait()
		close(tw.done)
	}()
	select {
	case <- tw.done:
	case <-time.After(tw.timeout):
	}
}

测试代码如下：

func main() {
	const TIMEOUT = 5 * time.Second
	const TASK_EXECUTE_TIME = 3 * time.Second

	tw := NewWaitGroupTimeout(TIMEOUT)
	go func() {
		tw.Add(1)
		defer tw.Done()
		fmt.Println("start")
		time.Sleep(TASK_EXECUTE_TIME)
		fmt.Println("end")
	}()
	fmt.Println("wait exit")
	tw.Wait()
	fmt.Println("exit")
}

当超时时间大于任务时间时，任务可以正常完成然后退出，
当超时时间小于任务时间时，任务没有执行完成就退出了。

sync.WaitGroup增加timeout

上一篇文章

下一篇文章