From 4e31380c33f94b2cf4099dc4f26db3d1a7f25f88 Mon Sep 17 00:00:00 2001
From: wuqixuan <wuqixuan@huawei.com>
Date: Mon, 6 Jul 2015 20:10:50 +0800
Subject: [PATCH] fleetctl: support overwrite command

If unit file modified, can use "overwrite" command to update.
First stop and remove the old unit, and submit the new unit and restore
the status as the previous one.

Fixes #76
---
 fleetctl/fleetctl.go  |   7 ++-
 fleetctl/overwrite.go | 114 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 119 insertions(+), 2 deletions(-)
 create mode 100644 fleetctl/overwrite.go

diff --git a/fleetctl/fleetctl.go b/fleetctl/fleetctl.go
index b71a9baa9..702918ba4 100644
--- a/fleetctl/fleetctl.go
+++ b/fleetctl/fleetctl.go
@@ -175,6 +175,7 @@ func init() {
 		cmdUnloadUnit,
 		cmdVerifyUnit,
 		cmdVersion,
+		cmdOverwriteUnit,
 	}
 }
 
@@ -664,13 +665,13 @@ func lazyCreateUnits(args []string) error {
 	return nil
 }
 
-func warnOnDifferentLocalUnit(loc string, su *schema.Unit) {
+func warnOnDifferentLocalUnit(loc string, su *schema.Unit) bool {
 	suf := schema.MapSchemaUnitOptionsToUnitFile(su.Options)
 	if _, err := os.Stat(loc); !os.IsNotExist(err) {
 		luf, err := getUnitFromFile(loc)
 		if err == nil && luf.Hash() != suf.Hash() {
 			stderr("WARNING: Unit %s in registry differs from local unit file %s", su.Name, loc)
-			return
+			return true
 		}
 	}
 	if uni := unit.NewUnitNameInfo(path.Base(loc)); uni != nil && uni.IsInstance() {
@@ -679,9 +680,11 @@ func warnOnDifferentLocalUnit(loc string, su *schema.Unit) {
 			tmpl, err := getUnitFromFile(file)
 			if err == nil && tmpl.Hash() != suf.Hash() {
 				stderr("WARNING: Unit %s in registry differs from local template unit file %s", su.Name, uni.Template)
+				return true
 			}
 		}
 	}
+	return false
 }
 
 func lazyLoadUnits(args []string) ([]*schema.Unit, error) {
diff --git a/fleetctl/overwrite.go b/fleetctl/overwrite.go
new file mode 100644
index 000000000..415048126
--- /dev/null
+++ b/fleetctl/overwrite.go
@@ -0,0 +1,114 @@
+package main
+
+import (
+	"os"
+
+	"github.com/coreos/fleet/job"
+)
+
+var cmdOverwriteUnit = &Command{
+	Name:    "overwrite",
+	Summary: "Overwrite one or more units in the cluster",
+	Usage:   "UNIT...",
+	Description: `Overwrite one or more running or submitted units from the cluster.
+
+Act as a procedure of stop-->destroy-->commit-->start(if needed) on unit(s)`,
+
+	Run: runOverwriteUnits,
+}
+
+func runOverwriteUnits(args []string) (exit int) {
+	for _, v := range args {
+		v = maybeAppendDefaultUnitType(v)
+		name := unitNameMangle(v)
+
+		u, err := cAPI.Unit(name)
+		if err != nil {
+			stderr("error retrieving Unit(%s) from Registry: %v", name, err)
+			return 1
+		}
+		if u == nil {
+			stdout("Unit(%s) in can not be found in Registry, nothing to do with it", name)
+			continue
+		}
+
+		hash_mismatch := warnOnDifferentLocalUnit(v, u)
+		if hash_mismatch == false {
+			stdout("Nothing different between Unit(%s) in registory and local unit file, just skip", name)
+			continue
+		}
+
+		stopping := make([]string, 0)
+		stopping = append(stopping, u.Name)
+		if job.JobState(u.CurrentState) == job.JobStateLaunched || suToGlobal(*u) {
+			stdout("Stop Unit(%s) first, so that we can overwrite unit file", u.Name)
+			cAPI.SetUnitTargetState(u.Name, string(job.JobStateLoaded))
+			errchan := waitForUnitStates(stopping, job.JobStateLoaded, 0, os.Stdout)
+			for err := range errchan {
+				stderr("Error waiting for units: %v", err)
+				exit = 1
+			}
+		}
+
+		err = cAPI.DestroyUnit(name)
+		if err != nil {
+			continue
+		}
+
+		recreating := make([]string, 0)
+		recreating = append(recreating, v)
+		stdout("Recreating unit(%s)", v)
+		if err := lazyCreateUnits(recreating); err != nil {
+			stderr("Error creating units: %v", err)
+			return
+		}
+		if job.JobState(u.CurrentState) == job.JobStateInactive {
+			continue
+		} else if job.JobState(u.CurrentState) == job.JobStateLoaded {
+			stdout("Rescheduling unit(%s)", v)
+			triggered, err := lazyLoadUnits(recreating)
+			if err != nil {
+				stderr("Error loading units: %v", err)
+				return 1
+			}
+
+			var loading []string
+			for _, u := range triggered {
+				if suToGlobal(*u) {
+					stdout("Triggered global unit %s load", u.Name)
+				} else {
+					loading = append(loading, u.Name)
+				}
+			}
+
+			errchan := waitForUnitStates(loading, job.JobStateLoaded, sharedFlags.BlockAttempts, os.Stdout)
+			for err := range errchan {
+				stderr("Error waiting for units: %v", err)
+			}
+		} else if job.JobState(u.CurrentState) == job.JobStateLaunched {
+			stdout("Restarting unit(%s)", v)
+			triggered, err := lazyStartUnits(recreating)
+			if err != nil {
+				stderr("Error starting units: %v", err)
+				return 1
+			}
+
+			var starting []string
+			for _, u := range triggered {
+				if suToGlobal(*u) {
+					stdout("Triggered global unit %s start", u.Name)
+				} else {
+					starting = append(starting, u.Name)
+				}
+			}
+
+			errchan := waitForUnitStates(starting, job.JobStateLaunched, sharedFlags.BlockAttempts, os.Stdout)
+			for err := range errchan {
+				stderr("Error waiting for units: %v", err)
+				exit = 1
+			}
+		}
+
+	}
+	return
+}