github.com/viant/toolbox@v0.34.5/waitgroup_helper.go (about)

     1  package toolbox
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  )
     8  
     9  // WaitGroup that waits with a timeout
    10  // Returns true if timeout exceeded and false if there was no timeout
    11  func WaitTimeout(wg *sync.WaitGroup, duration time.Duration) bool {
    12  	done := make(chan bool, 1)
    13  	closed := int32(0)
    14  	defer func() {
    15  		if atomic.CompareAndSwapInt32(&closed, 0, 1) {
    16  			close(done)
    17  		}
    18  	}()
    19  	go func() {
    20  		wg.Wait()
    21  		if atomic.LoadInt32(&closed) == 0 {
    22  			done <- true
    23  		}
    24  	}()
    25  
    26  	select {
    27  	case <-done: //Wait till the task is complete and channel get unblocked
    28  		return false //No durationToken. Normal execution of task completion
    29  	case <-time.After(duration): //Wait till durationToken to elapse
    30  		//TODO: time.After() creates a timer that does not get GC until timer durationToken gets elapsed. Need to use AfterFunc
    31  		return true //Timed out
    32  	}
    33  }