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

     1  package syncx
     2  
     3  import (
     4  	"errors"
     5  
     6  	"github.com/lingyao2333/mo-zero/core/lang"
     7  )
     8  
     9  // ErrLimitReturn indicates that the more than borrowed elements were returned.
    10  var ErrLimitReturn = errors.New("discarding limited token, resource pool is full, someone returned multiple times")
    11  
    12  // Limit controls the concurrent requests.
    13  type Limit struct {
    14  	pool chan lang.PlaceholderType
    15  }
    16  
    17  // NewLimit creates a Limit that can borrow n elements from it concurrently.
    18  func NewLimit(n int) Limit {
    19  	return Limit{
    20  		pool: make(chan lang.PlaceholderType, n),
    21  	}
    22  }
    23  
    24  // Borrow borrows an element from Limit in blocking mode.
    25  func (l Limit) Borrow() {
    26  	l.pool <- lang.Placeholder
    27  }
    28  
    29  // Return returns the borrowed resource, returns error only if returned more than borrowed.
    30  func (l Limit) Return() error {
    31  	select {
    32  	case <-l.pool:
    33  		return nil
    34  	default:
    35  		return ErrLimitReturn
    36  	}
    37  }
    38  
    39  // TryBorrow tries to borrow an element from Limit, in non-blocking mode.
    40  // If success, true returned, false for otherwise.
    41  func (l Limit) TryBorrow() bool {
    42  	select {
    43  	case l.pool <- lang.Placeholder:
    44  		return true
    45  	default:
    46  		return false
    47  	}
    48  }