diff --git a/src/api/chartmuseum/controller.go b/src/api/chartmuseum/controller.go index 6c1cb7061..4cd1087ee 100644 --- a/src/api/chartmuseum/controller.go +++ b/src/api/chartmuseum/controller.go @@ -17,7 +17,6 @@ package chartmuseum import ( "context" "errors" - "fmt" "net/http" "net/url" "strings" @@ -25,8 +24,9 @@ import ( "github.com/goharbor/harbor/src/api/project" "github.com/goharbor/harbor/src/chartserver" + "github.com/goharbor/harbor/src/common" + "github.com/goharbor/harbor/src/common/config" commonhttp "github.com/goharbor/harbor/src/common/http" - "github.com/goharbor/harbor/src/core/config" ) var ( @@ -51,46 +51,54 @@ func NewController() Controller { } type controller struct { - projectCtl project.Controller - cc *chartserver.Controller - ccError error - ccOnce sync.Once + projectCtl project.Controller + cc *chartserver.Controller + withChartMuseum bool + + initializeError error + initializeOnce sync.Once } -func (c *controller) initialize() (*chartserver.Controller, error) { - c.ccOnce.Do(func() { - addr, err := config.GetChartMuseumEndpoint() - if err != nil { - c.ccError = fmt.Errorf("failed to get the endpoint URL of chart storage server: %s", err.Error()) +func (c *controller) initialize() error { + c.initializeOnce.Do(func() { + cfg := config.NewDBCfgManager() + + c.withChartMuseum = cfg.Get(common.WithChartMuseum).GetBool() + if !c.withChartMuseum { return } - addr = strings.TrimSuffix(addr, "/") - url, err := url.Parse(addr) + chartEndpoint := strings.TrimSpace(cfg.Get(common.ChartRepoURL).GetString()) + if len(chartEndpoint) == 0 { + c.initializeError = errors.New("empty chartmuseum endpoint") + return + } + + url, err := url.Parse(strings.TrimSuffix(chartEndpoint, "/")) if err != nil { - c.ccError = errors.New("endpoint URL of chart storage server is malformed") + c.initializeError = errors.New("endpoint URL of chart storage server is malformed") return } ctr, err := chartserver.NewController(url) if err != nil { - c.ccError = errors.New("failed to initialize chart API controller") + c.initializeError = errors.New("failed to initialize chart API controller") + return } c.cc = ctr }) - return c.cc, c.ccError + return c.initializeError } func (c *controller) Count(ctx context.Context, projectID int64) (int64, error) { - if !config.WithChartMuseum() { - return 0, nil + if err := c.initialize(); err != nil { + return 0, err } - cc, err := c.initialize() - if err != nil { - return 0, err + if !c.withChartMuseum { + return 0, nil } proj, err := c.projectCtl.Get(ctx, projectID) @@ -98,7 +106,7 @@ func (c *controller) Count(ctx context.Context, projectID int64) (int64, error) return 0, err } - count, err := cc.GetCountOfCharts([]string{proj.Name}) + count, err := c.cc.GetCountOfCharts([]string{proj.Name}) if err != nil { return 0, err } @@ -107,13 +115,12 @@ func (c *controller) Count(ctx context.Context, projectID int64) (int64, error) } func (c *controller) Exist(ctx context.Context, projectID int64, chartName, version string) (bool, error) { - if !config.WithChartMuseum() { - return false, nil + if err := c.initialize(); err != nil { + return false, err } - cc, err := c.initialize() - if err != nil { - return false, err + if !c.withChartMuseum { + return false, nil } proj, err := c.projectCtl.Get(ctx, projectID) @@ -121,7 +128,7 @@ func (c *controller) Exist(ctx context.Context, projectID int64, chartName, vers return false, err } - chartVersion, err := cc.GetChartVersion(proj.Name, chartName, version) + chartVersion, err := c.cc.GetChartVersion(proj.Name, chartName, version) if err != nil { var httpErr *commonhttp.Error if errors.As(err, &httpErr) { diff --git a/src/api/project/controller.go b/src/api/project/controller.go index 156bd533d..9495b097b 100644 --- a/src/api/project/controller.go +++ b/src/api/project/controller.go @@ -36,6 +36,8 @@ type Controller interface { Get(ctx context.Context, projectID int64, options ...Option) (*models.Project, error) // GetByName get the project by project name GetByName(ctx context.Context, projectName string, options ...Option) (*models.Project, error) + // List list projects + List(ctx context.Context, query *models.ProjectQueryParam, options ...Option) ([]*models.Project, error) } // NewController creates an instance of the default project controller @@ -62,7 +64,7 @@ func (c *controller) Get(ctx context.Context, projectID int64, options ...Option return nil, ierror.NotFoundError(nil).WithMessage("project %d not found", projectID) } - return c.assembleProject(ctx, p, options...) + return c.assembleProject(ctx, p, newOptions(options...)) } func (c *controller) GetByName(ctx context.Context, projectName string, options ...Option) (*models.Project, error) { @@ -78,12 +80,26 @@ func (c *controller) GetByName(ctx context.Context, projectName string, options return nil, ierror.NotFoundError(nil).WithMessage("project %s not found", projectName) } - return c.assembleProject(ctx, p, options...) + return c.assembleProject(ctx, p, newOptions(options...)) } -func (c *controller) assembleProject(ctx context.Context, p *models.Project, options ...Option) (*models.Project, error) { - opts := newOptions(options...) +func (c *controller) List(ctx context.Context, query *models.ProjectQueryParam, options ...Option) ([]*models.Project, error) { + projects, err := c.projectMgr.List(query) + if err != nil { + return nil, err + } + opts := newOptions(options...) + for _, p := range projects { + if _, err := c.assembleProject(ctx, p, opts); err != nil { + return nil, err + } + } + + return projects, nil +} + +func (c *controller) assembleProject(ctx context.Context, p *models.Project, opts *Options) (*models.Project, error) { if opts.Metadata { meta, err := c.metaMgr.Get(p.ProjectID) if err != nil { diff --git a/src/api/quota/util.go b/src/api/quota/util.go index 69735f281..51317848d 100644 --- a/src/api/quota/util.go +++ b/src/api/quota/util.go @@ -15,8 +15,14 @@ package quota import ( + "context" "fmt" "strconv" + + "github.com/goharbor/harbor/src/api/project" + "github.com/goharbor/harbor/src/common/models" + "github.com/goharbor/harbor/src/common/utils/log" + ierror "github.com/goharbor/harbor/src/internal/error" ) const ( @@ -39,3 +45,67 @@ func ReferenceID(i interface{}) string { return fmt.Sprintf("%v", i) } } + +// RefreshForProjects refresh quotas of all projects +func RefreshForProjects(ctx context.Context) error { + log := log.G(ctx) + + driver, err := Driver(ctx, ProjectReference) + if err != nil { + return err + } + + projects := func(chunkSize int) <-chan *models.Project { + ch := make(chan *models.Project, chunkSize) + + go func() { + defer close(ch) + + params := &models.ProjectQueryParam{ + Pagination: &models.Pagination{Page: 1, Size: int64(chunkSize)}, + } + + for { + results, err := project.Ctl.List(ctx, params, project.Metadata(false)) + if err != nil { + log.Errorf("list projects failed, error: %v", err) + return + } + + for _, p := range results { + ch <- p + } + + if len(results) < chunkSize { + break + } + + params.Pagination.Page++ + } + + }() + + return ch + }(50) // default chunk size is 50 + + for p := range projects { + referenceID := ReferenceID(p.ProjectID) + + _, err := Ctl.GetByRef(ctx, ProjectReference, referenceID) + if ierror.IsNotFoundErr(err) { + if _, err := Ctl.Create(ctx, ProjectReference, referenceID, driver.HardLimits(ctx)); err != nil { + log.Warningf("initialize quota for project %s failed, error: %v", p.Name, err) + continue + } + } else if err != nil { + log.Warningf("get quota of the project %s failed, error: %v", p.Name, err) + continue + } + + if err := Ctl.Refresh(ctx, ProjectReference, referenceID, IgnoreLimitation(true)); err != nil { + log.Warningf("refresh quota usage for project %s failed, error: %v", p.Name, err) + } + } + + return nil +} diff --git a/src/api/quota/util_test.go b/src/api/quota/util_test.go new file mode 100644 index 000000000..18c8e4a71 --- /dev/null +++ b/src/api/quota/util_test.go @@ -0,0 +1,122 @@ +// 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 ( + "context" + "math/rand" + "testing" + "time" + + "github.com/goharbor/harbor/src/api/project" + "github.com/goharbor/harbor/src/common/models" + "github.com/goharbor/harbor/src/internal/orm" + "github.com/goharbor/harbor/src/pkg/quota" + "github.com/goharbor/harbor/src/pkg/quota/driver" + "github.com/goharbor/harbor/src/pkg/types" + projecttesting "github.com/goharbor/harbor/src/testing/api/project" + ormtesting "github.com/goharbor/harbor/src/testing/lib/orm" + "github.com/goharbor/harbor/src/testing/mock" + quotatesting "github.com/goharbor/harbor/src/testing/pkg/quota" + drivertesting "github.com/goharbor/harbor/src/testing/pkg/quota/driver" + "github.com/stretchr/testify/suite" +) + +type RefreshForProjectsTestSuite struct { + suite.Suite + + originalProjectCtl project.Controller + projectCtl *projecttesting.Controller + + originalQuotaCtl Controller + quotaMgr quota.Manager + + originalDriver driver.Driver + driver *drivertesting.Driver +} + +func (suite *RefreshForProjectsTestSuite) SetupTest() { + suite.originalDriver, _ = Driver(context.TODO(), ProjectReference) + suite.driver = &drivertesting.Driver{} + driver.Register(ProjectReference, suite.driver) + + suite.originalProjectCtl = project.Ctl + suite.projectCtl = &projecttesting.Controller{} + project.Ctl = suite.projectCtl + + suite.originalQuotaCtl = Ctl + + suite.quotaMgr = "atesting.Manager{} + Ctl = &controller{ + quotaMgr: suite.quotaMgr, + } +} + +func (suite *RefreshForProjectsTestSuite) TearDownTest() { + project.Ctl = suite.originalProjectCtl + Ctl = suite.originalQuotaCtl + + driver.Register(ProjectReference, suite.originalDriver) +} + +func (suite *RefreshForProjectsTestSuite) TestRefreshForProjects() { + rand.Seed(time.Now().UnixNano()) + + startProjectID := rand.Int63() + var firstPageProjects, secondPageProjects []*models.Project + for i := 0; i < 50; i++ { + firstPageProjects = append(firstPageProjects, &models.Project{ + ProjectID: startProjectID + int64(i), + }) + } + + for i := 0; i < 10; i++ { + secondPageProjects = append(secondPageProjects, &models.Project{ + ProjectID: startProjectID + 50 + int64(i), + }) + } + + page := 1 + mock.OnAnything(suite.projectCtl, "List").Return(func(context.Context, *models.ProjectQueryParam, ...project.Option) []*models.Project { + defer func() { + page++ + }() + + if page == 1 { + return firstPageProjects + } else if page == 2 { + return secondPageProjects + } else { + return nil + } + }, nil) + + q := "a.Quota{} + q.SetHard(types.ResourceList{types.ResourceCount: 10}) + q.SetUsed(types.ResourceList{types.ResourceCount: 0}) + + mock.OnAnything(suite.quotaMgr, "GetByRef").Return(q, nil) + mock.OnAnything(suite.quotaMgr, "GetByRefForUpdate").Return(q, nil) + mock.OnAnything(suite.quotaMgr, "Update").Return(nil) + mock.OnAnything(suite.driver, "CalculateUsage").Return(types.ResourceList{types.ResourceCount: 1}, nil) + + ctx := orm.NewContext(context.TODO(), &ormtesting.FakeOrmer{}) + RefreshForProjects(ctx) + suite.Equal(3, page) +} + +func TestRefreshForProjectsTestSuite(t *testing.T) { + suite.Run(t, &RefreshForProjectsTestSuite{}) +} diff --git a/src/core/api/internal.go b/src/core/api/internal.go index bdd33ea63..29d186722 100644 --- a/src/core/api/internal.go +++ b/src/core/api/internal.go @@ -88,7 +88,7 @@ func (ia *InternalAPI) SwitchQuota() { } ctx := orm.NewContext(ia.Ctx.Request.Context(), o.NewOrm()) - if err := ia.refreshQuotas(ctx); err != nil { + if err := quota.RefreshForProjects(ctx); err != nil { ia.SendInternalServerError(err) return } @@ -101,38 +101,6 @@ func (ia *InternalAPI) SwitchQuota() { return } -func (ia *InternalAPI) refreshQuotas(ctx context.Context) error { - driver, err := quota.Driver(ctx, quota.ProjectReference) - if err != nil { - return err - } - - projects, err := dao.GetProjects(nil) - if err != nil { - return err - } - - for _, project := range projects { - referenceID := quota.ReferenceID(project.ProjectID) - - _, err := quota.Ctl.GetByRef(ctx, quota.ProjectReference, referenceID) - if ierror.IsNotFoundErr(err) { - if _, err := quota.Ctl.Create(ctx, quota.ProjectReference, referenceID, driver.HardLimits(ctx)); err != nil { - log.Warningf("initialize quota for project %s failed, error: %v", project.Name, err) - continue - } - } else if err != nil { - log.Warningf("get quota of the project %s failed, error: %v", project.Name, err) - continue - } - - if err := quota.Ctl.Refresh(ctx, quota.ProjectReference, referenceID, quota.IgnoreLimitation(true)); err != nil { - log.Warningf("refresh quota usage for project %s failed, error: %v", project.Name, err) - } - } - return nil -} - // SyncQuota ... func (ia *InternalAPI) SyncQuota() { if !config.QuotaPerProjectEnable() { @@ -154,7 +122,7 @@ func (ia *InternalAPI) SyncQuota() { }() log.Info("start to sync quota(API), the system will be set to ReadOnly and back it normal once it done.") ctx := orm.NewContext(context.TODO(), o.NewOrm()) - err := ia.refreshQuotas(ctx) + err := quota.RefreshForProjects(ctx) if err != nil { log.Errorf("fail to sync quota(API), but with error: %v, please try to do it again.", err) return diff --git a/src/core/service/notifications/admin/handler.go b/src/core/service/notifications/admin/handler.go index a7a50b72e..4032621eb 100644 --- a/src/core/service/notifications/admin/handler.go +++ b/src/core/service/notifications/admin/handler.go @@ -19,12 +19,14 @@ import ( "encoding/json" o "github.com/astaxie/beego/orm" + "github.com/goharbor/harbor/src/api/quota" "github.com/goharbor/harbor/src/api/scan" "github.com/goharbor/harbor/src/common/dao" "github.com/goharbor/harbor/src/common/job" job_model "github.com/goharbor/harbor/src/common/job/models" "github.com/goharbor/harbor/src/common/models" "github.com/goharbor/harbor/src/common/utils/log" + "github.com/goharbor/harbor/src/core/config" "github.com/goharbor/harbor/src/core/service/notifications" "github.com/goharbor/harbor/src/internal/orm" j "github.com/goharbor/harbor/src/jobservice/job" @@ -109,5 +111,11 @@ func (h *Handler) HandleAdminJob() { // For scan all job if h.jobName == job.ImageScanAllJob && h.checkIn != "" { go scan.HandleCheckIn(orm.NewContext(context.TODO(), o.NewOrm()), h.checkIn) + } else if h.jobName == job.ImageGC && h.status == models.JobFinished { + go func() { + if config.QuotaPerProjectEnable() { + quota.RefreshForProjects(orm.NewContext(context.TODO(), o.NewOrm())) + } + }() } } diff --git a/src/testing/api/project/controller.go b/src/testing/api/project/controller.go index b3bd33ac0..37941a450 100644 --- a/src/testing/api/project/controller.go +++ b/src/testing/api/project/controller.go @@ -75,3 +75,33 @@ func (_m *Controller) GetByName(ctx context.Context, projectName string, options return r0, r1 } + +// List provides a mock function with given fields: ctx, query, options +func (_m *Controller) List(ctx context.Context, query *models.ProjectQueryParam, options ...project.Option) ([]*models.Project, error) { + _va := make([]interface{}, len(options)) + for _i := range options { + _va[_i] = options[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, query) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + var r0 []*models.Project + if rf, ok := ret.Get(0).(func(context.Context, *models.ProjectQueryParam, ...project.Option) []*models.Project); ok { + r0 = rf(ctx, query, options...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*models.Project) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func(context.Context, *models.ProjectQueryParam, ...project.Option) error); ok { + r1 = rf(ctx, query, options...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +}