feat(scheduler): gracefully shutdown

wait for workers to finish before exiting should fix tests reporting they couldn't remove rootDir because it's being written by tasks Signed-off-by: Petu Eusebiu <[email protected]>
project-zot · Nov 17, 2023 · 5ad1c4f · 5ad1c4f
1 parent 8dd06c6
commit 5ad1c4f
Show file tree

Hide file tree

Showing 46 changed files with 571 additions and 350 deletions.
diff --git a/pkg/api/controller.go b/pkg/api/controller.go
@@ -49,6 +49,7 @@ type Controller struct {
 	SyncOnDemand    SyncOnDemand
 	RelyingParties  map[string]rp.RelyingParty
 	CookieStore     *CookieStore
+	taskScheduler   *scheduler.Scheduler
 	// runtime params
 	chosenPort int // kernel-chosen port
 }
@@ -366,13 +367,14 @@ func (c *Controller) LoadNewConfig(reloadCtx context.Context, newConfig *config.
 }
 
 func (c *Controller) Shutdown() {
+	c.taskScheduler.Shutdown()
 	ctx := context.Background()
 	_ = c.Server.Shutdown(ctx)
 }
 
 func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
-	taskScheduler := scheduler.NewScheduler(c.Config, c.Log)
-	taskScheduler.RunScheduler(reloadCtx)
+	c.taskScheduler = scheduler.NewScheduler(c.Config, c.Log)
+	c.taskScheduler.RunScheduler(reloadCtx)
 
 	// Enable running garbage-collect periodically for DefaultStore
 	if c.Config.Storage.GC {
@@ -381,20 +383,20 @@ func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
 			ImageRetention: c.Config.Storage.Retention,
 		}, c.Audit, c.Log)
 
-		gc.CleanImageStorePeriodically(c.Config.Storage.GCInterval, taskScheduler)
+		gc.CleanImageStorePeriodically(c.Config.Storage.GCInterval, c.taskScheduler)
 	}
 
 	// Enable running dedupe blobs both ways (dedupe or restore deduped blobs)
-	c.StoreController.DefaultStore.RunDedupeBlobs(time.Duration(0), taskScheduler)
+	c.StoreController.DefaultStore.RunDedupeBlobs(time.Duration(0), c.taskScheduler)
 
 	// Enable extensions if extension config is provided for DefaultStore
 	if c.Config != nil && c.Config.Extensions != nil {
 		ext.EnableMetricsExtension(c.Config, c.Log, c.Config.Storage.RootDirectory)
-		ext.EnableSearchExtension(c.Config, c.StoreController, c.MetaDB, taskScheduler, c.CveScanner, c.Log)
+		ext.EnableSearchExtension(c.Config, c.StoreController, c.MetaDB, c.taskScheduler, c.CveScanner, c.Log)
 	}
 	// runs once if metrics are enabled & imagestore is local
 	if c.Config.IsMetricsEnabled() && c.Config.Storage.StorageDriver == nil {
-		c.StoreController.DefaultStore.PopulateStorageMetrics(time.Duration(0), taskScheduler)
+		c.StoreController.DefaultStore.PopulateStorageMetrics(time.Duration(0), c.taskScheduler)
 	}
 
 	if c.Config.Storage.SubPaths != nil {
@@ -407,7 +409,7 @@ func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
 						ImageRetention: storageConfig.Retention,
 					}, c.Audit, c.Log)
 
-				gc.CleanImageStorePeriodically(storageConfig.GCInterval, taskScheduler)
+				gc.CleanImageStorePeriodically(storageConfig.GCInterval, c.taskScheduler)
 			}
 
 			// Enable extensions if extension config is provided for subImageStore
