github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/misc.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/zhiqiangxu/util/logger"
     9  	"go.uber.org/zap"
    10  )
    11  
    12  // GoFunc runs a goroutine under WaitGroup
    13  func GoFunc(routinesGroup *sync.WaitGroup, f func()) {
    14  	routinesGroup.Add(1)
    15  	go func() {
    16  		defer routinesGroup.Done()
    17  		f()
    18  	}()
    19  }
    20  
    21  // TryUntilSuccess will try f until success
    22  func TryUntilSuccess(f func() bool, duration time.Duration) {
    23  	var r bool
    24  	for {
    25  		r = f()
    26  		if r {
    27  			return
    28  		}
    29  		time.Sleep(duration)
    30  	}
    31  }
    32  
    33  // RunWithRecovery for run exec, calls recoverFn when panic
    34  func RunWithRecovery(exec func(), recoverFn func(r interface{})) {
    35  	defer func() {
    36  		r := recover()
    37  		if r != nil {
    38  			logger.Instance().Error("panic in the recoverable goroutine",
    39  				zap.Reflect("r", r),
    40  				zap.Stack("stack trace"))
    41  
    42  			if recoverFn != nil {
    43  				recoverFn(r)
    44  			}
    45  		}
    46  	}()
    47  	exec()
    48  }
    49  
    50  // RunWithRetry will run the f with backoff and retry.
    51  // retryCnt: Max retry count
    52  // backoff: When run f failed, it will sleep backoff * triedCount time.Millisecond.
    53  // Function f should have two return value. The first one is an bool which indicate if the err if retryable.
    54  // The second is if the f meet any error.
    55  func RunWithRetry(retryCnt int, backoff uint64, f func() (bool, error)) (err error) {
    56  	var retryAble bool
    57  	for i := 1; i <= retryCnt; i++ {
    58  		retryAble, err = f()
    59  		if err == nil || !retryAble {
    60  			return
    61  		}
    62  		sleepTime := time.Duration(backoff*uint64(i)) * time.Millisecond
    63  		time.Sleep(sleepTime)
    64  	}
    65  	return
    66  }
    67  
    68  // RunWithCancel for run a job with cancel-ability
    69  func RunWithCancel(ctx context.Context, f, cancel func()) {
    70  	var wg sync.WaitGroup
    71  
    72  	doneCh := make(chan struct{})
    73  	GoFunc(&wg, func() {
    74  		f()
    75  		close(doneCh)
    76  	})
    77  
    78  	select {
    79  	case <-ctx.Done():
    80  		cancel()
    81  	case <-doneCh:
    82  	}
    83  }