mirror of
https://github.com/goharbor/harbor.git
synced 2024-12-24 17:47:46 +01:00
Manager for quota
Signed-off-by: He Weiwei <hweiwei@vmware.com>
This commit is contained in:
parent
5ec6e39a65
commit
41ba410bb2
@ -187,3 +187,29 @@ func Escape(str string) string {
|
|||||||
str = strings.Replace(str, `_`, `\_`, -1)
|
str = strings.Replace(str, `_`, `\_`, -1)
|
||||||
return str
|
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
|
||||||
|
}
|
||||||
|
@ -47,8 +47,8 @@ func cleanByUser(username string) {
|
|||||||
o := GetOrmer()
|
o := GetOrmer()
|
||||||
o.Begin()
|
o.Begin()
|
||||||
|
|
||||||
err = execUpdate(o, `delete
|
err = execUpdate(o, `delete
|
||||||
from project_member
|
from project_member
|
||||||
where entity_id = (
|
where entity_id = (
|
||||||
select user_id
|
select user_id
|
||||||
from harbor_user
|
from harbor_user
|
||||||
@ -59,7 +59,7 @@ func cleanByUser(username string) {
|
|||||||
log.Error(err)
|
log.Error(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
err = execUpdate(o, `delete
|
err = execUpdate(o, `delete
|
||||||
from project_member
|
from project_member
|
||||||
where project_id = (
|
where project_id = (
|
||||||
select project_id
|
select project_id
|
||||||
@ -71,8 +71,8 @@ func cleanByUser(username string) {
|
|||||||
log.Error(err)
|
log.Error(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
err = execUpdate(o, `delete
|
err = execUpdate(o, `delete
|
||||||
from access_log
|
from access_log
|
||||||
where username = ?
|
where username = ?
|
||||||
`, username)
|
`, username)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@ -80,7 +80,7 @@ func cleanByUser(username string) {
|
|||||||
log.Error(err)
|
log.Error(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
err = execUpdate(o, `delete
|
err = execUpdate(o, `delete
|
||||||
from access_log
|
from access_log
|
||||||
where project_id = (
|
where project_id = (
|
||||||
select project_id
|
select project_id
|
||||||
@ -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.True(t, isDupRecErr(fmt.Errorf("pq: duplicate key value violates unique constraint \"properties_k_key\"")))
|
||||||
assert.False(t, isDupRecErr(fmt.Errorf("other error")))
|
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("a)
|
||||||
|
|
||||||
|
return fmt.Errorf("failed")
|
||||||
|
}
|
||||||
|
|
||||||
|
var quotaID int64
|
||||||
|
success := func(o orm.Ormer) error {
|
||||||
|
id, err := o.Insert("a)
|
||||||
|
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("a, "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("a, "reference", "reference_id")
|
||||||
|
assert.Nil(err)
|
||||||
|
assert.True(quota.ID != 0)
|
||||||
|
assert.Equal(quotaID, quota.ID)
|
||||||
|
|
||||||
|
GetOrmer().Delete(&models.Quota{ID: quotaID}, "id")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
@ -40,7 +40,7 @@ func (h QuotaHard) Copy() QuotaHard {
|
|||||||
// Quota model for quota
|
// Quota model for quota
|
||||||
type Quota struct {
|
type Quota struct {
|
||||||
ID int64 `orm:"pk;auto;column(id)" json:"id"`
|
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"`
|
ReferenceID string `orm:"column(reference_id)" json:"reference_id"`
|
||||||
Hard string `orm:"column(hard);type(jsonb)" json:"-"`
|
Hard string `orm:"column(hard);type(jsonb)" json:"-"`
|
||||||
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`
|
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`
|
||||||
|
@ -40,7 +40,7 @@ func (u QuotaUsed) Copy() QuotaUsed {
|
|||||||
// QuotaUsage model for quota usage
|
// QuotaUsage model for quota usage
|
||||||
type QuotaUsage struct {
|
type QuotaUsage struct {
|
||||||
ID int64 `orm:"pk;auto;column(id)" json:"id"`
|
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"`
|
ReferenceID string `orm:"column(reference_id)" json:"reference_id"`
|
||||||
Used string `orm:"column(used);type(jsonb)" json:"-"`
|
Used string `orm:"column(used);type(jsonb)" json:"-"`
|
||||||
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`
|
CreationTime time.Time `orm:"column(creation_time);auto_now_add" json:"creation_time"`
|
||||||
|
162
src/common/quota/manager.go
Normal file
162
src/common/quota/manager.go
Normal 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
|
||||||
|
}
|
226
src/common/quota/manager_test.go
Normal file
226
src/common/quota/manager_test.go
Normal 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()
|
||||||
|
}
|
103
src/common/quota/resources.go
Normal file
103
src/common/quota/resources.go
Normal 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
|
||||||
|
}
|
72
src/common/quota/resources_test.go
Normal file
72
src/common/quota/resources_test.go
Normal 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
65
src/common/quota/util.go
Normal 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
|
||||||
|
}
|
94
src/common/quota/util_test.go
Normal file
94
src/common/quota/util_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user