limit.go 1.1 KB

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