add statistics metrics collector

Signed-off-by: Maksym Trofimenko <maksym@container-registry.com>

src/pkg/exporter/collector_test.go

@@ -0,0 +1,29 @@
+// 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 exporter
+
+import (
+	"github.com/stretchr/testify/suite"
+	"testing"
+)
+
+func TestCollectorsTestSuite(t *testing.T) {
+	setupTest(t)
+	defer tearDownTest(t)
+	suite.Run(t, new(ProjectCollectorTestSuite))
+	suite.Run(t, &StatisticsCollectorTestSuite{
+		collector: NewStatisticsCollector(),
+	})
+}

src/pkg/exporter/exporter.go

@@ -50,7 +50,9 @@ func NewExporter(opt *Opt) *Exporter {
 	err := exporter.RegisterCollector(NewHealthCollect(hbrCli),
 		NewSystemInfoCollector(hbrCli),
 		NewProjectCollector(),
-		NewJobServiceCollector())
+		NewJobServiceCollector(),
+		NewStatisticsCollector(),
+	)
 	if err != nil {
 		log.Warningf("calling RegisterCollector() errored out, error: %v", err)
 	}

src/pkg/exporter/project_collector_test.go

@@ -1,6 +1,7 @@
 package exporter
 
 import (
+	"github.com/goharbor/harbor/src/common/utils/test"
 	"strconv"
 	"testing"
 	"time"
@@ -10,7 +11,6 @@ import (
 	"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/utils/test"
 	proctl "github.com/goharbor/harbor/src/controller/project"
 	quotactl "github.com/goharbor/harbor/src/controller/quota"
 	"github.com/goharbor/harbor/src/lib/orm"
@@ -41,8 +41,8 @@ var (
 
 func setupTest(t *testing.T) {
 	test.InitDatabaseFromEnv()
-	ctx := orm.Context()
 
+	ctx := orm.Context()
 	// register projAdmin and assign project admin role
 	aliceID, err := user.Mgr.Create(ctx, &alice)
 	if err != nil {
@@ -137,11 +137,11 @@ func tearDownTest(t *testing.T) {
 	dao.GetOrmer().Raw("delete from harbor_user where user_id in (?, ?, ?)", []int{alice.UserID, bob.UserID, eve.UserID}).Exec()
 }
 
-type PorjectCollectorTestSuite struct {
+type ProjectCollectorTestSuite struct {
 	suite.Suite
 }
 
-func (c *PorjectCollectorTestSuite) TestProjectCollector() {
+func (c *ProjectCollectorTestSuite) TestProjectCollector() {
 	pMap := make(map[int64]*projectInfo)
 	updateProjectBasicInfo(pMap)
 	updateProjectMemberInfo(pMap)
@@ -169,9 +169,3 @@ func (c *PorjectCollectorTestSuite) TestProjectCollector() {
 	c.Equalf(pMap[testPro2.ProjectID].Artifact["IMAGE"].ArtifactTotal, float64(1), "pMap %v", pMap)
 
 }
-
-func TestPorjectCollectorTestSuite(t *testing.T) {
-	setupTest(t)
-	defer tearDownTest(t)
-	suite.Run(t, new(PorjectCollectorTestSuite))
-}

src/pkg/exporter/statistics_collector.go

@@ -0,0 +1,176 @@
+// 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 exporter
+
+import (
+	"context"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/goharbor/harbor/src/controller/blob"
+	"github.com/goharbor/harbor/src/controller/project"
+	"github.com/goharbor/harbor/src/controller/repository"
+	"github.com/goharbor/harbor/src/lib/log"
+	"github.com/goharbor/harbor/src/lib/orm"
+	"github.com/goharbor/harbor/src/lib/q"
+	"github.com/goharbor/harbor/src/pkg/systemartifact"
+)
+
+const StatisticsCollectorName = "StatisticsCollector"
+
+var (
+	totalUsage = typedDesc{
+		desc:      newDescWithLables("", "statistics_total_storage_consumption", "Total storage used"),
+		valueType: prometheus.GaugeValue,
+	}
+	totalProjectAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_total_project_amount", "Total amount of projects"),
+		valueType: prometheus.GaugeValue,
+	}
+	publicProjectAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_public_project_amount", "Amount of public projects"),
+		valueType: prometheus.GaugeValue,
+	}
+	privateProjectAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_private_project_amount", "Amount of private projects"),
+		valueType: prometheus.GaugeValue,
+	}
+	totalRepoAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_total_repo_amount", "Total amount of repositories"),
+		valueType: prometheus.GaugeValue,
+	}
+	publicRepoAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_public_repo_amount", "Amount of public repositories"),
+		valueType: prometheus.GaugeValue,
+	}
+	privateRepoAmount = typedDesc{
+		desc:      newDescWithLables("", "statistics_private_repo_amount", "Amount of private repositories"),
+		valueType: prometheus.GaugeValue,
+	}
+)
+
+type StatisticsCollector struct {
+	proCtl            project.Controller
+	repoCtl           repository.Controller
+	blobCtl           blob.Controller
+	systemArtifactMgr systemartifact.Manager
+}
+
+func NewStatisticsCollector() *StatisticsCollector {
+	return &StatisticsCollector{
+		blobCtl:           blob.Ctl,
+		systemArtifactMgr: systemartifact.Mgr,
+		proCtl:            project.Ctl,
+		repoCtl:           repository.Ctl,
+	}
+}
+
+func (g StatisticsCollector) GetName() string {
+	return StatisticsCollectorName
+}
+
+func (g StatisticsCollector) Describe(c chan<- *prometheus.Desc) {
+	c <- totalUsage.Desc()
+}
+
+func (g StatisticsCollector) getTotalUsageMetric(ctx context.Context) prometheus.Metric {
+	sum, _ := g.blobCtl.CalculateTotalSize(ctx, true)
+	sysArtifactStorageSize, _ := g.systemArtifactMgr.GetStorageSize(ctx)
+	return totalUsage.MustNewConstMetric(float64(sum + sysArtifactStorageSize))
+}
+
+func (g StatisticsCollector) getTotalRepoAmount(ctx context.Context) int64 {
+	n, err := g.repoCtl.Count(ctx, nil)
+	if err != nil {
+		log.Errorf("get total repositories error: %v", err)
+		return 0
+	}
+	return n
+}
+
+func (g StatisticsCollector) getTotalProjectsAmount(ctx context.Context) int64 {
+	count, err := g.proCtl.Count(ctx, nil)
+	if err != nil {
+		log.Errorf("get total projects error: %v", err)
+		return 0
+	}
+	return count
+}
+
+func (g StatisticsCollector) getPublicProjectsAndRepositories(ctx context.Context) (int64, int64) {
+	pubProjects, err := g.proCtl.List(ctx, q.New(q.KeyWords{"public": true}), project.Metadata(false))
+	if err != nil {
+		log.Errorf("get public projects error: %v", err)
+	}
+	pubProjectsAmount := int64(len(pubProjects))
+
+	if pubProjectsAmount == 0 {
+		return pubProjectsAmount, 0
+	}
+	var ids []interface{}
+	for _, p := range pubProjects {
+		ids = append(ids, p.ProjectID)
+	}
+	n, err := g.repoCtl.Count(ctx, &q.Query{
+		Keywords: map[string]interface{}{
+			"ProjectID": q.NewOrList(ids),
+		},
+	})
+	if err != nil {
+		log.Errorf("get public repo error: %v", err)
+		return pubProjectsAmount, 0
+	}
+	return pubProjectsAmount, n
+}
+
+// Collect implements prometheus.Collector
+func (g StatisticsCollector) Collect(c chan<- prometheus.Metric) {
+	for _, m := range g.getStatistics() {
+		c <- m
+	}
+}
+
+func (g StatisticsCollector) getStatistics() []prometheus.Metric {
+	if CacheEnabled() {
+		value, ok := CacheGet(StatisticsCollectorName)
+		if ok {
+			return value.([]prometheus.Metric)
+		}
+	}
+	var (
+		result []prometheus.Metric
+		ctx    = orm.Context()
+	)
+
+	var (
+		publicProjects, publicRepos = g.getPublicProjectsAndRepositories(ctx)
+		totalProjects               = g.getTotalProjectsAmount(ctx)
+		totalRepos                  = g.getTotalRepoAmount(ctx)
+	)
+
+	result = []prometheus.Metric{
+		totalRepoAmount.MustNewConstMetric(float64(totalRepos)),
+		publicRepoAmount.MustNewConstMetric(float64(publicRepos)),
+		privateRepoAmount.MustNewConstMetric(float64(totalRepos) - float64(publicRepos)),
+		totalProjectAmount.MustNewConstMetric(float64(totalProjects)),
+		publicProjectAmount.MustNewConstMetric(float64(publicProjects)),
+		privateProjectAmount.MustNewConstMetric(float64(totalProjects) - float64(publicProjects)),
+		g.getTotalUsageMetric(ctx),
+	}
+	if CacheEnabled() {
+		CachePut(StatisticsCollectorName, result)
+	}
+	return result
+}

src/pkg/exporter/statistics_collector_test.go

@@ -0,0 +1,58 @@
+package exporter
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/suite"
+)
+
+type StatisticsCollectorTestSuite struct {
+	suite.Suite
+	collector *StatisticsCollector
+}
+
+func (c *StatisticsCollectorTestSuite) TestStatisticsCollector() {
+	metrics := c.collector.getStatistics()
+	c.Equalf(7, len(metrics), "statistics collector should return %d metrics", 7)
+	c.testGaugeMetric(metrics[0], 2, "total repo amount mismatch")     // total repo amount
+	c.testGaugeMetric(metrics[1], 1, "public repo amount mismatch")    // only one project is public so its single repo is public too
+	c.testGaugeMetric(metrics[2], 1, "primate repo amount mismatch")   //
+	c.testGaugeMetric(metrics[3], 3, "total project amount mismatch")  // including library, project by default
+	c.testGaugeMetric(metrics[4], 2, "public project amount mismatch") // including library, project by default
+	c.testGaugeMetric(metrics[5], 1, "private project amount mismatch")
+	c.testGaugeMetric(metrics[6], 0, "total storage usage mismatch") // still zero
+}
+
+func (c *StatisticsCollectorTestSuite) getMetricDTO(m prometheus.Metric) *dto.Metric {
+	d := &dto.Metric{}
+	c.NoError(m.Write(d))
+	return d
+}
+
+func (c *StatisticsCollectorTestSuite) testCounterMetric(m prometheus.Metric, value float64) {
+	d := c.getMetricDTO(m)
+	if !c.NotNilf(d, "write metric error") {
+		return
+	}
+	if !c.NotNilf(d.Counter, "counter is nil") {
+		return
+	}
+	if !c.NotNilf(d.Counter.Value, "counter value is nil") {
+		return
+	}
+	c.Equalf(value, *d.Counter.Value, "expected counter value does not match: expected: %v actual: %v", value, *d.Counter.Value)
+}
+
+func (c *StatisticsCollectorTestSuite) testGaugeMetric(m prometheus.Metric, value float64, msg string) {
+	d := c.getMetricDTO(m)
+	if !c.NotNilf(d, "write metric error") {
+		return
+	}
+	if !c.NotNilf(d.Gauge, "gauge is nil") {
+		return
+	}
+	if !c.NotNilf(d.Gauge.Value, "gauge value is nil") {
+		return
+	}
+	c.Equalf(value, *d.Gauge.Value, "%s expected: %v actual: %v", msg, value, *d.Gauge.Value)
+}