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
 
 }

internal/kustomization/build.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/fluxcd/flux2/internal/utils"
@@ -53,6 +54,9 @@ type Builder struct {
 	name          string
 	namespace     string
 	resourcesPath string
+	// mu is used to synchronize access to the kustomization file
+	mu            sync.Mutex
+	action        action
 	kustomization *kustomizev1.Kustomization
 	timeout       time.Duration
 }
@@ -154,8 +158,10 @@ func (b *Builder) build() (m resmap.ResMap, err error) {
 		return
 	}
 
+	b.action = action
+
 	defer func() {
-		errf := CleanDirectory(b.resourcesPath, action)
+		errf := b.Cancel()
 		if err == nil {
 			err = errf
 		}
@@ -191,11 +197,21 @@ func (b *Builder) generate(kustomization kustomizev1.Kustomization, dirPath stri
 		return "", err
 	}
 	gen := NewGenerator(unstructured.Unstructured{Object: data})
+
+	// acuire the lock
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
 	return gen.WriteFile(dirPath, WithSaveOriginalKustomization())
 }
 
 func (b *Builder) do(ctx context.Context, kustomization kustomizev1.Kustomization, dirPath string) (resmap.ResMap, error) {
 	fs := filesys.MakeFsOnDisk()
+
+	// acuire the lock
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
 	m, err := BuildKustomization(fs, dirPath)
 	if err != nil {
 		return nil, fmt.Errorf("kustomize build failed: %w", err)
@@ -263,3 +279,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 := CleanDirectory(b.resourcesPath, b.action)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

internal/kustomization/diff.go

@@ -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()