Add State Persistence for Crash Recovery (container-registry#321)

Signed-off-by: bupd <bupdprasanth@gmail.com>

Author: bupd

Taskfile.yml

@@ -83,6 +83,16 @@ tasks:
     cmds:
       - task: _e2e:test-spiffe
 
+  e2e-crash-recovery:
+    desc: Run crash recovery E2E test
+    cmds:
+      - task: _e2e:test-crash-recovery
+
+  e2e-crash-recovery-byo:
+    desc: Run crash recovery BYO registry E2E test
+    cmds:
+      - task: _e2e:test-crash-recovery-byo
+
   publish:
     desc: "Publish both components to registry"
     aliases: [p]

cmd/main.go

@@ -249,7 +249,7 @@ func run(opts SatelliteOptions, pathConfig *config.PathConfig, shutdownTimeout s
 		}
 	})
 
-	s := satellite.NewSatellite(cm)
+	s := satellite.NewSatellite(cm, pathConfig.StateFile)
 	err = s.Run(ctx)
 	if err != nil {
 		return fmt.Errorf("unable to start satellite: %w", err)

internal/satellite/satellite.go

@@ -10,22 +10,24 @@ import (
 )
 
 type Satellite struct {
-	cm         *config.ConfigManager
-	schedulers []*scheduler.Scheduler
+	cm            *config.ConfigManager
+	schedulers    []*scheduler.Scheduler
+	stateFilePath string
 }
 
-func NewSatellite(cm *config.ConfigManager) *Satellite {
+func NewSatellite(cm *config.ConfigManager, stateFilePath string) *Satellite {
 	return &Satellite{
-		cm:         cm,
-		schedulers: make([]*scheduler.Scheduler, 0),
+		cm:            cm,
+		schedulers:    make([]*scheduler.Scheduler, 0),
+		stateFilePath: stateFilePath,
 	}
 }
 
 func (s *Satellite) Run(ctx context.Context) error {
 	log := logger.FromContext(ctx)
 	log.Info().Msg("Starting Satellite")
 
-	fetchAndReplicateStateProcess := state.NewFetchAndReplicateStateProcess(s.cm)
+	fetchAndReplicateStateProcess := state.NewFetchAndReplicateStateProcess(s.cm, s.stateFilePath, log)
 
 	// Create ZTR scheduler if not already done
 	if !s.cm.IsZTRDone() {

internal/state/replicator.go

@@ -11,7 +11,10 @@ import (
 	"github.com/container-registry/harbor-satellite/pkg/config"
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/crane"
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
@@ -52,10 +55,10 @@ func NewBasicReplicatorWithTLS(sourceUsername, sourcePassword, sourceRegistry, r
 
 // Entity represents an image or artifact which needs to be handled by the replicator
 type Entity struct {
-	Name       string
-	Repository string
-	Tag        string
-	Digest     string
+	Name       string `json:"name"`
+	Repository string `json:"repository"`
+	Tag        string `json:"tag"`
+	Digest     string `json:"digest"`
 }
 
 func (e Entity) GetName() string {
@@ -70,35 +73,38 @@ func (e Entity) GetTag() string {
 	return e.Tag
 }
 
-// Replicate replicates images from the source registry to the Zot registry.
+// Replicate replicates images from the source registry to the local registry.
+// Before pulling, it checks which blobs already exist at the destination and
+// only downloads missing layers from source, saving bandwidth on crash recovery.
 func (r *BasicReplicator) Replicate(ctx context.Context, replicationEntities []Entity) error {
 	log := logger.FromContext(ctx)
-	pullAuthConfig := authn.FromConfig(authn.AuthConfig{
+	pullAuth := authn.FromConfig(authn.AuthConfig{
 		Username: r.sourceUsername,
 		Password: r.sourcePassword,
 	})
-	pushAuthConfig := authn.FromConfig(authn.AuthConfig{
+	pushAuth := authn.FromConfig(authn.AuthConfig{
 		Username: r.remoteUsername,
 		Password: r.remotePassword,
 	})
 
-	pullOptions := []crane.Option{crane.WithAuth(pullAuthConfig), crane.WithContext(ctx)}
-	pushOptions := []crane.Option{crane.WithAuth(pushAuthConfig), crane.WithContext(ctx)}
+	var nameOpts []name.Option
+	pullOpts := []remote.Option{remote.WithAuth(pullAuth), remote.WithContext(ctx)}
+	pushOpts := []remote.Option{remote.WithAuth(pushAuth), remote.WithContext(ctx)}
 
 	if r.useUnsecure {
-		pullOptions = append(pullOptions, crane.Insecure)
-		pushOptions = append(pushOptions, crane.Insecure)
+		nameOpts = append(nameOpts, name.Insecure)
 	} else {
 		transport, err := r.buildTLSTransport()
 		if err != nil {
 			return fmt.Errorf("build TLS transport: %w", err)
 		}
 		if transport != nil {
-			pullOptions = append(pullOptions, crane.WithTransport(transport))
+			pullOpts = append(pullOpts, remote.WithTransport(transport))
+			pushOpts = append(pushOpts, remote.WithTransport(transport))
 		}
 	}
 
-	for _, replicationEntity := range replicationEntities {
+	for _, entity := range replicationEntities {
 		// Check context cancellation before processing each image
 		select {
 		case <-ctx.Done():
@@ -107,29 +113,105 @@ func (r *BasicReplicator) Replicate(ctx context.Context, replicationEntities []E
 		default:
 		}
 
-		log.Info().Msgf("Pulling image %s from repository %s at registry %s with tag %s", replicationEntity.GetName(), replicationEntity.GetRepository(), r.sourceRegistry, replicationEntity.GetTag())
-		// Pull the image from the source registry
-		srcImage, err := crane.Pull(fmt.Sprintf("%s/%s/%s:%s", r.sourceRegistry, replicationEntity.GetRepository(), replicationEntity.GetName(), replicationEntity.GetTag()), pullOptions...)
+		srcRef := fmt.Sprintf("%s/%s/%s:%s", r.sourceRegistry, entity.GetRepository(), entity.GetName(), entity.GetTag())
+		dstRef := fmt.Sprintf("%s/%s/%s:%s", r.remoteRegistryURL, entity.GetRepository(), entity.GetName(), entity.GetTag())
+
+		src, err := name.ParseReference(srcRef, nameOpts...)
 		if err != nil {
-			log.Error().Msgf("Failed to pull image: %v", err)
-			return err
+			return fmt.Errorf("parse source ref %s: %w", srcRef, err)
 		}
 
-		// Convert Docker manifest to OCI manifest
-		ociImage := mutate.MediaType(srcImage, types.OCIManifestSchema1)
+		dst, err := name.ParseReference(dstRef, nameOpts...)
+		if err != nil {
+			return fmt.Errorf("parse dest ref %s: %w", dstRef, err)
+		}
 
-		// Push the converted OCI image to the Zot registry
-		err = crane.Push(ociImage, fmt.Sprintf("%s/%s/%s:%s", r.remoteRegistryURL, replicationEntity.GetRepository(), replicationEntity.GetName(), replicationEntity.GetTag()), pushOptions...)
+		// Lazy fetch: only the manifest is downloaded, no layer data yet
+		desc, err := remote.Get(src, pullOpts...)
 		if err != nil {
-			log.Error().Msgf("Failed to push image: %v", err)
+			log.Error().Msgf("Failed to fetch image descriptor: %v", err)
 			return err
 		}
-		log.Info().Msgf("Image %s pushed successfully", replicationEntity.GetName())
 
+		img, err := desc.Image()
+		if err != nil {
+			log.Error().Msgf("Failed to resolve image: %v", err)
+			return err
+		}
+
+		// Lazy OCI conversion, no data materialized
+		ociImage := mutate.MediaType(img, types.OCIManifestSchema1)
+
+		// Check if image already exists at destination with same digest
+		srcDigest, err := ociImage.Digest()
+		if err != nil {
+			return fmt.Errorf("compute source digest: %w", err)
+		}
+
+		dstDesc, dstErr := remote.Head(dst, pushOpts...)
+		if dstErr == nil && dstDesc.Digest == srcDigest {
+			log.Info().Msgf("Image %s already up-to-date at destination, skipping", entity.GetName())
+			continue
+		}
+
+		// Log which layers need pulling vs already present
+		srcLayers, err := ociImage.Layers()
+		if err != nil {
+			return fmt.Errorf("get source layers: %w", err)
+		}
+
+		missing := r.countMissingLayers(dst, srcLayers, pushOpts)
+		log.Info().Msgf("Replicating image %s: %d/%d layers to pull", entity.GetName(), missing, len(srcLayers))
+
+		// remote.Write streams layers one-by-one. For each layer it HEAD-checks
+		// the destination first; only missing blobs are pulled from source.
+		// Manifest is pushed last.
+		if err := remote.Write(dst, ociImage, pushOpts...); err != nil {
+			log.Error().Msgf("Failed to replicate image: %v", err)
+			return err
+		}
+		log.Info().Msgf("Image %s replicated successfully", entity.GetName())
 	}
 	return nil
 }
 
+// countMissingLayers checks which source layers are absent from the destination
+// by comparing against the existing image's layer digests (if any).
+func (r *BasicReplicator) countMissingLayers(dst name.Reference, srcLayers []v1.Layer, pushOpts []remote.Option) int {
+	dstImg, err := remote.Image(dst, pushOpts...)
+	if err != nil {
+		// No image at destination, all layers are missing
+		return len(srcLayers)
+	}
+
+	dstLayers, err := dstImg.Layers()
+	if err != nil {
+		return len(srcLayers)
+	}
+
+	existing := make(map[v1.Hash]struct{}, len(dstLayers))
+	for _, l := range dstLayers {
+		d, err := l.Digest()
+		if err != nil {
+			continue
+		}
+		existing[d] = struct{}{}
+	}
+
+	missing := 0
+	for _, l := range srcLayers {
+		d, err := l.Digest()
+		if err != nil {
+			missing++
+			continue
+		}
+		if _, ok := existing[d]; !ok {
+			missing++
+		}
+	}
+	return missing
+}
+
 func (r *BasicReplicator) DeleteReplicationEntity(ctx context.Context, replicationEntity []Entity) error {
 	log := logger.FromContext(ctx)
 	auth := authn.FromConfig(authn.AuthConfig{