From 9e0addee55c06f3e1e7eaf692730541d2496d9a3 Mon Sep 17 00:00:00 2001 From: wang yan Date: Tue, 13 Aug 2019 14:03:15 +0800 Subject: [PATCH] Enable usage sync when switch quota setting Signed-off-by: wang yan --- src/common/dao/project_blob.go | 19 +++++ src/common/dao/project_blob_test.go | 28 ++++++++ src/common/quota/manager.go | 49 ++++++++++++- src/common/quota/manager_test.go | 43 ++++++++++++ src/core/api/harborapi_test.go | 2 + src/core/api/internal.go | 104 +++++++++++++++++++++++++++- src/core/api/internal_test.go | 56 +++++++++++++++ src/core/router.go | 1 + 8 files changed, 299 insertions(+), 3 deletions(-) create mode 100644 src/core/api/internal_test.go diff --git a/src/common/dao/project_blob.go b/src/common/dao/project_blob.go index 7c6821b426..a13422278f 100644 --- a/src/common/dao/project_blob.go +++ b/src/common/dao/project_blob.go @@ -18,7 +18,9 @@ import ( "fmt" "time" + "github.com/astaxie/beego/orm" "github.com/goharbor/harbor/src/common/models" + "strconv" ) // AddBlobToProject ... @@ -103,3 +105,20 @@ func GetBlobsNotInProject(projectID int64, blobDigests ...string) ([]*models.Blo return blobs, nil } + +// CountSizeOfProject ... +func CountSizeOfProject(pid int64) (int64, error) { + var res []orm.Params + num, err := GetOrmer().Raw(`SELECT sum(bb.size) FROM project_blob pb LEFT JOIN blob bb ON pb.blob_id = bb.id WHERE pb.project_id = ? `, pid).Values(&res) + if err != nil { + return -1, err + } + if num > 0 { + size, err := strconv.ParseInt(res[0]["sum"].(string), 0, 64) + if err != nil { + return -1, err + } + return size, nil + } + return -1, err +} diff --git a/src/common/dao/project_blob_test.go b/src/common/dao/project_blob_test.go index 071bfdd3d0..3d3643aee3 100644 --- a/src/common/dao/project_blob_test.go +++ b/src/common/dao/project_blob_test.go @@ -38,3 +38,31 @@ func TestHasBlobInProject(t *testing.T) { require.Nil(t, err) assert.True(t, has) } + +func TestCountSizeOfProject(t *testing.T) { + id1, err := AddBlob(&models.Blob{ + Digest: "CountSizeOfProject_blob1", + Size: 101, + }) + require.Nil(t, err) + + id2, err := AddBlob(&models.Blob{ + Digest: "CountSizeOfProject_blob2", + Size: 202, + }) + require.Nil(t, err) + + pid1, err := AddProject(models.Project{ + Name: "CountSizeOfProject_project1", + OwnerID: 1, + }) + require.Nil(t, err) + + _, err = AddBlobToProject(id1, pid1) + require.Nil(t, err) + _, err = AddBlobToProject(id2, pid1) + require.Nil(t, err) + + pSize, err := CountSizeOfProject(pid1) + assert.Equal(t, pSize, int64(303)) +} diff --git a/src/common/quota/manager.go b/src/common/quota/manager.go index 43d70777b3..9e477f680a 100644 --- a/src/common/quota/manager.go +++ b/src/common/quota/manager.go @@ -176,16 +176,63 @@ func (m *Manager) DeleteQuota() error { // UpdateQuota update the quota resource spec func (m *Manager) UpdateQuota(hardLimits types.ResourceList) error { + o := dao.GetOrmer() if err := m.driver.Validate(hardLimits); err != nil { return err } sql := `UPDATE quota SET hard = ? WHERE reference = ? AND reference_id = ?` - _, err := dao.GetOrmer().Raw(sql, hardLimits.String(), m.reference, m.referenceID).Exec() + _, err := o.Raw(sql, hardLimits.String(), m.reference, m.referenceID).Exec() return err } +// EnsureQuota ensures the reference has quota and usage, +// if non-existent, will create new quota and usage. +// if existent, update the quota and usage. +func (m *Manager) EnsureQuota(usages types.ResourceList) error { + query := &models.QuotaQuery{ + Reference: m.reference, + ReferenceID: m.referenceID, + } + quotas, err := dao.ListQuotas(query) + if err != nil { + return err + } + + // non-existent: create quota and usage + defaultHardLimit := m.driver.HardLimits() + if len(quotas) == 0 { + _, err := m.NewQuota(defaultHardLimit, usages) + if err != nil { + return err + } + return nil + } + + // existent + used := usages + quotaUsed, err := types.NewResourceList(quotas[0].Used) + if types.Equals(quotaUsed, used) { + return nil + } + dao.WithTransaction(func(o orm.Ormer) error { + usage, err := m.getUsageForUpdate(o) + if err != nil { + return err + } + usage.Used = used.String() + usage.UpdateTime = time.Now() + _, err = o.Update(usage) + if err != nil { + return err + } + return nil + }) + + return nil +} + // AddResources add resources to usage func (m *Manager) AddResources(resources types.ResourceList) error { return dao.WithTransaction(func(o orm.Ormer) error { diff --git a/src/common/quota/manager_test.go b/src/common/quota/manager_test.go index 7de96d9980..fde4b7d82e 100644 --- a/src/common/quota/manager_test.go +++ b/src/common/quota/manager_test.go @@ -21,6 +21,7 @@ import ( "testing" "github.com/goharbor/harbor/src/common/dao" + "github.com/goharbor/harbor/src/common/models" "github.com/goharbor/harbor/src/common/quota/driver" "github.com/goharbor/harbor/src/common/quota/driver/mocks" "github.com/goharbor/harbor/src/pkg/types" @@ -131,6 +132,48 @@ func (suite *ManagerSuite) TestUpdateQuota() { } } +func (suite *ManagerSuite) TestEnsureQuota() { + // non-existent + nonExistRefID := "3" + mgr := suite.quotaManager(nonExistRefID) + infinite := types.ResourceList{types.ResourceCount: -1, types.ResourceStorage: -1} + usage := types.ResourceList{types.ResourceCount: 10, types.ResourceStorage: 10} + err := mgr.EnsureQuota(usage) + suite.Nil(err) + query := &models.QuotaQuery{ + Reference: reference, + ReferenceID: nonExistRefID, + } + quotas, err := dao.ListQuotas(query) + suite.Nil(err) + suite.Equal(usage, mustResourceList(quotas[0].Used)) + suite.Equal(infinite, mustResourceList(quotas[0].Hard)) + + // existent + existRefID := "4" + mgr = suite.quotaManager(existRefID) + used := types.ResourceList{types.ResourceCount: 11, types.ResourceStorage: 11} + if id, err := mgr.NewQuota(hardLimits, used); suite.Nil(err) { + quota, _ := dao.GetQuota(id) + suite.Equal(hardLimits, mustResourceList(quota.Hard)) + + usage, _ := dao.GetQuotaUsage(id) + suite.Equal(used, mustResourceList(usage.Used)) + } + + usage2 := types.ResourceList{types.ResourceCount: 12, types.ResourceStorage: 12} + err = mgr.EnsureQuota(usage2) + suite.Nil(err) + query2 := &models.QuotaQuery{ + Reference: reference, + ReferenceID: existRefID, + } + quotas2, err := dao.ListQuotas(query2) + suite.Equal(usage2, mustResourceList(quotas2[0].Used)) + suite.Equal(hardLimits, mustResourceList(quotas2[0].Hard)) + +} + func (suite *ManagerSuite) TestQuotaAutoCreation() { for i := 0; i < 10; i++ { mgr := suite.quotaManager(fmt.Sprintf("%d", i)) diff --git a/src/core/api/harborapi_test.go b/src/core/api/harborapi_test.go index 5357a6579b..f76530f88e 100644 --- a/src/core/api/harborapi_test.go +++ b/src/core/api/harborapi_test.go @@ -202,6 +202,8 @@ func init() { beego.Router("/api/quotas", quotaAPIType, "get:List") beego.Router("/api/quotas/:id([0-9]+)", quotaAPIType, "get:Get;put:Put") + beego.Router("/api/internal/switchquota", &InternalAPI{}, "put:SwitchQuota") + // syncRegistry if err := SyncRegistry(config.GlobalProjectMgr); err != nil { log.Fatalf("failed to sync repositories from registry: %v", err) diff --git a/src/core/api/internal.go b/src/core/api/internal.go index 71f1f317ee..841802f65a 100644 --- a/src/core/api/internal.go +++ b/src/core/api/internal.go @@ -15,12 +15,19 @@ package api import ( - "errors" - + "fmt" + "github.com/goharbor/harbor/src/chartserver" "github.com/goharbor/harbor/src/common" "github.com/goharbor/harbor/src/common/dao" "github.com/goharbor/harbor/src/common/models" + "github.com/goharbor/harbor/src/common/quota" "github.com/goharbor/harbor/src/common/utils/log" + "github.com/goharbor/harbor/src/core/config" + "github.com/goharbor/harbor/src/jobservice/logger" + "github.com/pkg/errors" + "net/url" + "strconv" + "strings" ) // InternalAPI handles request of harbor admin... @@ -69,3 +76,96 @@ func (ia *InternalAPI) RenameAdmin() { log.Debugf("The super user has been renamed to: %s", newName) ia.DestroySession() } + +type QuotaSwitcher struct { + Disabled bool +} + +// SwitchQuota ... +func (ia *InternalAPI) SwitchQuota() { + var req QuotaSwitcher + if err := ia.DecodeJSONReq(&req); err != nil { + ia.SendBadRequestError(err) + return + } + // From disable to enable, update the quota usage bases on the DB records. + if config.QuotaPerProjectEnable() == false && req.Disabled == true { + if err := ensureQuota(); err != nil { + ia.SendBadRequestError(err) + return + } + } + defer config.GetCfgManager().Set(common.QuotaPerProjectEnable, req.Disabled) + return +} + +func ensureQuota() error { + projects, err := dao.GetProjects(nil) + if err != nil { + return err + } + for _, project := range projects { + pSize, err := dao.CountSizeOfProject(project.ProjectID) + if err != nil { + logger.Warningf("error happen on counting size of project:%d , error:%v, just skip it.", project.ProjectID, err) + continue + } + afQuery := &models.ArtifactQuery{ + PID: project.ProjectID, + } + afs, err := dao.ListArtifacts(afQuery) + if err != nil { + logger.Warningf("error happen on counting number of project:%d , error:%v, just skip it.", project.ProjectID, err) + continue + } + pCount := int64(len(afs)) + + // it needs to append the chart count + if config.GetCfgManager().Get(common.WithChartMuseum).GetBool() { + chartEndpoint := strings.TrimSpace(config.GetCfgManager().Get(common.ChartRepoURL).GetString()) + if len(chartEndpoint) == 0 { + return errors.New("empty chartmuseum endpoint") + } + chartCtr, err := getChartCtr(chartEndpoint) + if err != nil { + return err + } + count, err := chartCtr.GetCountOfCharts([]string{project.Name}) + if err != nil { + err = errors.Wrap(err, fmt.Sprintf("get chart count of project %d failed", project.ProjectID)) + return err + } + pCount = pCount + int64(count) + } + + quotaMgr, err := quota.NewManager("project", strconv.FormatInt(project.ProjectID, 10)) + if err != nil { + logger.Errorf("Error occurred when to new quota manager %v, just skip it.", err) + continue + } + used := quota.ResourceList{ + quota.ResourceStorage: pSize, + quota.ResourceCount: pCount, + } + if err := quotaMgr.EnsureQuota(used); err != nil { + logger.Errorf("cannot ensure quota for the project: %d, err: %v, just skip it.", project.ProjectID, err) + continue + } + } + return nil +} + +func getChartCtr(chartEndpoint string) (*chartserver.Controller, error) { + chartEndpoint = strings.TrimSuffix(chartEndpoint, "/") + url, err := url.Parse(chartEndpoint) + if err != nil { + err = errors.New("endpoint URL of chart storage server is malformed") + return nil, err + } + ctr, err := chartserver.NewController(url) + if err != nil { + err = errors.New("failed to initialize chart API controller") + return nil, err + } + return ctr, nil +} diff --git a/src/core/api/internal_test.go b/src/core/api/internal_test.go new file mode 100644 index 0000000000..de30fbec0e --- /dev/null +++ b/src/core/api/internal_test.go @@ -0,0 +1,56 @@ +// Copyright 2018 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 api + +import ( + "net/http" + "testing" +) + +// cannot verify the real scenario here +func TestSwitchQuota(t *testing.T) { + cases := []*codeCheckingCase{ + // 401 + { + request: &testingRequest{ + method: http.MethodPut, + url: "/api/internal/switchquota", + }, + code: http.StatusUnauthorized, + }, + // 200 + { + request: &testingRequest{ + method: http.MethodPut, + url: "/api/internal/switchquota", + credential: sysAdmin, + bodyJSON: &QuotaSwitcher{ + Disabled: true, + }, + }, + code: http.StatusOK, + }, + // 403 + { + request: &testingRequest{ + url: "/api/internal/switchquota", + method: http.MethodPut, + credential: nonSysAdmin, + }, + code: http.StatusForbidden, + }, + } + runCodeCheckingCases(t, cases...) +} diff --git a/src/core/router.go b/src/core/router.go index 04fd1a1736..39729174ae 100755 --- a/src/core/router.go +++ b/src/core/router.go @@ -134,6 +134,7 @@ func initRouters() { beego.Router("/api/internal/syncregistry", &api.InternalAPI{}, "post:SyncRegistry") beego.Router("/api/internal/renameadmin", &api.InternalAPI{}, "post:RenameAdmin") + beego.Router("/api/internal/switchquota", &api.InternalAPI{}, "put:SwitchQuota") // external service that hosted on harbor process: beego.Router("/service/notifications", ®istry.NotificationHandler{})