feat: add context to pull and data injections (#2654)

go.mod

@@ -23,7 +23,6 @@ require (
 	github.com/fluxcd/pkg/apis/meta v1.3.0
 	github.com/fluxcd/source-controller/api v1.2.4
 	github.com/go-git/go-git/v5 v5.11.0
-	github.com/go-logr/logr v1.4.1
 	github.com/goccy/go-yaml v1.11.3
 	github.com/gofrs/flock v0.8.1
 	github.com/google/go-containerregistry v0.19.0
@@ -63,8 +62,6 @@ require (
 	sigs.k8s.io/yaml v1.4.0
 )
 
-require github.com/evanphx/json-patch/v5 v5.6.0 // indirect
-
 require (
 	atomicgo.dev/cursor v0.2.0 // indirect
 	atomicgo.dev/keyboard v0.2.9 // indirect
@@ -225,6 +222,7 @@ require (
 	github.com/emicklei/proto v1.12.1 // indirect
 	github.com/emirpasic/gods v1.18.1 // indirect
 	github.com/evanphx/json-patch v5.7.0+incompatible // indirect
+	github.com/evanphx/json-patch/v5 v5.6.0 // indirect
 	github.com/exponent-io/jsonpath v0.0.0-20151013193312-d6023ce2651d // indirect
 	github.com/facebookincubator/nvdtools v0.1.5 // indirect
 	github.com/fatih/camelcase v1.0.0 // indirect
@@ -246,6 +244,7 @@ require (
 	github.com/go-gorp/gorp/v3 v3.1.0 // indirect
 	github.com/go-ini/ini v1.67.0 // indirect
 	github.com/go-jose/go-jose/v3 v3.0.3 // indirect
+	github.com/go-logr/logr v1.4.1 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-openapi/analysis v0.22.0 // indirect
 	github.com/go-openapi/errors v0.21.0 // indirect

src/internal/packager/images/pull.go

@@ -240,10 +240,10 @@ func Pull(ctx context.Context, cfg PullConfig) (map[transform.Image]v1.Image, er
 		return err
 	}
 
-	if err := helpers.Retry(sc, 2, 5*time.Second, message.Warnf); err != nil {
+	if err := helpers.RetryWithContext(ctx, sc, 2, 5*time.Second, message.Warnf); err != nil {
 		message.Warnf("Failed to save images in parallel, falling back to sequential save: %s", err.Error())
 
-		if err := helpers.Retry(ss, 2, 5*time.Second, message.Warnf); err != nil {
+		if err := helpers.RetryWithContext(ctx, ss, 2, 5*time.Second, message.Warnf); err != nil {
 			return nil, err
 		}
 	}
@@ -348,30 +348,35 @@ func SaveConcurrent(ctx context.Context, cl clayout.Path, m map[transform.Image]
 	for info, img := range m {
 		info, img := info, img
 		eg.Go(func() error {
-			desc, err := partial.Descriptor(img)
-			if err != nil {
-				return err
-			}
+			select {
+			case <-ectx.Done():
+				return ectx.Err()
+			default:
+				desc, err := partial.Descriptor(img)
+				if err != nil {
+					return err
+				}
 
-			if err := cl.WriteImage(img); err != nil {
-				if err := CleanupInProgressLayers(ectx, img); err != nil {
-					message.WarnErr(err, "failed to clean up in-progress layers, please run `zarf tools clear-cache`")
+				if err := cl.WriteImage(img); err != nil {
+					if err := CleanupInProgressLayers(ectx, img); err != nil {
+						message.WarnErr(err, "failed to clean up in-progress layers, please run `zarf tools clear-cache`")
+					}
+					return err
 				}
-				return err
-			}
 
-			mu.Lock()
-			defer mu.Unlock()
-			annotations := map[string]string{
-				ocispec.AnnotationBaseImageName: info.Reference,
-			}
-			desc.Annotations = annotations
-			if err := cl.AppendDescriptor(*desc); err != nil {
-				return err
-			}
+				mu.Lock()
+				defer mu.Unlock()
+				annotations := map[string]string{
+					ocispec.AnnotationBaseImageName: info.Reference,
+				}
+				desc.Annotations = annotations
+				if err := cl.AppendDescriptor(*desc); err != nil {
+					return err
+				}
 
-			saved[info] = img
-			return nil
+				saved[info] = img
+				return nil
+			}
 		})
 	}
 

src/pkg/packager/deploy.go

@@ -566,7 +566,7 @@ func (p *Packager) pushReposToRepository(ctx context.Context, reposPath string,
 		}
 
 		// Try repo push up to retry limit
-		if err := helpers.Retry(tryPush, p.cfg.PkgOpts.Retries, 5*time.Second, message.Warnf); err != nil {
+		if err := helpers.RetryWithContext(ctx, tryPush, p.cfg.PkgOpts.Retries, 5*time.Second, message.Warnf); err != nil {
 			return fmt.Errorf("unable to push repo %s to the Git Server: %w", repoURL, err)
 		}
 	}