gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/pool/limiter.go (about)

     1  package pool
     2  
     3  import "sync/atomic"
     4  
     5  type limiter struct {
     6  	limit     int64
     7  	allocated int64
     8  }
     9  
    10  func (l *limiter) setLimit(limit int64) {
    11  	l.limit = limit
    12  }
    13  
    14  func (l *limiter) alloc() bool {
    15  	if l.limit > 0 {
    16  		if l.allocated >= l.limit {
    17  			return false
    18  		}
    19  		if l.allocated++; l.allocated <= 0 {
    20  			panic(NewInvalidReferenceError(l.allocated, l.allocated-1))
    21  		}
    22  	}
    23  	return true
    24  }
    25  
    26  func (l *limiter) release() {
    27  	if l.limit > 0 {
    28  		if l.allocated--; l.allocated < 0 {
    29  			panic(NewInvalidReferenceError(l.allocated, l.allocated+1))
    30  		}
    31  	}
    32  }
    33  
    34  type syncLimiter struct {
    35  	limit     int64
    36  	allocated atomic.Int64
    37  }
    38  
    39  func (l *syncLimiter) setLimit(limit int64) {
    40  	l.limit = limit
    41  }
    42  
    43  func (l *syncLimiter) alloc() bool {
    44  	if l.limit > 0 {
    45  		if allocated := l.allocated.Add(1); allocated > l.limit {
    46  			l.allocated.Add(-1)
    47  			return false
    48  		} else if allocated <= 0 {
    49  			panic(NewInvalidReferenceError(allocated, allocated-1))
    50  		}
    51  	}
    52  	return true
    53  }
    54  
    55  func (l *syncLimiter) release() {
    56  	if l.limit > 0 {
    57  		if allocated := l.allocated.Add(-1); allocated < 0 {
    58  			panic(NewInvalidReferenceError(allocated, allocated+1))
    59  		}
    60  	}
    61  }