fix: more accurate storage metrics after zot restart (#1972)

Signed-off-by: Alexei Dodon <adodon@cisco.com>

pkg/extensions/monitoring/common.go

@@ -16,7 +16,7 @@ type MetricServer interface {
 	IsEnabled() bool
 }
 
-func getDirSize(path string) (int64, error) {
+func GetDirSize(path string) (int64, error) {
 	var size int64
 
 	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {

pkg/extensions/monitoring/extension.go

@@ -171,7 +171,7 @@ func IncDownloadCounter(ms MetricServer, repo string) {
 func SetStorageUsage(ms MetricServer, rootDir, repo string) {
 	ms.SendMetric(func() {
 		dir := path.Join(rootDir, repo)
-		repoSize, err := getDirSize(dir)
+		repoSize, err := GetDirSize(dir)
 
 		if err == nil {
 			repoStorageBytes.WithLabelValues(repo).Set(float64(repoSize))

pkg/extensions/monitoring/minimal.go

@@ -486,7 +486,7 @@ func IncUploadCounter(ms MetricServer, repo string) {
 func SetStorageUsage(ms MetricServer, rootDir, repo string) {
 	dir := path.Join(rootDir, repo)
 
-	repoSize, err := getDirSize(dir)
+	repoSize, err := GetDirSize(dir)
 	if err != nil {
 		ms.(*metricServer).log.Error().Err(err).Msg("failed to set storage usage")
 	}

pkg/extensions/monitoring/monitoring_test.go

@@ -4,9 +4,12 @@
 package monitoring_test
 
 import (
+	"context"
+	"fmt"
 	"math/rand"
 	"net/http"
 	"os"
+	"path"
 	"testing"
 	"time"
 
@@ -17,6 +20,8 @@ import (
 	"zotregistry.io/zot/pkg/api/config"
 	extconf "zotregistry.io/zot/pkg/extensions/config"
 	"zotregistry.io/zot/pkg/extensions/monitoring"
+	"zotregistry.io/zot/pkg/scheduler"
+	common "zotregistry.io/zot/pkg/storage/common"
 	test "zotregistry.io/zot/pkg/test/common"
 	. "zotregistry.io/zot/pkg/test/image-utils"
 	ociutils "zotregistry.io/zot/pkg/test/oci-utils"
@@ -413,6 +418,70 @@ func TestMetricsAuthorization(t *testing.T) {
 	})
 }
 
+func TestPopulateStorageMetrics(t *testing.T) {
+	Convey("Start a scheduler when metrics enabled", t, func() {
+		port := test.GetFreePort()
+		baseURL := test.GetBaseURL(port)
+		conf := config.New()
+		conf.HTTP.Port = port
+
+		rootDir := t.TempDir()
+
+		conf.Storage.RootDirectory = rootDir
+		conf.Extensions = &extconf.ExtensionConfig{}
+		enabled := true
+		conf.Extensions.Metrics = &extconf.MetricsConfig{
+			BaseConfig: extconf.BaseConfig{Enable: &enabled},
+			Prometheus: &extconf.PrometheusConfig{Path: "/metrics"},
+		}
+
+		ctlr := api.NewController(conf)
+		So(ctlr, ShouldNotBeNil)
+
+		cm := test.NewControllerManager(ctlr)
+		cm.StartAndWait(port)
+		defer cm.StopServer()
+
+		// write a couple of images
+		srcStorageCtlr := ociutils.GetDefaultStoreController(rootDir, ctlr.Log)
+		err := WriteImageToFileSystem(CreateDefaultImage(), "alpine", "0.0.1", srcStorageCtlr)
+		So(err, ShouldBeNil)
+		err = WriteImageToFileSystem(CreateDefaultImage(), "busybox", "0.0.1", srcStorageCtlr)
+		So(err, ShouldBeNil)
+
+		sch := scheduler.NewScheduler(conf, ctlr.Log)
+		ctx, cancel := context.WithCancel(context.Background())
+		sch.RunScheduler(ctx)
+
+		generator := &common.StorageMetricsInitGenerator{
+			ImgStore: ctlr.StoreController.DefaultStore,
+			Metrics:  ctlr.Metrics,
+			Log:      ctlr.Log,
+			MaxDelay: 1, // maximum delay between jobs (each job computes repo's storage size)
+		}
+
+		sch.SubmitGenerator(generator, time.Duration(0), scheduler.LowPriority)
+
+		time.Sleep(5 * time.Second)
+		cancel()
+		alpineSize, err := monitoring.GetDirSize(path.Join(rootDir, "alpine"))
+		So(err, ShouldBeNil)
+		busyboxSize, err := monitoring.GetDirSize(path.Join(rootDir, "busybox"))
+		So(err, ShouldBeNil)
+
+		resp, err := resty.R().Get(baseURL + "/metrics")
+		So(err, ShouldBeNil)
+		So(resp, ShouldNotBeNil)
+		So(resp.StatusCode(), ShouldEqual, http.StatusOK)
+
+		alpineMetric := fmt.Sprintf("zot_repo_storage_bytes{repo=\"alpine\"} %d", alpineSize)
+		busyboxMetric := fmt.Sprintf("zot_repo_storage_bytes{repo=\"busybox\"} %d", busyboxSize)
+		respStr := string(resp.Body())
+		So(respStr, ShouldContainSubstring, alpineMetric)
+		So(respStr, ShouldContainSubstring, busyboxMetric)
+	})
+}
+
 func generateRandomString() string {
 	//nolint: gosec
 	seededRand := rand.New(rand.NewSource(time.Now().UnixNano()))