frostfs-qos/limiting/limiter.go

package limiting

import (
	"fmt"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting/semaphore"
)

type ReleaseFunc func()

type Limiter interface {
	// Acquire attempts to reserve a slot without blocking.
	//
	// Returns a release function and true if successful. The function must be
	// called to release the limiter. The function must be called exactly once.
	// Calling the function more that once will cause incorrect behavior of the
	// limiter.
	//
	// Returns nil and false if fails.
	//
	// If the key was not defined in the limiter, no limit is applied.
	Acquire(key string) (ReleaseFunc, bool)
}

type SemaphoreLimiter struct {
	m map[string]*semaphore.Semaphore
}

// KeyLimit defines a concurrency limit for a set of keys.
//
// All keys of one set share the same limit.
// Keys of different sets have separate limits.
//
// Sets must not overlap.
type KeyLimit struct {
	Keys  []string
	Limit int64
}

func NewSemaphoreLimiter(limits []KeyLimit) (*SemaphoreLimiter, error) {
	lr := SemaphoreLimiter{make(map[string]*semaphore.Semaphore)}
	for _, limit := range limits {
		if limit.Limit < 0 {
			return nil, fmt.Errorf("invalid limit %d", limit.Limit)
		}
		sem := semaphore.NewSemaphore(limit.Limit)

		if err := lr.addLimit(&limit, sem); err != nil {
			return nil, err
		}
	}
	return &lr, nil
}

func (lr *SemaphoreLimiter) addLimit(limit *KeyLimit, sem *semaphore.Semaphore) error {
	for _, key := range limit.Keys {
		if _, exists := lr.m[key]; exists {
			return fmt.Errorf("duplicate key %q", key)
		}
		lr.m[key] = sem
	}
	return nil
}

func (lr *SemaphoreLimiter) Acquire(key string) (ReleaseFunc, bool) {
	sem, ok := lr.m[key]
	if !ok {
		return func() {}, true
	}

	if ok := sem.Acquire(); ok {
		return sem.Release, true
	}
	return nil, false
}