fix: make registry push cleanup per-call to prevent loss on repeated pushes

The cleanup callback stored on the Pusher was cleared after each Push call,
causing user-provided cleanup (WithImageCleanup) to be lost on subsequent
pushes. Changed the internal imageSource to return cleanup per-call as a
local variable, keeping the user cleanup immutable on the receiver.

commit: bbc425bc31fcd58ce9bfcaf38360d8daf1af8fbf
author: brice · 7mo ago · 2026-03-15 07:07 UTC
parent: 18c86d4
signature: unsigned

2 files · +17 -17

M pkg/registrypush/push.go +11 -14


@@ -39,7 +39,7 @@ type Pusher struct {
39	39		timeout time.Duration
40	40		transport http.RoundTripper
41	41		insecureTLS bool
42		-	imageSource ImageSource
	42	+	imageSource func(ctx context.Context, ref string) (v1.Image, func(), error)
43	43		cleanupFn func()
44	44		progress io.Writer
45	45		}

@@ -70,7 +70,12 @@ func WithInsecureTLS(insecure bool) Option {
70	70
71	71		// WithImageSource overrides the default daemon-based image reader.
72	72		func WithImageSource(src ImageSource) Option {
73		-	return func(p *Pusher) { p.imageSource = src }
	73	+	return func(p *Pusher) {
	74	+	p.imageSource = func(ctx context.Context, ref string) (v1.Image, func(), error) {
	75	+	img, err := src(ctx, ref)
	76	+	return img, nil, err
	77	+	}
	78	+	}
74	79		}
75	80
76	81		// WithImageCleanup sets a function called after Push completes to clean up

@@ -98,13 +103,8 @@ func New(opts ...Option) *Pusher {
98	103		}
99	104		}
100	105		if p.imageSource == nil {
101		-	p.imageSource = func(ctx context.Context, ref string) (v1.Image, error) {
102		-	img, cleanup, err := defaultImageSourceWithCleanup(ctx, ref)
103		-	if err != nil {
104		-	return nil, err
105		-	}
106		-	p.cleanupFn = cleanup
107		-	return img, nil
	106	+	p.imageSource = func(ctx context.Context, ref string) (v1.Image, func(), error) {
	107	+	return defaultImageSourceWithCleanup(ctx, ref)
108	108		}
109	109		}
110	110		return p

@@ -125,19 +125,16 @@ func (p *Pusher) Push(ctx context.Context, ref string) error {
125	125		return fmt.Errorf("image ref %s must include a tag", ref)
126	126		}
127	127
128		-	userCleanup := p.cleanupFn
129		-	p.cleanupFn = nil
130		-
131	128		authHeader, err := resolveAuthHeader(parsedRef)
132	129		if err != nil {
133	130		return err
134	131		}
135	132
136		-	img, err := p.imageSource(ctx, ref)
	133	+	img, sourceCleanup, err := p.imageSource(ctx, ref)
137	134		if err != nil {
138	135		return err
139	136		}
140		-	sourceCleanup := p.cleanupFn
	137	+	userCleanup := p.cleanupFn
141	138		switch {
142	139		case sourceCleanup != nil && userCleanup != nil:
143	140		defer func() {

 	timeout     time.Duration
 	transport   http.RoundTripper
 	insecureTLS bool
-	imageSource ImageSource
+	imageSource func(ctx context.Context, ref string) (v1.Image, func(), error)
 	cleanupFn   func()
 	progress    io.Writer
 }
 // WithImageSource overrides the default daemon-based image reader.
 func WithImageSource(src ImageSource) Option {
-	return func(p *Pusher) { p.imageSource = src }
+	return func(p *Pusher) {
+		p.imageSource = func(ctx context.Context, ref string) (v1.Image, func(), error) {
+			img, err := src(ctx, ref)
+			return img, nil, err
+		}
+	}
 }
 // WithImageCleanup sets a function called after Push completes to clean up
 		}
 	}
 	if p.imageSource == nil {
-		p.imageSource = func(ctx context.Context, ref string) (v1.Image, error) {
+		p.imageSource = func(ctx context.Context, ref string) (v1.Image, func(), error) {
-			img, cleanup, err := defaultImageSourceWithCleanup(ctx, ref)
+			return defaultImageSourceWithCleanup(ctx, ref)
-			if err != nil {
-				return nil, err
-			}
-			p.cleanupFn = cleanup
-			return img, nil
 		}
 	}
 	return p
 		return fmt.Errorf("image ref %s must include a tag", ref)
 	}
-	userCleanup := p.cleanupFn
-	p.cleanupFn = nil
 	authHeader, err := resolveAuthHeader(parsedRef)
 	if err != nil {
 		return err
 	}
-	img, err := p.imageSource(ctx, ref)
+	img, sourceCleanup, err := p.imageSource(ctx, ref)
 	if err != nil {
 		return err
 	}
-	sourceCleanup := p.cleanupFn
+	userCleanup := p.cleanupFn
 	switch {
 	case sourceCleanup != nil && userCleanup != nil:
 		defer func() {

M pkg/registrypush/push_test.go +6 -3


@@ -219,7 +219,7 @@ func TestPusher_Push(t *testing.T) {
219	219		}
220	220
221	221		func TestPusher_Push_InvokesCleanup(t *testing.T) {
222		-	var cleanupCalled bool
	222	+	var cleanupCount int
223	223		img, err := random.Image(1024, 1)
224	224		require.NoError(t, err)
225	225

@@ -232,15 +232,18 @@ func TestPusher_Push_InvokesCleanup(t *testing.T) {
232	232		return img, nil
233	233		}),
234	234		registrypush.WithImageCleanup(func() {
235		-	cleanupCalled = true
	235	+	cleanupCount++
236	236		}),
237	237		)
238	238
239	239		registryHost := strings.TrimPrefix(registry.URL, "http://")
240	240		err = p.Push(context.Background(), registryHost+"/demo/app:v1")
	241	+	require.NoError(t, err)
	242	+	assert.Equal(t, 1, cleanupCount, "cleanup callback should have been invoked after first push")
241	243
	244	+	err = p.Push(context.Background(), registryHost+"/demo/app:v1")
242	245		require.NoError(t, err)
243		-	assert.True(t, cleanupCalled, "cleanup callback should have been invoked after push")
	246	+	assert.Equal(t, 2, cleanupCount, "cleanup callback should have been invoked after second push")
244	247		}
245	248
246	249		func TestPusher_WithInsecureTLS_AllowsSelfSignedCertificate(t *testing.T) {

 }
 func TestPusher_Push_InvokesCleanup(t *testing.T) {
-	var cleanupCalled bool
+	var cleanupCount int
 	img, err := random.Image(1024, 1)
 	require.NoError(t, err)
 			return img, nil
 		}),
 		registrypush.WithImageCleanup(func() {
-			cleanupCalled = true
+			cleanupCount++
 		}),
 	)
 	registryHost := strings.TrimPrefix(registry.URL, "http://")
 	err = p.Push(context.Background(), registryHost+"/demo/app:v1")
+	require.NoError(t, err)
+	assert.Equal(t, 1, cleanupCount, "cleanup callback should have been invoked after first push")
+	err = p.Push(context.Background(), registryHost+"/demo/app:v1")
 	require.NoError(t, err)
-	assert.True(t, cleanupCalled, "cleanup callback should have been invoked after push")
+	assert.Equal(t, 2, cleanupCount, "cleanup callback should have been invoked after second push")
 }
 func TestPusher_WithInsecureTLS_AllowsSelfSignedCertificate(t *testing.T) {