github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/util/sync.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // WaitGroup calls Wait() on a sync.WaitGroup and return once the Wait() completed
     9  // or the context is cancelled or times out, whatever occurs first. Returns the
    10  // specific context error if the context is cancelled or times out before Wait()
    11  // completes.
    12  func WaitGroup(ctx context.Context, wg *sync.WaitGroup) error {
    13  	c := make(chan struct{})
    14  
    15  	go func() {
    16  		defer close(c)
    17  		wg.Wait()
    18  	}()
    19  
    20  	select {
    21  	case <-c:
    22  		return nil
    23  	case <-ctx.Done():
    24  		return ctx.Err()
    25  	}
    26  }