74 lines
1.2 KiB
Go
74 lines
1.2 KiB
Go
package limiting
|
|
|
|
import (
|
|
"sync/atomic"
|
|
)
|
|
|
|
type atomicSemaphore struct {
|
|
countDown atomic.Int64
|
|
}
|
|
|
|
func newAtomicSemaphore(size int64) *atomicSemaphore {
|
|
sem := new(atomicSemaphore)
|
|
sem.countDown.Store(size)
|
|
return sem
|
|
}
|
|
|
|
func (s *atomicSemaphore) acquire() bool {
|
|
for {
|
|
v := s.countDown.Load()
|
|
if v == 0 {
|
|
return false
|
|
}
|
|
if s.countDown.CompareAndSwap(v, v-1) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *atomicSemaphore) release() {
|
|
s.countDown.Add(1)
|
|
}
|
|
|
|
type burstAtomicSemaphore struct {
|
|
count atomic.Int64
|
|
limit int64
|
|
}
|
|
|
|
func newBurstAtomicSemaphore(size int64) *burstAtomicSemaphore {
|
|
return &burstAtomicSemaphore{limit: size}
|
|
}
|
|
|
|
func (s *burstAtomicSemaphore) acquire() bool {
|
|
v := s.count.Add(1)
|
|
if v > s.limit {
|
|
s.count.Add(-1)
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (s *burstAtomicSemaphore) release() {
|
|
s.count.Add(-1)
|
|
}
|
|
|
|
type channelSemaphore struct {
|
|
ch chan struct{}
|
|
}
|
|
|
|
func newChannelSemaphore(size int64) *channelSemaphore {
|
|
return &channelSemaphore{make(chan struct{}, size)}
|
|
}
|
|
|
|
func (s *channelSemaphore) acquire() bool {
|
|
select {
|
|
case s.ch <- struct{}{}:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (s *channelSemaphore) release() {
|
|
<-s.ch
|
|
}
|