-
Notifications
You must be signed in to change notification settings - Fork 0
/
waitgroup.go
36 lines (31 loc) · 815 Bytes
/
waitgroup.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.
package utils
import (
"errors"
"sync"
"sync/atomic"
"time"
)
// for keeping track of how many goroutines are waiting
var waiting int32
// WaitGroup is an extension of sync.WaitGroup that allows waiting with a maximum duration
type WaitGroup struct {
sync.WaitGroup
}
// WaitForMax waits until the WaitGroup is Done or the specified duration has elapsed
func (wg *WaitGroup) WaitForMax(d time.Duration) error {
waitChan := make(chan struct{})
go func() {
atomic.AddInt32(&waiting, 1)
wg.Wait()
atomic.AddInt32(&waiting, -1)
close(waitChan)
}()
select {
case <-waitChan:
return nil
case <-time.After(d):
return errors.New("Wait timeout expired")
}
}