@@ -418,19 +420,19 @@ func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
 			// Enable running dedupe blobs both ways (dedupe or restore deduped blobs) for subpaths
 			substore := c.StoreController.SubStore[route]
 			if substore != nil {
-				substore.RunDedupeBlobs(time.Duration(0), taskScheduler)
+				substore.RunDedupeBlobs(time.Duration(0), c.taskScheduler)
 
 				if c.Config.IsMetricsEnabled() && c.Config.Storage.StorageDriver == nil {
-					substore.PopulateStorageMetrics(time.Duration(0), taskScheduler)
+					substore.PopulateStorageMetrics(time.Duration(0), c.taskScheduler)
 				}
 			}
 		}
 	}
 
 	if c.Config.Extensions != nil {
-		ext.EnableScrubExtension(c.Config, c.Log, c.StoreController, taskScheduler)
+		ext.EnableScrubExtension(c.Config, c.Log, c.StoreController, c.taskScheduler)
 		//nolint: contextcheck
-		syncOnDemand, err := ext.EnableSyncExtension(c.Config, c.MetaDB, c.StoreController, taskScheduler, c.Log)
+		syncOnDemand, err := ext.EnableSyncExtension(c.Config, c.MetaDB, c.StoreController, c.taskScheduler, c.Log)
 		if err != nil {
 			c.Log.Error().Err(err).Msg("unable to start sync extension")
 		}
@@ -439,11 +441,11 @@ func (c *Controller) StartBackgroundTasks(reloadCtx context.Context) {
 	}
 
 	if c.CookieStore != nil {
-		c.CookieStore.RunSessionCleaner(taskScheduler)
+		c.CookieStore.RunSessionCleaner(c.taskScheduler)
 	}
 
 	// we can later move enabling the other scheduled tasks inside the call below
-	ext.EnableScheduledTasks(c.Config, taskScheduler, c.MetaDB, c.Log) //nolint: contextcheck
+	ext.EnableScheduledTasks(c.Config, c.taskScheduler, c.MetaDB, c.Log) //nolint: contextcheck
 }
 
 type SyncOnDemand interface {

diff --git a/pkg/api/controller_test.go b/pkg/api/controller_test.go
@@ -8378,7 +8378,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 			So(len(index.Manifests), ShouldEqual, 1)
 
 			// shouldn't do anything
-			err = gc.CleanRepo(repoName) //nolint: contextcheck
+			err = gc.CleanRepo(ctx, repoName) //nolint: contextcheck
 			So(err, ShouldBeNil)
 
 			// make sure both signatures are stored in repodb
@@ -8404,7 +8404,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 				err = UploadImage(img, baseURL, repoName, img.DigestStr())
 				So(err, ShouldBeNil)
 
-				err = gc.CleanRepo(repoName)
+				err = gc.CleanRepo(ctx, repoName)
 				So(err, ShouldNotBeNil)
 
 				err = os.Chmod(path.Join(dir, repoName, "blobs", "sha256", refs.Manifests[0].Digest.Encoded()), 0o755)
@@ -8418,7 +8418,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 				err = UploadImage(img, baseURL, repoName, tag)
 				So(err, ShouldBeNil)
 
-				err = gc.CleanRepo(repoName)
+				err = gc.CleanRepo(ctx, repoName)
 				So(err, ShouldNotBeNil)
 
 				err = os.WriteFile(path.Join(dir, repoName, "blobs", "sha256", refs.Manifests[0].Digest.Encoded()), content, 0o600)
@@ -8469,7 +8469,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 				So(err, ShouldBeNil)
 				newManifestDigest := godigest.FromBytes(manifestBuf)
 
-				err = gc.CleanRepo(repoName) //nolint: contextcheck
+				err = gc.CleanRepo(ctx, repoName) //nolint: contextcheck
 				So(err, ShouldBeNil)
 
 				// make sure both signatures are removed from metaDB and repo reference for untagged is removed
@@ -8548,7 +8548,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 			So(err, ShouldBeNil)
 
 			cm := test.NewControllerManager(ctlr)
-			cm.StartAndWait(port)
+			cm.StartAndWait(port) //nolint: contextcheck
 			defer cm.StopServer()
 
 			gc := gc.NewGarbageCollect(ctlr.StoreController.DefaultStore, ctlr.MetaDB,
@@ -8606,7 +8606,7 @@ func TestGCSignaturesAndUntaggedManifestsWithMetaDB(t *testing.T) {
 			So(err, ShouldBeNil)
 			So(resp.StatusCode(), ShouldEqual, http.StatusCreated)
 
-			err = gc.CleanRepo(repoName)
+			err = gc.CleanRepo(ctx, repoName)
 			So(err, ShouldBeNil)
 
 			resp, err = resty.R().SetHeader("Content-Type", ispec.MediaTypeImageIndex).

diff --git a/pkg/cli/client/cve_cmd_test.go b/pkg/cli/client/cve_cmd_test.go
@@ -583,7 +583,7 @@ func TestCVESort(t *testing.T) {
 	}
 
 	ctlr.CveScanner = mocks.CveScannerMock{
-		ScanImageFn: func(image string) (map[string]cvemodel.CVE, error) {
+		ScanImageFn: func(ctx context.Context, image string) (map[string]cvemodel.CVE, error) {
 			return map[string]cvemodel.CVE{
 				"CVE-2023-1255": {
 					ID:       "CVE-2023-1255",
@@ -687,7 +687,7 @@ func getMockCveScanner(metaDB mTypes.MetaDB) cveinfo.Scanner {
 	// MetaDB loaded with initial data now mock the scanner
 	// Setup test CVE data in mock scanner
 	scanner := mocks.CveScannerMock{
-		ScanImageFn: func(image string) (map[string]cvemodel.CVE, error) {
+		ScanImageFn: func(ctx context.Context, image string) (map[string]cvemodel.CVE, error) {
 			if strings.Contains(image, "zot-cve-test@sha256:db573b01") ||
 				image == "zot-cve-test:0.0.1" {
 				return map[string]cvemodel.CVE{

diff --git a/pkg/common/common.go b/pkg/common/common.go
@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -135,3 +136,12 @@ func IsReferrersTag(tag string) bool {
 
 	return referrersTagRule.MatchString(tag)
 }
+
+func IsContextDone(ctx context.Context) bool {
+	select {
+	case <-ctx.Done():
+		return true
+	default:
+		return false
+	}
+}
diff --git a/pkg/extensions/imagetrust/image_trust.go b/pkg/extensions/imagetrust/image_trust.go
@@ -256,8 +256,12 @@ func (validityT *validityTask) DoWork(ctx context.Context) error {
 	validityT.log.Info().Msg("update signatures validity")
 
 	for signedManifest, sigs := range validityT.repo.Signatures {
+		if zcommon.IsContextDone(ctx) {
+			return ctx.Err()
+		}
+
 		if len(sigs[zcommon.CosignSignature]) != 0 || len(sigs[zcommon.NotationSignature]) != 0 {
-			err := validityT.metaDB.UpdateSignaturesValidity(validityT.repo.Name, godigest.Digest(signedManifest))
+			err := validityT.metaDB.UpdateSignaturesValidity(ctx, validityT.repo.Name, godigest.Digest(signedManifest))
 			if err != nil {
 				validityT.log.Info().Msg("error while verifying signatures")
 

diff --git a/pkg/extensions/search/convert/convert_internal_test.go b/pkg/extensions/search/convert/convert_internal_test.go
@@ -61,7 +61,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{}, ErrTestError
 					},
@@ -99,7 +99,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{
 							Count:       1,
@@ -126,7 +126,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{}, ErrTestError
 					},
@@ -149,7 +149,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{
 							Count:       1,
@@ -179,7 +179,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{
 							Count:       1,
@@ -207,7 +207,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{
 							Count:       1,
@@ -248,7 +248,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{
 							Count:       1,
@@ -271,7 +271,7 @@ func TestCVEConvert(t *testing.T) {
 					Vulnerabilities: false,
 				},
 				mocks.CveInfoMock{
-					GetCVESummaryForImageMediaFn: func(repo string, digest, mediaType string,
+					GetCVESummaryForImageMediaFn: func(ctx context.Context, repo string, digest, mediaType string,
 					) (cvemodel.ImageCVESummary, error) {
 						return cvemodel.ImageCVESummary{}, ErrTestError
 					},

diff --git a/pkg/extensions/search/convert/cve.go b/pkg/extensions/search/convert/cve.go
@@ -47,7 +47,7 @@ func updateImageSummaryVulnerabilities(
 		return
 	}
 
-	imageCveSummary, err := cveInfo.GetCVESummaryForImageMedia(*imageSummary.RepoName, *imageSummary.Digest,
+	imageCveSummary, err := cveInfo.GetCVESummaryForImageMedia(ctx, *imageSummary.RepoName, *imageSummary.Digest,
 		*imageSummary.MediaType)
 	if err != nil {
 		// Log the error, but we should still include the image in results
@@ -91,7 +91,7 @@ func updateManifestSummaryVulnerabilities(
 		return
 	}
 
-	imageCveSummary, err := cveInfo.GetCVESummaryForImageMedia(repoName, *manifestSummary.Digest,
+	imageCveSummary, err := cveInfo.GetCVESummaryForImageMedia(ctx, repoName, *manifestSummary.Digest,
 		ispec.MediaTypeImageManifest)
 	if err != nil {
 		// Log the error, but we should still include the manifest in results