Manager for quota

Signed-off-by: He Weiwei <hweiwei@vmware.com>
This commit is contained in:
He Weiwei 2019-07-09 13:59:48 +08:00
parent 5ec6e39a65
commit 41ba410bb2
10 changed files with 804 additions and 8 deletions

View File

@ -187,3 +187,29 @@ func Escape(str string) string {
str = strings.Replace(str, `_`, `\_`, -1)
return str
}
// WithTransaction helper for transaction
func WithTransaction(handler func(o orm.Ormer) error) error {
o := orm.NewOrm()
if err := o.Begin(); err != nil {
log.Errorf("begin transaction failed: %v", err)
return err
}
if err := handler(o); err != nil {
if e := o.Rollback(); e != nil {
log.Errorf("rollback transaction failed: %v", e)
return e
}
return err
}
if err := o.Commit(); err != nil {
log.Errorf("commit transaction failed: %v", err)
return err
}
return nil
}

View File

@ -1035,3 +1035,51 @@ func TestIsDupRecError(t *testing.T) {
assert.True(t, isDupRecErr(fmt.Errorf("pq: duplicate key value violates unique constraint \"properties_k_key\"")))
assert.False(t, isDupRecErr(fmt.Errorf("other error")))
}
func TestWithTransaction(t *testing.T) {
quota := models.Quota{
Reference: "project",
ReferenceID: "1",
Hard: "{}",
}
failed := func(o orm.Ormer) error {
o.Insert(&quota)
return fmt.Errorf("failed")
}
var quotaID int64
success := func(o orm.Ormer) error {
id, err := o.Insert(&quota)
if err != nil {
return err
}
quotaID = id
return nil
}
assert := assert.New(t)
if assert.Error(WithTransaction(failed)) {
var quota models.Quota
quota.Reference = "project"
quota.ReferenceID = "1"
err := GetOrmer().Read(&quota, "reference", "reference_id")
assert.Error(err)
assert.False(quota.ID != 0)
}
if assert.Nil(WithTransaction(success)) {
var quota models.Quota
quota.Reference = "project"
quota.ReferenceID = "1"
err := GetOrmer().Read(&quota, "reference", "reference_id")
assert.Nil(err)
assert.True(quota.ID != 0)
assert.Equal(quotaID, quota.ID)
GetOrmer().Delete(&models.Quota{ID: quotaID}, "id")
}
}

View File

