From b9713c4ac402e2aa952101a25d81a49f86a1a483 Mon Sep 17 00:00:00 2001
From: Andy Zhou <andy@andyz.dev>
Date: Mon, 20 May 2019 20:50:01 -0700
Subject: [PATCH] Add --auto-restart flag for the create command

When creating clusters with the --auto-restart flag, any running cluster
will remain "running" up on docker daemon restart.

By default, without this flag, a "running" cluster becomes "stopped"
after docker daemon restart.

Clusters stopped with 'k3d stop' command will remain stopped after
docker daemon restart regardless the settings of this flag.
---
 cli/commands.go  |  2 ++
 cli/container.go | 12 ++++++++++--
 main.go          |  4 ++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/cli/commands.go b/cli/commands.go
index 3290af51..0f9fd389 100644
--- a/cli/commands.go
+++ b/cli/commands.go
@@ -126,6 +126,7 @@ func CreateCluster(c *cli.Context) error {
 		c.String("name"),
 		c.StringSlice("volume"),
 		portmap,
+		c.Bool("auto-restart"),
 	)
 	if err != nil {
 		log.Printf("ERROR: failed to create cluster\n%+v", err)
@@ -196,6 +197,7 @@ func CreateCluster(c *cli.Context) error {
 				c.String("api-port"),
 				portmap,
 				c.Int("port-auto-offset"),
+				c.Bool("auto-restart"),
 			)
 			if err != nil {
 				return fmt.Errorf("ERROR: failed to create worker node for cluster %s\n%+v", c.String("name"), err)
diff --git a/cli/container.go b/cli/container.go
index 58643bae..898f8cee 100644
--- a/cli/container.go
+++ b/cli/container.go
@@ -63,7 +63,7 @@ func startContainer(verbose bool, config *container.Config, hostConfig *containe
 }
 
 func createServer(verbose bool, image string, apiPort string, args []string, env []string,
-	name string, volumes []string, nodeToPortSpecMap map[string][]string) (string, error) {
+	name string, volumes []string, nodeToPortSpecMap map[string][]string, autoRestart bool) (string, error) {
 	log.Printf("Creating server using %s...\n", image)
 
 	containerLabels := make(map[string]string)
@@ -95,6 +95,10 @@ func createServer(verbose bool, image string, apiPort string, args []string, env
 		Privileged:   true,
 	}
 
+	if autoRestart {
+		hostConfig.RestartPolicy.Name = "unless-stopped"
+	}
+
 	if len(volumes) > 0 && volumes[0] != "" {
 		hostConfig.Binds = volumes
 	}
@@ -125,7 +129,7 @@ func createServer(verbose bool, image string, apiPort string, args []string, env
 
 // createWorker creates/starts a k3s agent node that connects to the server
 func createWorker(verbose bool, image string, args []string, env []string, name string, volumes []string,
-	postfix int, serverPort string, nodeToPortSpecMap map[string][]string, portAutoOffset int) (string, error) {
+	postfix int, serverPort string, nodeToPortSpecMap map[string][]string, portAutoOffset int, autoRestart bool) (string, error) {
 	containerLabels := make(map[string]string)
 	containerLabels["app"] = "k3d"
 	containerLabels["component"] = "worker"
@@ -161,6 +165,10 @@ func createWorker(verbose bool, image string, args []string, env []string, name
 		Privileged:   true,
 	}
 
+	if autoRestart {
+		hostConfig.RestartPolicy.Name = "unless-stopped"
+	}
+
 	if len(volumes) > 0 && volumes[0] != "" {
 		hostConfig.Binds = volumes
 	}
diff --git a/main.go b/main.go
index 90f6d30b..c109511e 100644
--- a/main.go
+++ b/main.go
@@ -107,6 +107,10 @@ func main() {
 					Value: 0,
 					Usage: "Specify how many worker nodes you want to spawn",
 				},
+				cli.BoolFlag{
+					Name:  "auto-restart",
+					Usage: "Set docker's --restart=unless-stopped flag on the containers",
+				},
 			},
 			Action: run.CreateCluster,
 		},