From db4de12767ed9afc6ee022e6c2ea8415a7449c1d Mon Sep 17 00:00:00 2001
From: James Shubin <james@shubin.ca>
Date: Mon, 29 Aug 2016 02:54:04 -0400
Subject: [PATCH] Add more flexibility around the prefixes available

This allows you to specify a custom prefix, or a tmp prefix which is
chosen automatically.
---
 DOCUMENTATION.md | 18 ++++++++++++++++++
 main.go          | 23 ++++++++++++++++++++---
 test/shell/t2.sh |  2 +-
 test/shell/t3.sh |  6 +++---
 test/shell/t4.sh |  2 +-
 test/shell/t5.sh |  2 +-
 test/shell/t6.sh |  2 +-
 test/shell/t7.sh |  2 +-
 test/shell/t8.sh |  4 ++--
 9 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md
index db06759f..c99ed0e2 100644
--- a/DOCUMENTATION.md
+++ b/DOCUMENTATION.md
@@ -249,6 +249,24 @@ Point to a graph file to run on the remote host specified within. This parameter
 can be used multiple times if you'd like to remotely run on multiple hosts in
 parallel.
 
+####`--prefix <path>`
+Specify a path to a custom working directory prefix. This directory will get
+created if it does not exist. This usually defaults to `/var/lib/mgmt/`. This
+can't be combined with the `--tmp-prefix` option. It can be combined with the
+`--allow-tmp-prefix` option.
+
+####`--tmp-prefix`
+If this option is specified, a temporary prefix will be used instead of the
+default prefix. This can't be combined with the `--prefix` option.
+
+####`--allow-tmp-prefix`
+If this option is specified, we will attempt to fall back to a temporary prefix
+if the primary prefix couldn't be created. This is useful for avoiding failures
+in environments where the primary prefix may or may not be available, but you'd
+like to try. The canonical example is when running `mgmt` with `--remote` there
+might be a cached copy of the binary in the primary prefix, but in case there's
+no binary available continue working in a temporary directory to avoid failure.
+
 ##Examples
 For example configurations, please consult the [examples/](https://github.com/purpleidea/mgmt/tree/master/examples) directory in the git
 source repository. It is available from:
diff --git a/main.go b/main.go
index 254d4a35..7bb339b5 100644
--- a/main.go
+++ b/main.go
@@ -122,10 +122,18 @@ func run(c *cli.Context) error {
 		return cli.NewExitError("", 1)
 	}
 
+	if c.IsSet("prefix") && c.Bool("tmp-prefix") {
+		log.Println("Main: Error: combining --prefix and the request for a tmp prefix is illogical!")
+		return cli.NewExitError("", 1)
+	}
+	if s := c.String("prefix"); c.IsSet("prefix") && s != "" {
+		prefix = s
+	}
+
 	// make sure the working directory prefix exists
-	if err := os.MkdirAll(prefix, 0770); err != nil {
-		if c.Bool("allow-tmp-prefix") {
-			if prefix, err = ioutil.TempDir("", program); err != nil {
+	if c.Bool("tmp-prefix") || os.MkdirAll(prefix, 0770) != nil {
+		if c.Bool("tmp-prefix") || c.Bool("allow-tmp-prefix") {
+			if prefix, err = ioutil.TempDir("", program+"-"); err != nil {
 				log.Printf("Main: Error: Can't create temporary prefix!")
 				return cli.NewExitError("", 1)
 			}
@@ -487,6 +495,15 @@ func main() {
 					Name:  "no-caching",
 					Usage: "don't allow remote caching of remote execution binary",
 				},
+				cli.StringFlag{
+					Name:   "prefix",
+					Usage:  "specify a path to the working prefix directory",
+					EnvVar: "MGMT_PREFIX",
+				},
+				cli.BoolFlag{
+					Name:  "tmp-prefix",
+					Usage: "request a pseudo-random, temporary prefix to be used",
+				},
 				cli.BoolFlag{
 					Name:  "allow-tmp-prefix",
 					Usage: "allow creation of a new temporary prefix if main prefix is unavailable",
diff --git a/test/shell/t2.sh b/test/shell/t2.sh
index d0d76647..c1a16be1 100755
--- a/test/shell/t2.sh
+++ b/test/shell/t2.sh
@@ -7,7 +7,7 @@ if env | grep -q -e '^TRAVIS=true$'; then
 fi
 
 # run till completion
-timeout --kill-after=15s 10s ./mgmt run --file t2.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=15s 10s ./mgmt run --file t2.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid=$!
 wait $pid	# get exit status
 e=$?
diff --git a/test/shell/t3.sh b/test/shell/t3.sh
index f79c4f50..b15e7394 100755
--- a/test/shell/t3.sh
+++ b/test/shell/t3.sh
@@ -10,11 +10,11 @@ fi
 mkdir -p "${MGMT_TMPDIR}"mgmt{A..C}
 
 # run till completion
-timeout --kill-after=15s 10s ./mgmt run --file t3-a.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=15s 10s ./mgmt run --file t3-a.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid1=$!
-timeout --kill-after=15s 10s ./mgmt run --file t3-b.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=15s 10s ./mgmt run --file t3-b.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid2=$!
-timeout --kill-after=15s 10s ./mgmt run --file t3-c.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=15s 10s ./mgmt run --file t3-c.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid3=$!
 
 wait $pid1	# get exit status
diff --git a/test/shell/t4.sh b/test/shell/t4.sh
index 6cca0cd5..d0f4cc6d 100755
--- a/test/shell/t4.sh
+++ b/test/shell/t4.sh
@@ -1,7 +1,7 @@
 #!/bin/bash -e
 
 # should take slightly more than 25s, but fail if we take 35s)
-timeout --kill-after=35s 30s ./mgmt run --file t4.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=35s 30s ./mgmt run --file t4.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid=$!
 wait $pid	# get exit status
 exit $?
diff --git a/test/shell/t5.sh b/test/shell/t5.sh
index 79d53616..f7d9506a 100755
--- a/test/shell/t5.sh
+++ b/test/shell/t5.sh
@@ -1,7 +1,7 @@
 #!/bin/bash -e
 
 # should take slightly more than 35s, but fail if we take 45s)
-timeout --kill-after=45s 40s ./mgmt run --file t5.yaml --converged-timeout=5 --no-watch --allow-tmp-prefix &
+timeout --kill-after=45s 40s ./mgmt run --file t5.yaml --converged-timeout=5 --no-watch --tmp-prefix &
 pid=$!
 wait $pid	# get exit status
 exit $?
diff --git a/test/shell/t6.sh b/test/shell/t6.sh
index c8805df7..2ee21024 100755
--- a/test/shell/t6.sh
+++ b/test/shell/t6.sh
@@ -7,7 +7,7 @@ if env | grep -q -e '^TRAVIS=true$'; then
 fi
 
 # run till completion
-timeout --kill-after=20s 15s ./mgmt run --file t6.yaml --no-watch --allow-tmp-prefix &
+timeout --kill-after=20s 15s ./mgmt run --file t6.yaml --no-watch --tmp-prefix &
 pid=$!
 sleep 1s	# let it converge
 test -e /tmp/mgmt/f1
diff --git a/test/shell/t7.sh b/test/shell/t7.sh
index 2a590905..655a05d6 100755
--- a/test/shell/t7.sh
+++ b/test/shell/t7.sh
@@ -1,7 +1,7 @@
 #!/bin/bash -e
 
 # run empty graph
-timeout --kill-after=20s 15s ./mgmt run --allow-tmp-prefix &
+timeout --kill-after=20s 15s ./mgmt run --tmp-prefix &
 pid=$!
 sleep 5s	# let it converge
 $(sleep 3s && killall -SIGINT mgmt)&	# send ^C to exit mgmt
diff --git a/test/shell/t8.sh b/test/shell/t8.sh
index 16e925e0..6db30e3b 100755
--- a/test/shell/t8.sh
+++ b/test/shell/t8.sh
@@ -1,11 +1,11 @@
 #!/bin/bash -e
 
 # run empty graphs, we're just testing etcd clustering
-timeout --kill-after=120s 90s ./mgmt run --hostname h1 --allow-tmp-prefix &
+timeout --kill-after=120s 90s ./mgmt run --hostname h1 --tmp-prefix &
 pid1=$!
 sleep 5s	# let it startup
 
-timeout --kill-after=120s 90s ./mgmt run --hostname h2 --seeds http://127.0.0.1:2379 --client-urls http://127.0.0.1:2381 --server-urls http://127.0.0.1:2382 --allow-tmp-prefix &
+timeout --kill-after=120s 90s ./mgmt run --hostname h2 --seeds http://127.0.0.1:2379 --client-urls http://127.0.0.1:2381 --server-urls http://127.0.0.1:2382 --tmp-prefix &
 pid2=$!
 sleep 5s