feat(cleanup): keep dependencies (#6704)

Signed-off-by: Yaroslav Pershin <62902094+iapershin@users.noreply.github.com>
Signed-off-by: Aleksei Igrychev <aleksei.igrychev@palark.com>
Co-authored-by: Aleksei Igrychev <aleksei.igrychev@palark.com>

Author: iapershin

Diff for: pkg/build/stage/dependencies.go

@@ -96,6 +96,7 @@ func (s *DependenciesStage) GetDependencies(ctx context.Context, c Conveyor, cb
 }
 
 func (s *DependenciesStage) prepareImageWithLegacyStapelBuilder(ctx context.Context, c Conveyor, cr container_backend.ContainerBackend, _, stageImage *StageImage) error {
+	imageServiceCommitChangeOptions := stageImage.Builder.LegacyStapelStageBuilder().Container().ServiceCommitChangeOptions()
 	for _, elm := range s.imports {
 		sourceImageName := getSourceImageName(elm)
 		srv, err := c.GetImportServer(ctx, s.targetPlatform, sourceImageName, elm.Stage)
@@ -106,8 +107,6 @@ func (s *DependenciesStage) prepareImageWithLegacyStapelBuilder(ctx context.Cont
 		command := srv.GetCopyCommand(ctx, elm)
 		stageImage.Builder.LegacyStapelStageBuilder().Container().AddServiceRunCommands(command)
 
-		imageServiceCommitChangeOptions := stageImage.Builder.LegacyStapelStageBuilder().Container().ServiceCommitChangeOptions()
-
 		checksumLabelKey := imagePkg.WerfImportChecksumLabelPrefix + getImportID(elm)
 		sourceStageIDLabelKey := imagePkg.WerfImportSourceStageIDLabelPrefix + getImportID(elm)
 
@@ -157,6 +156,11 @@ func (s *DependenciesStage) prepareImageWithLegacyStapelBuilder(ctx context.Cont
 				})
 			}
 		}
+
+		depStageID := c.GetStageIDForLastImageStage(s.targetPlatform, dep.ImageName)
+		imageServiceCommitChangeOptions.AddLabel(map[string]string{
+			dependencyLabelKey(depStageID): depStageID,
+		})
 	}
 
 	return nil
@@ -220,6 +224,10 @@ func (s *DependenciesStage) prepareImage(ctx context.Context, c Conveyor, cr con
 				})
 			}
 		}
+		depStageID := c.GetStageIDForLastImageStage(s.targetPlatform, dep.ImageName)
+		stageImage.Builder.StapelStageBuilder().AddLabels(map[string]string{
+			dependencyLabelKey(depStageID): depStageID,
+		})
 	}
 
 	return nil

Diff for: pkg/build/stage/dependencies_utils.go

@@ -1,6 +1,9 @@
 package stage
 
 import (
+	"fmt"
+
+	"github.com/werf/common-go/pkg/util"
 	"github.com/werf/werf/v2/pkg/config"
 	"github.com/werf/werf/v2/pkg/image"
 )
@@ -41,3 +44,7 @@ func ResolveDependenciesArgs(targetPlatform string, dependencies []*config.Depen
 
 	return resolved
 }
+
+func dependencyLabelKey(depStageID string) string {
+	return fmt.Sprintf("%s%s", image.WerfDependencySourceStageIDLabelPrefix, util.Sha256Hash(depStageID))
+}

Diff for: pkg/build/stage/full_dockerfile.go

