fix(handlers): cancel context on Shutdown to stop upload purger goroutine (#4956)

registry/handlers/app.go

@@ -89,15 +89,21 @@ type App struct {
 
 	// deleteEnabled is true if the registry is configured to enable deletions.
 	deleteEnabled bool
+
+	cancel context.CancelFunc
+
+	purgerDone <-chan struct{}
 }
 
 // NewApp takes a configuration and returns a configured app, ready to serve
 // requests. The app only implements ServeHTTP and can be wrapped in other
 // handlers accordingly.
 func NewApp(ctx context.Context, config *configuration.Configuration) *App {
+	ctx, cancel := context.WithCancel(ctx)
 	app := &App{
 		Config:  config,
 		Context: ctx,
+		cancel:  cancel,
 		router:  v2.RouterWithPrefix(config.HTTP.Prefix),
 		isCache: config.Proxy.RemoteURL != "",
 	}
@@ -153,7 +159,7 @@ func NewApp(ctx context.Context, config *configuration.Configuration) *App {
 		}
 	}
 
-	startUploadPurger(app, app.driver, dcontext.GetLogger(app), purgeConfig)
+	app.purgerDone = startUploadPurger(app, app.driver, dcontext.GetLogger(app), purgeConfig)
 
 	app.driver, err = applyStorageMiddleware(app, app.driver, config.Middleware["storage"])
 	if err != nil {
@@ -453,6 +459,9 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
 
 // Shutdown close the underlying registry
 func (app *App) Shutdown() error {
+	if app.cancel != nil {
+		defer app.cancel()
+	}
 	if r, ok := app.registry.(proxy.Closer); ok {
 		return r.Close()
 	}
@@ -1072,9 +1081,11 @@ func badPurgeUploadConfig(reason string) {
 
 // startUploadPurger schedules a goroutine which will periodically
 // check upload directories for old files and delete them
-func startUploadPurger(ctx context.Context, storageDriver storagedriver.StorageDriver, log dcontext.Logger, config map[any]any) {
+func startUploadPurger(ctx context.Context, storageDriver storagedriver.StorageDriver, log dcontext.Logger, config map[any]any) <-chan struct{} {
+	done := make(chan struct{})
 	if config["enabled"] == false {
-		return
+		close(done)
+		return done
 	}
 
 	var purgeAgeDuration time.Duration
@@ -1121,6 +1132,8 @@ func startUploadPurger(ctx context.Context, storageDriver storagedriver.StorageD
 	}
 
 	go func() {
+		defer close(done)
+
 		randInt, err := rand.Int(rand.Reader, new(big.Int).SetInt64(math.MaxInt64))
 		if err != nil {
 			log.Infof("Failed to generate random jitter: %v", err)
@@ -1129,12 +1142,27 @@ func startUploadPurger(ctx context.Context, storageDriver storagedriver.StorageD
 		}
 		jitter := time.Duration(randInt.Int64()%60) * time.Minute
 		log.Infof("Starting upload purge in %s", jitter)
-		time.Sleep(jitter)
+
+		select {
+		case <-time.After(jitter):
+		case <-ctx.Done():
+			return
+		}
+
+		ticker := time.NewTicker(intervalDuration)
+		defer ticker.Stop()
 
 		for {
 			storage.PurgeUploads(ctx, storageDriver, time.Now().Add(-purgeAgeDuration), !dryRunBool)
 			log.Infof("Starting upload purge in %s", intervalDuration)
-			time.Sleep(intervalDuration)
+
+			select {
+			case <-ticker.C:
+			case <-ctx.Done():
+				return
+			}
 		}
 	}()
+
+	return done
 }

registry/handlers/app_test.go

@@ -7,7 +7,9 @@ import (
 	"net/url"
 	"reflect"
 	"testing"
+	"time"
 
+	"github.com/distribution/distribution/v3"
 	"github.com/distribution/distribution/v3/configuration"
 	"github.com/distribution/distribution/v3/internal/dcontext"
 	"github.com/distribution/distribution/v3/registry/api/errcode"
@@ -211,6 +213,93 @@ func TestNewApp(t *testing.T) {
 	}
 }
 
+type testProxyCloser struct {
+	distribution.Namespace
+	closeCalled         bool
+	ctxCancelledOnClose bool
+	app                 *App
+}
+
+func (m *testProxyCloser) Close() error {
+	m.closeCalled = true
+	select {
+	case <-m.app.Context.Done():
+		m.ctxCancelledOnClose = true
+	default:
+		m.ctxCancelledOnClose = false
+	}
+	return nil
+}
+
+func TestAppShutdownStopsUploadPurger(t *testing.T) {
+	ctx := dcontext.Background()
+	config := configuration.Configuration{
+		Storage: configuration.Storage{
+			"inmemory": configuration.Parameters{},
+			"maintenance": configuration.Parameters{
+				"uploadpurging": map[any]any{
+					"enabled":  true,
+					"age":      "168h",
+					"interval": "24h",
+					"dryrun":   false,
+				},
+			},
+		},
+	}
+	app := NewApp(ctx, &config)
+	if app.purgerDone == nil {
+		t.Fatal("expected purgerDone channel to be initialized")
+	}
+
+	if err := app.Shutdown(); err != nil {
+		t.Fatalf("unexpected error on Shutdown: %v", err)
+	}
+
+	select {
+	case <-app.purgerDone:
+		// upload purger goroutine exited
+	case <-time.After(5 * time.Second):
+		t.Fatal("upload purger goroutine did not exit within timeout after Shutdown")
+	}
+
+	select {
+	case <-app.Context.Done():
+	default:
+		t.Fatal("expected app context to be cancelled after Shutdown")
+	}
+}
+
+func TestAppShutdownDefersCancelAfterRegistryClose(t *testing.T) {
+	ctx := dcontext.Background()
+	config := configuration.Configuration{
+		Storage: configuration.Storage{
+			"inmemory": configuration.Parameters{},
+		},
+	}
+	app := NewApp(ctx, &config)
+	closer := &testProxyCloser{
+		Namespace: app.registry,
+		app:       app,
+	}
+	app.registry = closer
+
+	if err := app.Shutdown(); err != nil {
+		t.Fatalf("unexpected error on Shutdown: %v", err)
+	}
+
+	if !closer.closeCalled {
+		t.Fatal("expected registry Close to be called")
+	}
+	if closer.ctxCancelledOnClose {
+		t.Fatal("expected context NOT to be cancelled when registry Close is called")
+	}
+	select {
+	case <-app.Context.Done():
+	default:
+		t.Fatal("expected app context to be cancelled after Shutdown completes")
+	}
+}
+
 // Test the access record accumulator
 func TestAppendAccessRecords(t *testing.T) {
 	repo := "testRepo"