Add graceful shutdown when interrupted

If implemented this permit restoring a clean state in case of signal
interruption.

Signed-off-by: Soule BA <soule@weave.works>

cmd/flux/build_kustomization.go

@@ -19,6 +19,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"os/signal"
 
 	"github.com/spf13/cobra"
 
@@ -69,13 +70,31 @@ func buildKsCmdRun(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	manifests, err := builder.Build()
-	if err != nil {
-		return err
+	// create a signal channel
+	sigc := make(chan os.Signal, 1)
+	signal.Notify(sigc, os.Interrupt)
+
+	errChan := make(chan error)
+	go func() {
+		manifests, err := builder.Build()
+		if err != nil {
+			errChan <- err
+		}
+
+		cmd.Print(string(manifests))
+		errChan <- nil
+	}()
+
+	select {
+	case <-sigc:
+		fmt.Println("Build cancelled... exiting.")
+		return builder.Cancel()
+	case err := <-errChan:
+		if err != nil {
+			return err
+		}
 	}
 
-	cmd.Print(string(manifests))
-
 	return nil
 
 }

cmd/flux/diff_kustomization.go

@@ -19,6 +19,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"os/signal"
 
 	"github.com/spf13/cobra"
 
@@ -67,13 +68,31 @@ func diffKsCmdRun(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	output, err := builder.Diff()
-	if err != nil {
-		return err
+	// create a signal channel
+	sigc := make(chan os.Signal, 1)
+	signal.Notify(sigc, os.Interrupt)
+
+	errChan := make(chan error)
+	go func() {
+		output, err := builder.Diff()
+		if err != nil {
+			errChan <- err
+		}
+
+		cmd.Print(output)
+		errChan <- nil
+	}()
+
+	select {
+	case <-sigc:
+		fmt.Println("Build cancelled... exiting.")
+		return builder.Cancel()
+	case err := <-errChan:
+		if err != nil {
+			return err
+		}
 	}
 
-	cmd.Print(output)
-
 	return nil
 
 }

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/fluxcd/image-reflector-controller/api v0.15.0
 	github.com/fluxcd/kustomize-controller/api v0.19.1
 	github.com/fluxcd/notification-controller/api v0.20.1
-	github.com/fluxcd/pkg/apis/kustomize v0.3.1
+	github.com/fluxcd/pkg/apis/kustomize v0.3.1 // indirect
 	github.com/fluxcd/pkg/apis/meta v0.10.2
 	github.com/fluxcd/pkg/runtime v0.12.3
 	github.com/fluxcd/pkg/ssa v0.10.0
@@ -25,7 +25,7 @@ require (
 	github.com/gonvenience/ytbx v1.4.2
 	github.com/google/go-cmp v0.5.6
 	github.com/google/go-containerregistry v0.2.0
-	github.com/hashicorp/go-multierror v1.1.1
+	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/hashicorp/go-retryablehttp v0.7.0 // indirect
 	github.com/homeport/dyff v1.4.6
 	github.com/lucasb-eyer/go-colorful v1.2.0
@@ -49,7 +49,7 @@ require (
 )
 
 require (
-	github.com/drone/envsubst v1.0.3
+	github.com/fluxcd/pkg/kustomize v0.0.2
 	sigs.k8s.io/kustomize/kyaml v0.13.0
 )
 
@@ -72,6 +72,7 @@ require (
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/drone/envsubst/v2 v2.0.0-20210730161058-179042472c46 // indirect
 	github.com/emirpasic/gods v1.12.0 // indirect
 	github.com/evanphx/json-patch v4.12.0+incompatible // indirect
 	github.com/exponent-io/jsonpath v0.0.0-20151013193312-d6023ce2651d // indirect

go.sum

@@ -194,8 +194,8 @@ github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5Xh
 github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZgvJUkLughtfhJv5dyTYa91l1fOUCrgjqmcifM=
 github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
-github.com/drone/envsubst v1.0.3 h1:PCIBwNDYjs50AsLZPYdfhSATKaRg/FJmDc2D6+C2x8g=
-github.com/drone/envsubst v1.0.3/go.mod h1:N2jZmlMufstn1KEqvbHjw40h1KyTmnVzHcSc9bFiJ2g=
+github.com/drone/envsubst/v2 v2.0.0-20210730161058-179042472c46 h1:7QPwrLT79GlD5sizHf27aoY2RTvw62mO6x7mxkScNk0=
+github.com/drone/envsubst/v2 v2.0.0-20210730161058-179042472c46/go.mod h1:esf2rsHFNlZlxsqsZDojNBcnNs5REqIvRrWRHqX0vEU=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/elazarl/goproxy v0.0.0-20180725130230-947c36da3153 h1:yUdfgN0XgIJw7foRItutHYUIhlcKzcSf5vDpdhQAKTc=
@@ -240,6 +240,8 @@ github.com/fluxcd/pkg/apis/kustomize v0.3.1 h1:wmb5D9e1+Rr3/5O3235ERuj+h2VKUArVf
 github.com/fluxcd/pkg/apis/kustomize v0.3.1/go.mod h1:k2HSRd68UwgNmOYBPOd6WbX6a2MH2X/Jeh7e3s3PFPc=
 github.com/fluxcd/pkg/apis/meta v0.10.2 h1:pnDBBEvfs4HaKiVAYgz+e/AQ8dLvcgmVfSeBroZ/KKI=
 github.com/fluxcd/pkg/apis/meta v0.10.2/go.mod h1:KQ2er9xa6koy7uoPMZjIjNudB5p4tXs+w0GO6fRcy7I=
+github.com/fluxcd/pkg/kustomize v0.0.2 h1:ipvQrxSeuGZDsPZrVUL6tYMlTR5xqYTZp6G0Tdy2hVs=
+github.com/fluxcd/pkg/kustomize v0.0.2/go.mod h1:AFwnf3OqQmpTCuwCARTGpPRMBf0ZFJNGCvW63KbgK04=
 github.com/fluxcd/pkg/runtime v0.12.3 h1:h21AZ3YG5MAP7DxFF9hfKrP+vFzys2L7CkUbPFjbP/0=
 github.com/fluxcd/pkg/runtime v0.12.3/go.mod h1:imJ2xYy/d4PbSinX2IefmZk+iS2c1P5fY0js8mCE4SM=
 github.com/fluxcd/pkg/ssa v0.10.0 h1:dhgWDeqz0/zAs5guzmPx/DMPCkzZdlEiPvCs1NChAQM=

internal/kustomization/build.go

@@ -21,10 +21,12 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/fluxcd/flux2/internal/utils"
 	kustomizev1 "github.com/fluxcd/kustomize-controller/api/v1beta2"
+	"github.com/fluxcd/pkg/kustomize"
 	"k8s.io/apimachinery/pkg/api/meta"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -37,9 +39,9 @@ import (
 )
 
 const (
-	controllerName         = "kustomize-controller"
-	controllerGroup        = "kustomize.toolkit.fluxcd.io"
-	mask            string = "**SOPS**"
+	controllerName  = "kustomize-controller"
+	controllerGroup = "kustomize.toolkit.fluxcd.io"
+	mask            = "**SOPS**"
 )
 
 var defaultTimeout = 80 * time.Second
@@ -53,6 +55,9 @@ type Builder struct {
 	name          string
 	namespace     string
 	resourcesPath string
+	// mu is used to synchronize access to the kustomization file
+	mu            sync.Mutex
+	action        kustomize.Action
 	kustomization *kustomizev1.Kustomization
 	timeout       time.Duration
 }
@@ -149,13 +154,15 @@ func (b *Builder) build() (m resmap.ResMap, err error) {
 	// generate kustomization.yaml if needed
 	action, er := b.generate(*k, b.resourcesPath)
 	if er != nil {
-		errf := CleanDirectory(b.resourcesPath, action)
+		errf := kustomize.CleanDirectory(b.resourcesPath, action)
 		err = fmt.Errorf("failed to generate kustomization.yaml: %w", fmt.Errorf("%v %v", er, errf))
 		return
 	}
 
+	b.action = action
+
 	defer func() {
-		errf := CleanDirectory(b.resourcesPath, action)
+		errf := b.Cancel()
 		if err == nil {
 			err = errf
 		}
@@ -185,18 +192,28 @@ func (b *Builder) build() (m resmap.ResMap, err error) {
 
 }
 
-func (b *Builder) generate(kustomization kustomizev1.Kustomization, dirPath string) (action, error) {
+func (b *Builder) generate(kustomization kustomizev1.Kustomization, dirPath string) (kustomize.Action, error) {
 	data, err := runtime.DefaultUnstructuredConverter.ToUnstructured(&kustomization)
 	if err != nil {
 		return "", err
 	}
-	gen := NewGenerator(unstructured.Unstructured{Object: data})
-	return gen.WriteFile(dirPath, WithSaveOriginalKustomization())
+	gen := kustomize.NewGenerator(unstructured.Unstructured{Object: data})
+
+	// acuire the lock
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	return gen.WriteFile(dirPath, kustomize.WithSaveOriginalKustomization())
 }
 
 func (b *Builder) do(ctx context.Context, kustomization kustomizev1.Kustomization, dirPath string) (resmap.ResMap, error) {
 	fs := filesys.MakeFsOnDisk()
-	m, err := BuildKustomization(fs, dirPath)
+
+	// acuire the lock
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	m, err := kustomize.BuildKustomization(fs, dirPath)
 	if err != nil {
 		return nil, fmt.Errorf("kustomize build failed: %w", err)
 	}
@@ -208,7 +225,7 @@ func (b *Builder) do(ctx context.Context, kustomization kustomizev1.Kustomizatio
 			if err != nil {
 				return nil, err
 			}
-			outRes, err := SubstituteVariables(ctx, b.client, unstructured.Unstructured{Object: data}, res)
+			outRes, err := kustomize.SubstituteVariables(ctx, b.client, unstructured.Unstructured{Object: data}, res)
 			if err != nil {
 				return nil, fmt.Errorf("var substitution failed for '%s': %w", res.GetName(), err)
 			}
@@ -263,3 +280,18 @@ func trimSopsData(res *resource.Resource) error {
 
 	return nil
 }
+
+// Cancel cancels the build
+// It restores a clean reprository
+func (b *Builder) Cancel() error {
+	// acuire the lock
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	err := kustomize.CleanDirectory(b.resourcesPath, b.action)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

internal/kustomization/diff.go

@@ -27,7 +27,7 @@ import (
 )
 
 func (b *Builder) Manager() (*ssa.ResourceManager, error) {
-	statusPoller := polling.NewStatusPoller(b.client, b.restMapper)
+	statusPoller := polling.NewStatusPoller(b.client, b.restMapper, nil)
 	owner := ssa.Owner{
 		Field: controllerName,
 		Group: controllerGroup,
@@ -53,8 +53,6 @@ func (b *Builder) Diff() (string, error) {
 		return "", err
 	}
 
-	resourceManager.SetOwnerLabels(objects, b.kustomization.GetName(), b.kustomization.GetNamespace())
-
 	ctx, cancel := context.WithTimeout(context.Background(), b.timeout)
 	defer cancel()