@ -40,7 +40,7 @@ func (h QuotaHard) Copy() QuotaHard {
// Quota model for quota
type Quota struct {
ID int64 `orm:"pk;auto;column(id)" json:"id"`
Reference string `orm:"column(reference)" json:"reference"`
Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota, eg: project, user
ReferenceID string `orm:"column(reference_id)" json:"reference_id"`
Hard string `orm:"column(hard);type(jsonb)" json:"-"`
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`

View File

@ -40,7 +40,7 @@ func (u QuotaUsed) Copy() QuotaUsed {
// QuotaUsage model for quota usage
type QuotaUsage struct {
ID int64 `orm:"pk;auto;column(id)" json:"id"`
Reference string `orm:"column(reference)" json:"reference"`
Reference string `orm:"column(reference)" json:"reference"` // The reference type for quota usage, eg: project, user
ReferenceID string `orm:"column(reference_id)" json:"reference_id"`
Used string `orm:"column(used);type(jsonb)" json:"-"`
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`

162
src/common/quota/manager.go Normal file
View File

@ -0,0 +1,162 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"time"
"github.com/astaxie/beego/orm"
"github.com/goharbor/harbor/src/common/dao"
"github.com/goharbor/harbor/src/common/models"
)
// Manager manager for quota
type Manager struct {
reference string
referenceID string
}
func (m *Manager) addQuota(o orm.Ormer, hardLimits ResourceList, now time.Time) (int64, error) {
quota := &models.Quota{
Reference: m.reference,
ReferenceID: m.referenceID,
Hard: hardLimits.String(),
CreationTime: now,
UpdateTime: now,
}
return o.Insert(quota)
}
func (m *Manager) addUsage(o orm.Ormer, hardLimits, used ResourceList, now time.Time, ids ...int64) (int64, error) {
usage := &models.QuotaUsage{
Reference: m.reference,
ReferenceID: m.referenceID,
Used: used.String(),
CreationTime: now,
UpdateTime: now,
}
if len(ids) > 0 {
usage.ID = ids[0]
}
return o.Insert(usage)
}
func (m *Manager) getQuotaForUpdate(o orm.Ormer) (*models.Quota, error) {
quota := &models.Quota{Reference: m.reference, ReferenceID: m.referenceID}
if err := o.ReadForUpdate(quota, "reference", "reference_id"); err != nil {
return nil, err
}
return quota, nil
}
func (m *Manager) getUsageForUpdate(o orm.Ormer) (*models.QuotaUsage, error) {
usage := &models.QuotaUsage{Reference: m.reference, ReferenceID: m.referenceID}
if err := o.ReadForUpdate(usage, "reference", "reference_id"); err != nil {
return nil, err
}
return usage, nil
}
func (m *Manager) updateUsage(o orm.Ormer, resources ResourceList, calculate func(ResourceList, ResourceList) ResourceList) error {
quota, err := m.getQuotaForUpdate(o)
if err != nil {
return err
}
hardLimits, err := NewResourceList(quota.Hard)
if err != nil {
return err
}
usage, err := m.getUsageForUpdate(o)
if err != nil {
return err
}
used, err := NewResourceList(usage.Used)
if err != nil {
return err
}
newUsed := calculate(used, resources)
if err := isSafe(hardLimits, newUsed); err != nil {
return err
}
usage.Used = newUsed.String()
usage.UpdateTime = time.Now()
_, err = o.Update(usage)
return err
}
// NewQuota create new quota for (reference, reference id)
func (m *Manager) NewQuota(hardLimits ResourceList, usages ...ResourceList) (int64, error) {
var quotaID int64
err := dao.WithTransaction(func(o orm.Ormer) error {
now := time.Now()
var err error
quotaID, err = m.addQuota(o, hardLimits, now)
if err != nil {
return err
}
var used ResourceList
if len(usages) > 0 {
used = usages[0]
} else {
used = ResourceList{}
}
if _, err := m.addUsage(o, hardLimits, used, now, quotaID); err != nil {
return err
}
return nil
})
if err != nil {
return 0, err
}
return quotaID, nil
}
// AddResources add resources to usage
func (m *Manager) AddResources(resources ResourceList) error {
return dao.WithTransaction(func(o orm.Ormer) error {
return m.updateUsage(o, resources, Add)
})
}
// SubtractResources subtract resources from usage
func (m *Manager) SubtractResources(resources ResourceList) error {
return dao.WithTransaction(func(o orm.Ormer) error {
return m.updateUsage(o, resources, Subtract)
})
}
// NewManager returns quota manager
func NewManager(reference string, referenceID string) (*Manager, error) {
return &Manager{
reference: reference,
referenceID: referenceID,
}, nil
}

View File

@ -0,0 +1,226 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"os"
"sync"
"testing"
"github.com/goharbor/harbor/src/common/dao"
"github.com/stretchr/testify/suite"
)
var (
hardLimits = ResourceList{ResourceStorage: 1000}
referenceProject = "project"
)
func mustResourceList(s string) ResourceList {
resources, _ := NewResourceList(s)
return resources
}
type ManagerSuite struct {
suite.Suite
}
func (suite *ManagerSuite) quotaManager(referenceIDs ...string) *Manager {
referenceID := "1"
if len(referenceIDs) > 0 {
referenceID = referenceIDs[0]
}
mgr, _ := NewManager(referenceProject, referenceID)
return mgr
}
func (suite *ManagerSuite) TearDownTest() {
dao.ClearTable("quota")
dao.ClearTable("quota_usage")
}
func (suite *ManagerSuite) TestNewQuota() {
mgr := suite.quotaManager()
if id, err := mgr.NewQuota(hardLimits); suite.Nil(err) {
quota, _ := dao.GetQuota(id)
suite.True(Equals(hardLimits, mustResourceList(quota.Hard)))
}
mgr = suite.quotaManager("2")
used := ResourceList{ResourceStorage: 100}
if id, err := mgr.NewQuota(hardLimits, used); suite.Nil(err) {
quota, _ := dao.GetQuota(id)
suite.True(Equals(hardLimits, mustResourceList(quota.Hard)))
usage, _ := dao.GetQuotaUsage(id)
suite.True(Equals(used, mustResourceList(usage.Used)))
}
}
func (suite *ManagerSuite) TestAddResources() {
mgr := suite.quotaManager()
id, _ := mgr.NewQuota(hardLimits)
resource := ResourceList{ResourceStorage: 100}
if suite.Nil(mgr.AddResources(resource)) {
usage, _ := dao.GetQuotaUsage(id)
suite.True(Equals(resource, mustResourceList(usage.Used)))
}
if suite.Nil(mgr.AddResources(resource)) {
usage, _ := dao.GetQuotaUsage(id)
suite.True(Equals(ResourceList{ResourceStorage: 200}, mustResourceList(usage.Used)))
}
if err := mgr.AddResources(ResourceList{ResourceStorage: 10000}); suite.Error(err) {
suite.True(IsUnsafeError(err))
}
}
func (suite *ManagerSuite) TestSubtractResources() {
mgr := suite.quotaManager()
id, _ := mgr.NewQuota(hardLimits)
resource := ResourceList{ResourceStorage: 100}
if suite.Nil(mgr.AddResources(resource)) {
usage, _ := dao.GetQuotaUsage(id)
suite.True(Equals(resource, mustResourceList(usage.Used)))
}
if suite.Nil(mgr.SubtractResources(resource)) {
usage, _ := dao.GetQuotaUsage(id)
suite.True(Equals(ResourceList{ResourceStorage: 0}, mustResourceList(usage.Used)))
}
}
func (suite *ManagerSuite) TestRaceAddResources() {
mgr := suite.quotaManager()
mgr.NewQuota(hardLimits)
resources := ResourceList{
ResourceStorage: 100,
}
var wg sync.WaitGroup
results := make([]bool, 100)
for i := 0; i < 100; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
results[i] = mgr.AddResources(resources) == nil
}(i)
}
wg.Wait()
var success int
for _, result := range results {
if result {
success++
}
}
suite.Equal(10, success)
}
func (suite *ManagerSuite) TestRaceSubtractResources() {
mgr := suite.quotaManager()
mgr.NewQuota(hardLimits, hardLimits)
resources := ResourceList{
ResourceStorage: 100,
}
var wg sync.WaitGroup
results := make([]bool, 100)
for i := 0; i < 100; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
results[i] = mgr.SubtractResources(resources) == nil
}(i)
}
wg.Wait()
var success int
for _, result := range results {
if result {
success++
}
}
suite.Equal(10, success)
}
func TestMain(m *testing.M) {
dao.PrepareTestForPostgresSQL()
if result := m.Run(); result != 0 {
os.Exit(result)
}
}
func TestRunManagerSuite(t *testing.T) {
suite.Run(t, new(ManagerSuite))
}
func BenchmarkAddResources(b *testing.B) {
defer func() {
dao.ClearTable("quota")
dao.ClearTable("quota_usage")
}()
mgr, _ := NewManager(referenceProject, "1")
mgr.NewQuota(ResourceList{ResourceStorage: int64(b.N)})
resource := ResourceList{
ResourceStorage: 1,
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
mgr.AddResources(resource)
}
b.StopTimer()
}
func BenchmarkAddResourcesParallel(b *testing.B) {
defer func() {
dao.ClearTable("quota")
dao.ClearTable("quota_usage")
}()
mgr, _ := NewManager(referenceProject, "1")
mgr.NewQuota(ResourceList{ResourceStorage: -1})
resource := ResourceList{
ResourceStorage: 1,
}
b.ResetTimer()
b.RunParallel(func(b *testing.PB) {
for b.Next() {
mgr.AddResources(resource)
}
})
b.StopTimer()
}

