frostfs-qos/limiting/semaphore.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)
}

// acquireWithReturnLimit calls [acquire] and returns current limit if failed.
// This is used for test purposes only.
//
//nolint:unused
func (s *burstAtomicSemaphore) acquireWithReturnLimit() (bool, int64) {
	v := s.count.Add(1)
	if v > s.limit {
		s.count.Add(-1)
		return false, v - 1
	}
	return true, v
}

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
}