@@ -606,6 +606,11 @@ func (s *FullDockerfileStage) PrepareImage(ctx context.Context, c Conveyor, cb c
 
 	stageImage.Builder.DockerfileBuilder().AppendLabels(fmt.Sprintf("%s=%s", image.WerfProjectRepoCommitLabel, c.GiterminismManager().HeadCommit()))
 
+	for _, dep := range s.dependencies {
+		depStageID := c.GetStageIDForLastImageStage(s.targetPlatform, dep.ImageName)
+		stageImage.Builder.DockerfileBuilder().AppendLabels(fmt.Sprintf("%s=%s", dependencyLabelKey(depStageID), depStageID))
+	}
+
 	if c.GiterminismManager().Dev() {
 		stageImage.Builder.DockerfileBuilder().AppendLabels(fmt.Sprintf("%s=true", image.WerfDevLabel))
 	}

Diff for: pkg/build/stage/stubs.go

@@ -49,11 +49,15 @@ func (c *LegacyContainerStub) ServiceCommitChangeOptions() container_backend.Leg
 type LegacyContainerOptionsStub struct {
 	container_backend.LegacyContainerOptions
 
-	Env map[string]string
+	Env    map[string]string
+	Labels map[string]string
 }
 
 func NewLegacyContainerOptionsStub() *LegacyContainerOptionsStub {
-	return &LegacyContainerOptionsStub{Env: make(map[string]string)}
+	return &LegacyContainerOptionsStub{
+		Env:    make(map[string]string),
+		Labels: make(map[string]string),
+	}
 }
 
 func (opts *LegacyContainerOptionsStub) AddEnv(envs map[string]string) {
@@ -62,6 +66,12 @@ func (opts *LegacyContainerOptionsStub) AddEnv(envs map[string]string) {
 	}
 }
 
+func (opts *LegacyContainerOptionsStub) AddLabel(labels map[string]string) {
+	for k, v := range labels {
+		opts.Labels[k] = v
+	}
+}
+
 type ConveyorStub struct {
 	Conveyor
 
@@ -104,6 +114,10 @@ func (c *ConveyorStub) GiterminismManager() giterminism_manager.Interface {
 	return c.giterminismManager
 }
 
+func (c *ConveyorStub) GetStageIDForLastImageStage(targetPlatform, imageName string) string {
+	return c.lastStageImageDigestByImageName[imageName]
+}
+
 type GiterminismInspectorStub struct {
 	giterminism_manager.Inspector
 }

Diff for: pkg/cleaning/cleanup.go

@@ -716,9 +716,9 @@ func (m *cleanupManager) cleanupUnusedStages(ctx context.Context) error {
 	{
 		for protectionReason, stageDescToKeepSet := range m.stageManager.GetProtectedStageDescSetByReason() {
 			// Git history based policy keeps import sources more effectively, other policies do not keep them.
-			withImportSources := protectionReason != stage_manager.ProtectionReasonGitPolicy
+			withImportOrDependencySources := protectionReason != stage_manager.ProtectionReasonGitPolicy
 			for stageDescToKeep := range stageDescToKeepSet.Iter() {
-				m.protectRelativeStageDescSetByStageDesc(stageDescToKeep, withImportSources)
+				m.protectRelativeStageDescSetByStageDesc(stageDescToKeep, withImportOrDependencySources)
 			}
 		}
 
@@ -882,7 +882,7 @@ func deleteImportsMetadata(ctx context.Context, projectName string, storageManag
 	})
 }
 
-func (m *cleanupManager) protectRelativeStageDescSetByStageDesc(targetStageDesc *image.StageDesc, withImportSource bool) {
+func (m *cleanupManager) protectRelativeStageDescSetByStageDesc(targetStageDesc *image.StageDesc, withImportOrDependencySources bool) {
 	targetStageDescSet := image.NewStageDescSet()
 	if targetStageDesc.Info.IsIndex {
 		for _, platformImageInfo := range targetStageDesc.Info.Index {
@@ -909,11 +909,11 @@ func (m *cleanupManager) protectRelativeStageDescSetByStageDesc(targetStageDesc
 				handledStageDescSet.Add(currentStageDesc)
 			}
 
-			// Import source checking.
-			if withImportSource {
-				for label, checksum := range currentStageDesc.Info.Labels {
+			// Import or Dependency source checking.
+			if withImportOrDependencySources {
+				for label, value := range currentStageDesc.Info.Labels {
 					if strings.HasPrefix(label, image.WerfImportChecksumLabelPrefix) {
-						sourceStageIDs, ok := m.checksumSourceStageIDs[checksum]
+						sourceStageIDs, ok := m.checksumSourceStageIDs[value]
 						if !ok {
 							continue
 						}
@@ -925,6 +925,12 @@ func (m *cleanupManager) protectRelativeStageDescSetByStageDesc(targetStageDesc
 								m.stageManager.MarkStageDescAsProtected(sourceStageDesc, stage_manager.ProtectionReasonImportSource, false)
 							}
 						}
+					} else if strings.HasPrefix(label, image.WerfDependencySourceStageIDLabelPrefix) {
+						sourceStageDesc := m.stageManager.GetStageDescByStageID(value)
+						if sourceStageDesc != nil {
+							currentStageDescSet.Add(sourceStageDesc)
+							m.stageManager.MarkStageDescAsProtected(sourceStageDesc, stage_manager.ProtectionReasonDependencySource, false)
+						}
 					}
 				}
 			}

Diff for: pkg/cleaning/stage_manager/stages.go

@@ -35,6 +35,7 @@ var (
 	ProtectionReasonGitPolicy                   = newProtectionReason("git policy")
 	ProtectionReasonBuiltWithinLastNHoursPolicy = newProtectionReason("built within last N hours")
 	ProtectionReasonImportSource                = newProtectionReason("import source")
+	ProtectionReasonDependencySource            = newProtectionReason("dependency source")
 	ProtectionReasonAncestor                    = newProtectionReason("ancestor")
 	ProtectionReasonNotFoundInRepo              = newProtectionReason("not found in repo")
 )

Diff for: pkg/image/const.go

@@ -1,20 +1,21 @@
 package image
 
 const (
-	WerfLabelPrefix                    = "werf"
-	WerfLabel                          = "werf"
-	WerfVersionLabel                   = "werf-version"
-	WerfCacheVersionLabel              = "werf-cache-version"
-	WerfImageLabel                     = "werf-image"
-	WerfDevLabel                       = "werf-dev"
-	WerfDockerImageName                = "werf-docker-image-name"
-	WerfStageDigestLabel               = "werf-stage-digest"
-	WerfStageContentDigestLabel        = "werf-stage-content-digest"
-	WerfProjectRepoCommitLabel         = "werf-project-repo-commit"
-	WerfImportChecksumLabelPrefix      = "werf-import-checksum-"
-	WerfImportSourceStageIDLabelPrefix = "werf-import-source-stage-id-"
-	WerfBaseImageIDLabel               = "werf.io/base-image-id"
-	WerfParentStageID                  = "werf.io/parent-stage-id"
+	WerfLabelPrefix                        = "werf"
+	WerfLabel                              = "werf"
+	WerfVersionLabel                       = "werf-version"
+	WerfCacheVersionLabel                  = "werf-cache-version"
+	WerfImageLabel                         = "werf-image"
+	WerfDevLabel                           = "werf-dev"
+	WerfDockerImageName                    = "werf-docker-image-name"
+	WerfStageDigestLabel                   = "werf-stage-digest"
+	WerfStageContentDigestLabel            = "werf-stage-content-digest"
+	WerfProjectRepoCommitLabel             = "werf-project-repo-commit"
+	WerfImportChecksumLabelPrefix          = "werf-import-checksum-"
+	WerfImportSourceStageIDLabelPrefix     = "werf-import-source-stage-id-"
+	WerfDependencySourceStageIDLabelPrefix = "werf-dependency-stage-id-"
+	WerfBaseImageIDLabel                   = "werf.io/base-image-id"
+	WerfParentStageID                      = "werf.io/parent-stage-id"
 
 	WerfImportMetadataChecksumLabel       = "checksum"
 	WerfImportMetadataSourceStageIDLabel  = "source-stage-id"