Skip to content

Commit f64329e

Browse files
AustinAbro321AustinAbro321
authored
feat: add retries to image push (#3718)
Signed-off-by: Austin Abro <[email protected]>
1 parent d8fd11d commit f64329e

File tree

2 files changed

+94
-67
lines changed

2 files changed

+94
-67
lines changed

‎src/internal/packager/images/push.go

+91-65
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,15 @@ package images
77
import (
88
"context"
99
"fmt"
10+
"time"
1011

12+
"github.com/avast/retry-go/v4"
1113
"oras.land/oras-go/v2"
1214
"oras.land/oras-go/v2/content/oci"
1315
"oras.land/oras-go/v2/registry"
1416
orasRemote "oras.land/oras-go/v2/registry/remote"
1517
"oras.land/oras-go/v2/registry/remote/auth"
16-
"oras.land/oras-go/v2/registry/remote/retry"
18+
orasRetry "oras.land/oras-go/v2/registry/remote/retry"
1719

1820
"github.com/defenseunicorns/pkg/helpers/v2"
1921
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -23,41 +25,20 @@ import (
2325
"github.com/zarf-dev/zarf/src/pkg/transform"
2426
)
2527

28+
const defaultRetries = 3
29+
2630
// Push pushes images to a registry.
2731
func Push(ctx context.Context, cfg PushConfig) error {
32+
if cfg.Retries < 1 {
33+
cfg.Retries = defaultRetries
34+
}
2835
cfg.ImageList = helpers.Unique(cfg.ImageList)
36+
toPush := map[string]struct{}{}
37+
for _, img := range cfg.ImageList {
38+
toPush[img.Reference] = struct{}{}
39+
}
2940
l := logger.From(ctx)
3041
registryURL := cfg.RegistryInfo.Address
31-
var tunnel *cluster.Tunnel
32-
c, _ := cluster.NewCluster()
33-
if c != nil {
34-
var err error
35-
registryURL, tunnel, err = c.ConnectToZarfRegistryEndpoint(ctx, cfg.RegistryInfo)
36-
if err != nil {
37-
return err
38-
}
39-
if tunnel != nil {
40-
defer tunnel.Close()
41-
}
42-
}
43-
client := &auth.Client{
44-
Client: retry.DefaultClient,
45-
Cache: auth.NewCache(),
46-
Credential: auth.StaticCredential(registryURL, auth.Credential{
47-
Username: cfg.RegistryInfo.PushUsername,
48-
Password: cfg.RegistryInfo.PushPassword,
49-
}),
50-
}
51-
52-
plainHTTP := cfg.PlainHTTP
53-
54-
if dns.IsLocalhost(registryURL) && !cfg.PlainHTTP {
55-
var err error
56-
plainHTTP, err = shouldUsePlainHTTP(ctx, registryURL, client)
57-
if err != nil {
58-
return err
59-
}
60-
}
6142
err := addRefNameAnnotationToImages(cfg.SourceDirectory)
6243
if err != nil {
6344
return err
@@ -68,53 +49,98 @@ func Push(ctx context.Context, cfg PushConfig) error {
6849
return fmt.Errorf("failed to instantiate oci directory: %w", err)
6950
}
7051

71-
pushImage := func(srcName, dstName string) error {
72-
remoteRepo := &orasRemote.Repository{
73-
PlainHTTP: plainHTTP,
74-
Client: client,
75-
}
76-
remoteRepo.Reference, err = registry.ParseReference(dstName)
77-
if err != nil {
78-
return fmt.Errorf("failed to parse ref %s: %w", dstName, err)
79-
}
80-
defaultPlatform := &ocispec.Platform{
81-
Architecture: cfg.Arch,
82-
OS: "linux",
52+
// The user may or may not have a cluster available, if it's available then use it to connect to the registry
53+
c, _ := cluster.NewCluster()
54+
err = retry.Do(func() error {
55+
// Include tunnel connection in case the port forward breaks, for example, a registry pod could spin down / restart
56+
var tunnel *cluster.Tunnel
57+
if c != nil {
58+
var err error
59+
registryURL, tunnel, err = c.ConnectToZarfRegistryEndpoint(ctx, cfg.RegistryInfo)
60+
if err != nil {
61+
return err
62+
}
63+
if tunnel != nil {
64+
defer tunnel.Close()
65+
}
8366
}
84-
if tunnel != nil {
85-
return tunnel.Wrap(func() error {
86-
return copyImage(ctx, src, remoteRepo, srcName, dstName, cfg.OCIConcurrency, defaultPlatform)
87-
})
67+
client := &auth.Client{
68+
Client: orasRetry.DefaultClient,
69+
Cache: auth.NewCache(),
70+
Credential: auth.StaticCredential(registryURL, auth.Credential{
71+
Username: cfg.RegistryInfo.PushUsername,
72+
Password: cfg.RegistryInfo.PushPassword,
73+
}),
8874
}
89-
return copyImage(ctx, src, remoteRepo, srcName, dstName, cfg.OCIConcurrency, defaultPlatform)
90-
}
9175

92-
for _, img := range cfg.ImageList {
93-
l.Info("pushing image", "name", img.Reference)
94-
// If this is not a no checksum image push it for use with the Zarf agent
95-
if !cfg.NoChecksum {
96-
offlineNameCRC, err := transform.ImageTransformHost(registryURL, img.Reference)
76+
plainHTTP := cfg.PlainHTTP
77+
78+
if dns.IsLocalhost(registryURL) && !cfg.PlainHTTP {
79+
var err error
80+
plainHTTP, err = shouldUsePlainHTTP(ctx, registryURL, client)
9781
if err != nil {
9882
return err
9983
}
84+
}
10085

101-
if err = pushImage(img.Reference, offlineNameCRC); err != nil {
102-
return err
86+
pushImage := func(srcName, dstName string) error {
87+
remoteRepo := &orasRemote.Repository{
88+
PlainHTTP: plainHTTP,
89+
Client: client,
10390
}
91+
remoteRepo.Reference, err = registry.ParseReference(dstName)
92+
if err != nil {
93+
return fmt.Errorf("failed to parse ref %s: %w", dstName, err)
94+
}
95+
defaultPlatform := &ocispec.Platform{
96+
Architecture: cfg.Arch,
97+
OS: "linux",
98+
}
99+
if tunnel != nil {
100+
return tunnel.Wrap(func() error {
101+
return copyImage(ctx, src, remoteRepo, srcName, dstName, cfg.OCIConcurrency, defaultPlatform)
102+
})
103+
}
104+
return copyImage(ctx, src, remoteRepo, srcName, dstName, cfg.OCIConcurrency, defaultPlatform)
104105
}
106+
pushed := []string{}
107+
// Delete the images that were already successfully pushed so that they aren't attempted on the next retry
108+
defer func() {
109+
for _, refInfo := range pushed {
110+
delete(toPush, refInfo)
111+
}
112+
}()
113+
for img := range toPush {
114+
l.Info("pushing image", "name", img)
115+
// If this is not a no checksum image push it for use with the Zarf agent
116+
if !cfg.NoChecksum {
117+
offlineNameCRC, err := transform.ImageTransformHost(registryURL, img)
118+
if err != nil {
119+
return err
120+
}
121+
122+
if err = pushImage(img, offlineNameCRC); err != nil {
123+
return err
124+
}
125+
}
105126

106-
// To allow for other non-zarf workloads to easily see the images upload a non-checksum version
107-
// (this may result in collisions but this is acceptable for this use case)
108-
offlineName, err := transform.ImageTransformHostWithoutChecksum(registryURL, img.Reference)
109-
if err != nil {
110-
return err
111-
}
127+
// To allow for other non-zarf workloads to easily see the images upload a non-checksum version
128+
// (this may result in collisions but this is acceptable for this use case)
129+
offlineName, err := transform.ImageTransformHostWithoutChecksum(registryURL, img)
130+
if err != nil {
131+
return err
132+
}
112133

113-
if err = pushImage(img.Reference, offlineName); err != nil {
114-
return err
134+
if err = pushImage(img, offlineName); err != nil {
135+
return err
136+
}
137+
pushed = append(pushed, img)
115138
}
139+
return nil
140+
}, retry.Context(ctx), retry.Attempts(uint(cfg.Retries)), retry.Delay(500*time.Millisecond))
141+
if err != nil {
142+
return err
116143
}
117-
118144
return nil
119145
}
120146

‎src/internal/packager2/mirror.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ type MirrorOptions struct {
3939

4040
// Mirror mirrors the package contents to the given registry and git server.
4141
func Mirror(ctx context.Context, opt MirrorOptions) error {
42-
err := pushImagesToRegistry(ctx, opt.PkgLayout, opt.RegistryInfo, opt.NoImageChecksum, opt.PlainHTTP, opt.OCIConcurrency)
42+
err := pushImagesToRegistry(ctx, opt.PkgLayout, opt.RegistryInfo, opt.NoImageChecksum, opt.PlainHTTP, opt.OCIConcurrency, opt.Retries)
4343
if err != nil {
4444
return err
4545
}
@@ -50,7 +50,7 @@ func Mirror(ctx context.Context, opt MirrorOptions) error {
5050
return nil
5151
}
5252

53-
func pushImagesToRegistry(ctx context.Context, pkgLayout *layout.PackageLayout, registryInfo types.RegistryInfo, noImgChecksum bool, plainHTTP bool, concurrency int) error {
53+
func pushImagesToRegistry(ctx context.Context, pkgLayout *layout.PackageLayout, registryInfo types.RegistryInfo, noImgChecksum bool, plainHTTP bool, concurrency int, retries int) error {
5454
refs := []transform.Image{}
5555
for _, component := range pkgLayout.Pkg.Components {
5656
for _, img := range component.Images {
@@ -72,6 +72,7 @@ func pushImagesToRegistry(ctx context.Context, pkgLayout *layout.PackageLayout,
7272
PlainHTTP: plainHTTP,
7373
NoChecksum: noImgChecksum,
7474
Arch: pkgLayout.Pkg.Build.Architecture,
75+
Retries: retries,
7576
}
7677
err := images.Push(ctx, pushConfig)
7778
if err != nil {

0 commit comments

Comments
 (0)