If no tags specified, then limiter could be optimized to use atomic semaphore. Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
39 lines
668 B
Go
39 lines
668 B
Go
package qos
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync/atomic"
|
|
|
|
"git.frostfs.info/TrueCloudLab/frostfs-qos/scheduling"
|
|
)
|
|
|
|
var (
|
|
_ scheduler = (*semaphore)(nil)
|
|
errSemaphoreLimitExceeded = errors.New("semaphore limit exceeded")
|
|
)
|
|
|
|
type semaphore struct {
|
|
count atomic.Int64
|
|
limit int64
|
|
}
|
|
|
|
func (s *semaphore) Close() {}
|
|
|
|
func (s *semaphore) RequestArrival(ctx context.Context, _ string) (scheduling.ReleaseFunc, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
default:
|
|
}
|
|
|
|
v := s.count.Add(1)
|
|
if v > s.limit {
|
|
s.count.Add(-1)
|
|
return nil, errSemaphoreLimitExceeded
|
|
}
|
|
|
|
return func() {
|
|
s.count.Add(-1)
|
|
}, nil
|
|
}
|