github.com/sandwich-go/boost@v1.3.29/xsync/wg.go (about)

     1  package xsync
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/sandwich-go/boost/internal/log"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  // WaitGroupTimeout is a wrapper of sync.WaitGroup that
    12  // supports wait with timeout.
    13  type WaitGroupTimeout struct {
    14  	sync.WaitGroup
    15  }
    16  
    17  func (w *WaitGroupTimeout) WaitTimeout(timeout time.Duration) bool {
    18  	return WaitTimeout(&w.WaitGroup, timeout)
    19  }
    20  
    21  // WaitContext performs a timed wait on a given sync.WaitGroup
    22  func WaitContext(waitGroup *sync.WaitGroup, ctx context.Context) bool {
    23  	success := make(chan struct{})
    24  	go func() {
    25  		defer func() {
    26  			if reason := recover(); reason != nil {
    27  				log.Error(fmt.Sprintf("wait context panic, reason: %v", reason))
    28  			}
    29  		}()
    30  		defer close(success)
    31  		waitGroup.Wait()
    32  	}()
    33  	select {
    34  	case <-success: // completed normally
    35  		return false
    36  	case <-ctx.Done():
    37  		return true
    38  	}
    39  }
    40  
    41  // WaitTimeout performs a timed wait on a given sync.WaitGroup
    42  func WaitTimeout(waitGroup *sync.WaitGroup, timeout time.Duration) bool {
    43  	ctx, cancelFunc := context.WithDeadline(context.Background(), time.Now().Add(timeout))
    44  	defer cancelFunc()
    45  	return WaitContext(waitGroup, ctx)
    46  }