View File

@ -0,0 +1,103 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"encoding/json"
)
const (
// ResourceCount count, in number
ResourceCount ResourceName = "count"
// ResourceStorage storage size, in bytes
ResourceStorage ResourceName = "storage"
)
// ResourceName is the name identifying various resources in a ResourceList.
type ResourceName string
// ResourceList is a set of (resource name, value) pairs.
type ResourceList map[ResourceName]int64
func (resources ResourceList) String() string {
bytes, _ := json.Marshal(resources)
return string(bytes)
}
// NewResourceList returns resource list from string
func NewResourceList(s string) (ResourceList, error) {
var resources ResourceList
if err := json.Unmarshal([]byte(s), &resources); err != nil {
return nil, err
}
return resources, nil
}
// Equals returns true if the two lists are equivalent
func Equals(a ResourceList, b ResourceList) bool {
if len(a) != len(b) {
return false
}
for key, value1 := range a {
value2, found := b[key]
if !found {
return false
}
if value1 != value2 {
return false
}
}
return true
}
// Add returns the result of a + b for each named resource
func Add(a ResourceList, b ResourceList) ResourceList {
result := ResourceList{}
for key, value := range a {
if other, found := b[key]; found {
value = value + other
}
result[key] = value
}
for key, value := range b {
if _, found := result[key]; !found {
result[key] = value
}
}
return result
}
// Subtract returns the result of a - b for each named resource
func Subtract(a ResourceList, b ResourceList) ResourceList {
result := ResourceList{}
for key, value := range a {
if other, found := b[key]; found {
value = value - other
}
result[key] = value
}
for key, value := range b {
if _, found := result[key]; !found {
result[key] = -value
}
}
return result
}

View File

