diff --git a/scheduling/mclock.go b/scheduling/mclock.go index f9bf2d2..64c62a8 100644 --- a/scheduling/mclock.go +++ b/scheduling/mclock.go @@ -22,6 +22,7 @@ var ( ErrMClockSchedulerUnknownTag = errors.New("unknown tag") ErrInvalidTagInfo = errors.New("invalid tag info: shares, limit and reservation must be greater than zero") ErrInvalidRunLimit = errors.New("invalid run limit: must be greater than zero") + ErrTagRequestsProhibited = errors.New("tag requests are prohibited") ) type request struct { @@ -49,6 +50,7 @@ type TagInfo struct { ReservedIOPS *float64 LimitIOPS *float64 Share float64 + Prohibited bool } // MClock is mClock scheduling algorithm implementation. @@ -196,6 +198,9 @@ func (q *MClock) pushRequest(tag string) (*request, ReleaseFunc, error) { if !ok { return nil, nil, ErrMClockSchedulerUnknownTag } + if tagInfo.Prohibited { + return nil, nil, ErrTagRequestsProhibited + } prev, ok := q.previous[tag] assert.Cond(ok, "undefined previous:", tag) diff --git a/scheduling/mclock_test.go b/scheduling/mclock_test.go index 81e1811..6433990 100644 --- a/scheduling/mclock_test.go +++ b/scheduling/mclock_test.go @@ -563,3 +563,22 @@ func TestMClockLimitTotalTime(t *testing.T) { require.True(t, time.Since(startedAt) >= 1500*time.Millisecond) require.True(t, time.Since(startedAt) <= 1600*time.Millisecond) // 100 ms offset to complete all requests } + +func TestMClockRestictTagRequests(t *testing.T) { + t.Parallel() + limit := 10.0 + q, err := NewMClock(100, 100, map[string]TagInfo{ + "class1": {Share: 50, LimitIOPS: &limit}, + "class2": {Share: 50, LimitIOPS: &limit, Prohibited: true}, + }, 5*time.Second) + require.NoError(t, err) + defer q.Close() + + release, err := q.RequestArrival(context.Background(), "class1") + require.NoError(t, err) + release() + + release, err = q.RequestArrival(context.Background(), "class2") + require.ErrorIs(t, err, ErrTagRequestsProhibited) + require.Nil(t, release) +}