frostfs-node/internal/qos/semaphore.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
}