github.com/lingyao2333/mo-zero@v1.4.1/core/syncx/timeoutlimit.go (about)

     1  package syncx
     2  
     3  import (
     4  	"errors"
     5  	"time"
     6  )
     7  
     8  // ErrTimeout is an error that indicates the borrow timeout.
     9  var ErrTimeout = errors.New("borrow timeout")
    10  
    11  // A TimeoutLimit is used to borrow with timeouts.
    12  type TimeoutLimit struct {
    13  	limit Limit
    14  	cond  *Cond
    15  }
    16  
    17  // NewTimeoutLimit returns a TimeoutLimit.
    18  func NewTimeoutLimit(n int) TimeoutLimit {
    19  	return TimeoutLimit{
    20  		limit: NewLimit(n),
    21  		cond:  NewCond(),
    22  	}
    23  }
    24  
    25  // Borrow borrows with given timeout.
    26  func (l TimeoutLimit) Borrow(timeout time.Duration) error {
    27  	if l.TryBorrow() {
    28  		return nil
    29  	}
    30  
    31  	var ok bool
    32  	for {
    33  		timeout, ok = l.cond.WaitWithTimeout(timeout)
    34  		if ok && l.TryBorrow() {
    35  			return nil
    36  		}
    37  
    38  		if timeout <= 0 {
    39  			return ErrTimeout
    40  		}
    41  	}
    42  }
    43  
    44  // Return returns a borrow.
    45  func (l TimeoutLimit) Return() error {
    46  	if err := l.limit.Return(); err != nil {
    47  		return err
    48  	}
    49  
    50  	l.cond.Signal()
    51  	return nil
    52  }
    53  
    54  // TryBorrow tries a borrow.
    55  func (l TimeoutLimit) TryBorrow() bool {
    56  	return l.limit.TryBorrow()
    57  }