wait.go 816 B

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. package wait
  2. import (
  3. "sync"
  4. "time"
  5. )
  6. // Wait is similar with sync.WaitGroup which can wait with timeout
  7. type Wait struct {
  8. wg sync.WaitGroup
  9. }
  10. // Add adds delta, which may be negative, to the WaitGroup counter.
  11. func (w *Wait) Add(delta int) {
  12. w.wg.Add(delta)
  13. }
  14. // Done decrements the WaitGroup counter by one
  15. func (w *Wait) Done() {
  16. w.wg.Done()
  17. }
  18. // Wait blocks until the WaitGroup counter is zero.
  19. func (w *Wait) Wait() {
  20. w.wg.Wait()
  21. }
  22. // WaitWithTimeout blocks until the WaitGroup counter is zero or timeout
  23. // returns true if timeout
  24. func (w *Wait) WaitWithTimeout(timeout time.Duration) bool {
  25. c := make(chan bool, 1)
  26. go func() {
  27. defer close(c)
  28. w.wg.Wait()
  29. c <- true
  30. }()
  31. select {
  32. case <-c:
  33. return false // completed normally
  34. case <-time.After(timeout):
  35. return true // timed out
  36. }
  37. }