@ -0,0 +1,72 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"testing"
"github.com/stretchr/testify/suite"
)
type ResourcesSuite struct {
suite.Suite
}
func (suite *ResourcesSuite) TestNewResourceList() {
res1, err1 := NewResourceList("")
suite.Error(err1)
suite.Nil(res1)
suite.Equal(0, len(res1))
res2, err2 := NewResourceList("{}")
suite.Nil(err2)
suite.NotNil(res2)
}
func (suite *ResourcesSuite) TestEquals() {
suite.True(Equals(ResourceList{}, ResourceList{}))
suite.True(Equals(ResourceList{ResourceStorage: 100}, ResourceList{ResourceStorage: 100}))
suite.False(Equals(ResourceList{ResourceStorage: 100}, ResourceList{ResourceStorage: 200}))
suite.False(Equals(ResourceList{ResourceStorage: 100}, ResourceList{ResourceStorage: 100, ResourceCount: 10}))
suite.False(Equals(ResourceList{ResourceStorage: 100, ResourceCount: 10}, ResourceList{ResourceStorage: 100}))
}
func (suite *ResourcesSuite) TestAdd() {
res1 := ResourceList{ResourceStorage: 100}
res2 := ResourceList{ResourceStorage: 100}
res3 := ResourceList{ResourceStorage: 100, ResourceCount: 10}
res4 := ResourceList{ResourceCount: 10}
suite.Equal(res1, Add(ResourceList{}, res1))
suite.Equal(ResourceList{ResourceStorage: 200}, Add(res1, res2))
suite.Equal(ResourceList{ResourceStorage: 200, ResourceCount: 10}, Add(res1, res3))
suite.Equal(ResourceList{ResourceStorage: 100, ResourceCount: 10}, Add(res1, res4))
}
func (suite *ResourcesSuite) TestSubtract() {
res1 := ResourceList{ResourceStorage: 100}
res2 := ResourceList{ResourceStorage: 100}
res3 := ResourceList{ResourceStorage: 100, ResourceCount: 10}
res4 := ResourceList{ResourceCount: 10}
suite.Equal(res1, Subtract(res1, ResourceList{}))
suite.Equal(ResourceList{ResourceStorage: 0}, Subtract(res1, res2))
suite.Equal(ResourceList{ResourceStorage: 0, ResourceCount: -10}, Subtract(res1, res3))
suite.Equal(ResourceList{ResourceStorage: 100, ResourceCount: -10}, Subtract(res1, res4))
}
func TestRunResourcesSuite(t *testing.T) {
suite.Run(t, new(ResourcesSuite))
}

65
src/common/quota/util.go Normal file
View File

@ -0,0 +1,65 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"fmt"
)
const (
// UNLIMITED unlimited quota value
UNLIMITED = int64(-1)
)
type unsafe struct {
message string
}
func (err *unsafe) Error() string {
return err.message
}
func newUnsafe(message string) error {
return &unsafe{message: message}
}
// IsUnsafeError returns true when the err is unsafe error
func IsUnsafeError(err error) bool {
_, ok := err.(*unsafe)
return ok
}
func isSafe(hardLimits ResourceList, used ResourceList) error {
for key, value := range used {
if value < 0 {
return newUnsafe(fmt.Sprintf("bad used value: %d", value))
}
if hard, found := hardLimits[key]; found {
if hard == UNLIMITED {
continue
}
if value > hard {
return newUnsafe(fmt.Sprintf("over the quota: used %d but only hard %d", value, hard))
}
} else {
return newUnsafe(fmt.Sprintf("hard limit not found: %s", key))
}
}
return nil
}

View File

@ -0,0 +1,94 @@
// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package quota
import (
"errors"
"testing"
)
func TestIsUnsafeError(t *testing.T) {
type args struct {
err error
}
tests := []struct {
name string
args args
want bool
}{
{
"is unsafe error",
args{err: newUnsafe("unsafe")},
true,
},
{
"is not unsafe error",
args{err: errors.New("unsafe")},
false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := IsUnsafeError(tt.args.err); got != tt.want {
t.Errorf("IsUnsafeError() = %v, want %v", got, tt.want)
}
})
}
}
func Test_checkQuotas(t *testing.T) {
type args struct {
hardLimits ResourceList
used ResourceList
}
tests := []struct {
name string
args args
wantErr bool
}{
{
"unlimited",
args{hardLimits: ResourceList{ResourceStorage: UNLIMITED}, used: ResourceList{ResourceStorage: 1000}},
false,
},
{
"ok",
args{hardLimits: ResourceList{ResourceStorage: 100}, used: ResourceList{ResourceStorage: 1}},
false,
},
{
"bad used value",
args{hardLimits: ResourceList{ResourceStorage: 100}, used: ResourceList{ResourceStorage: -1}},
true,
},
{
"over the hard limit",
args{hardLimits: ResourceList{ResourceStorage: 100}, used: ResourceList{ResourceStorage: 200}},
true,
},
{
"hard limit not found",
args{hardLimits: ResourceList{ResourceStorage: 100}, used: ResourceList{ResourceCount: 1}},
true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if err := isSafe(tt.args.hardLimits, tt.args.used); (err != nil) != tt.wantErr {
t.Errorf("isSafe() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}