From 9221e28391c238ddd94a1900731335abae242dd0 Mon Sep 17 00:00:00 2001
From: Johan Thomsen <jth@dbc.dk>
Date: Mon, 25 Nov 2019 17:22:48 +0100
Subject: [PATCH] misc: cleanup

- ensure that finalizers are run on all program exists
- register removal of nix result tmpdir as finalizer

- cleanup code in morph.go
-- don't panic
-- re-introduce validateEnvironment
-- simplify some error handling
---
 morph.go           | 73 ++++++++++++++++++++++------------------------
 nix/nix.go         |  7 ++---
 utils/finalizer.go |  8 +++--
 3 files changed, 44 insertions(+), 44 deletions(-)

diff --git a/morph.go b/morph.go
index 7d57882..2aa79e3 100644
--- a/morph.go
+++ b/morph.go
@@ -51,7 +51,7 @@ var (
 	executeCommand      []string
 	keepGCRoot          = app.Flag("keep-result", "Keep latest build in .gcroots to prevent it from being garbage collected").Default("False").Bool()
 
-	assetRoot, assetsErr = assets.Setup()
+	assetRoot string
 )
 
 func deploymentArg(cmd *kingpin.CmdClause) {
@@ -205,15 +205,34 @@ func listSecretsCmd(cmd *kingpin.CmdClause) *kingpin.CmdClause {
 	return cmd
 }
 
-func init() {
-	if err := validateEnvironment(); err != nil {
-		panic(err)
+func setup() {
+	handleError(validateEnvironment())
+
+	utils.AddFinalizer(func() {
+		assets.Teardown(assetRoot)
+	})
+	utils.SignalHandler()
+
+	var assetErr error
+	assetRoot, assetErr = assets.Setup()
+	handleError(assetErr)
+}
+
+func validateEnvironment() (err error) {
+	dependencies := []string{"nix", "scp", "ssh"}
+	missingDepencies := make([]string, 0)
+	for _, dependency := range dependencies {
+		_, err := exec.LookPath(dependency)
+		if err != nil {
+			missingDepencies = append(missingDepencies, dependency)
+		}
 	}
 
-	if assetsErr != nil {
-		fmt.Fprintln(os.Stderr, "Error unpacking assets:")
-		panic(assetsErr)
+	if len(missingDepencies) > 0 {
+		return errors.New("Missing dependencies: " + strings.Join(missingDepencies, ", "))
 	}
+
+	return nil
 }
 
 func main() {
@@ -226,15 +245,10 @@ func main() {
 	}
 
 	defer utils.RunFinalizers()
-	utils.AddFinalizer(func() {
-		assets.Teardown(assetRoot)
-	})
-	utils.SignalHandler()
+	setup()
 
 	hosts, err := getHosts(deployment)
-	if err != nil {
-		handleError(clause, hosts, err)
-	}
+	handleError(err)
 
 	switch clause {
 	case build.FullCommand():
@@ -257,15 +271,15 @@ func main() {
 		err = execExecute(hosts)
 	}
 
-	if err != nil {
-		handleError(clause, hosts, err)
-	}
+	handleError(err)
 }
 
-func handleError(cmd string, hosts []nix.Host, err error) {
+func handleError(err error) {
 	//Stupid handling of catch-all errors for now
-	fmt.Fprint(os.Stderr, err.Error())
-	os.Exit(1)
+	if err != nil {
+		fmt.Fprint(os.Stderr, err.Error())
+		utils.Exit(1)
+	}
 }
 
 func execExecute(hosts []nix.Host) error {
@@ -378,7 +392,7 @@ func execDeploy(hosts []nix.Host) (string, error) {
 			if err != nil {
 				fmt.Fprintln(os.Stderr)
 				fmt.Fprintln(os.Stderr, "Not deploying to additional hosts, since a host health check failed.")
-				os.Exit(1)
+				utils.Exit(1)
 			}
 		}
 
@@ -484,23 +498,6 @@ func execListSecretsAsJson(hosts []nix.Host) error {
 	return nil
 }
 
-func validateEnvironment() (err error) {
-	dependencies := []string{"nix", "scp", "ssh"}
-	missingDepencies := make([]string, 0)
-	for _, dependency := range dependencies {
-		_, err := exec.LookPath(dependency)
-		if err != nil {
-			missingDepencies = append(missingDepencies, dependency)
-		}
-	}
-
-	if len(missingDepencies) > 0 {
-		return errors.New("Missing dependencies: " + strings.Join(missingDepencies, ", "))
-	}
-
-	return nil
-}
-
 func getHosts(deploymentFile string) (hosts []nix.Host, err error) {
 
 	deployment, err := os.Open(deploymentFile)
diff --git a/nix/nix.go b/nix/nix.go
index 982e365..300634a 100644
--- a/nix/nix.go
+++ b/nix/nix.go
@@ -170,7 +170,9 @@ func (ctx *NixContext) BuildMachines(deploymentPath string, hosts []Host, nixArg
 		if err != nil {
 			return "", err
 		}
-		defer os.Remove(tmpdir)
+		utils.AddFinalizer(func() {
+			os.RemoveAll(tmpdir)
+		})
 		resultLinkPath = filepath.Join(tmpdir, "result")
 	}
 	args := []string{ctx.EvalMachines,
@@ -195,9 +197,6 @@ func (ctx *NixContext) BuildMachines(deploymentPath string, hosts []Host, nixArg
 	}
 
 	cmd := exec.Command("nix-build", args...)
-	if !ctx.KeepGCRoot {
-		defer os.Remove(resultLinkPath)
-	}
 
 	// show process output on attached stdout/stderr
 	cmd.Stdout = os.Stderr
diff --git a/utils/finalizer.go b/utils/finalizer.go
index a6eac85..29865bd 100644
--- a/utils/finalizer.go
+++ b/utils/finalizer.go
@@ -46,7 +46,11 @@ func SignalHandler() {
 	go func() {
 		sig := <-sigs
 		fmt.Fprintf(os.Stderr, "Received signal: %s\n", sig.String())
-		RunFinalizers()
-		os.Exit(130) // reserved exit code for "Interrupted"
+		Exit(130) // reserved exit code for "Interrupted"
 	}()
 }
+
+func Exit(exitCode int) {
+	RunFinalizers()
+	os.Exit(exitCode)
+}