From da02653a2afce04db6198fd6e44267dc387d1404 Mon Sep 17 00:00:00 2001
From: Casey Callendrello <casey.callendrello@coreos.com>
Date: Thu, 18 Jan 2018 17:25:50 +0100
Subject: [PATCH] agent: set a label when after successful os-update

This sets the label
`container-linux-update.v1.coreos.com/os-update-staged: "true"`
on nodes where the OS has been upgraded but not yet rebooted. This is so
daemonsets like tectonic-torcx can run ASAP.

Fixes: #167
---
 doc/labels-and-annotations.md |  1 +
 pkg/agent/agent.go            | 11 +++++++++--
 pkg/constants/constants.go    |  1 +
 pkg/k8sutil/metadata.go       |  4 ++--
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/doc/labels-and-annotations.md b/doc/labels-and-annotations.md
index f1619a66..80155c1b 100644
--- a/doc/labels-and-annotations.md
+++ b/doc/labels-and-annotations.md
@@ -30,6 +30,7 @@ A few labels may be set directly by admins to customize behavior. These are call
 | id   | coreos  |  update-agent    | Reflects the ID in `/etc/os-release` |
 | version | 1497.7.0 | update-agent | Reflects the VERSION in `/etc/os-release` |
 | group | stable | update-agent     | Reflects the GROUP in `/usr/share/coreos/update.conf` or `/etc/coreos/update.conf` |
+| reboot-needed | true | update-agent | Set to true to request a coordinated reboot |
 
 **Annotations**
 
diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
index 20a6a9f2..dc6fab09 100644
--- a/pkg/agent/agent.go
+++ b/pkg/agent/agent.go
@@ -92,10 +92,14 @@ func (k *Klocksmith) process(stop <-chan struct{}) error {
 		constants.AnnotationRebootInProgress: constants.False,
 		constants.AnnotationRebootNeeded:     constants.False,
 	}
+	labels := map[string]string{
+		constants.LabelRebootNeeded: constants.False,
+	}
 	glog.Infof("Setting annotations %#v", anno)
-	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
+	if err := k8sutil.SetNodeAnnotationsLabels(k.nc, k.node, anno, labels); err != nil {
 		return err
 	}
+
 	// Since we set 'reboot-needed=false', 'ok-to-reboot' should clear.
 	// Wait for it to do so, else we might start reboot-looping
 	if err := k.waitForNotOkToReboot(); err != nil {
@@ -199,14 +203,17 @@ func (k *Klocksmith) updateStatusCallback(s updateengine.Status) {
 		constants.AnnotationNewVersion:      s.NewVersion,
 	}
 
+	labels := map[string]string{}
+
 	// indicate we need a reboot
 	if s.CurrentOperation == updateengine.UpdateStatusUpdatedNeedReboot {
 		glog.Info("Indicating a reboot is needed")
 		anno[constants.AnnotationRebootNeeded] = constants.True
+		labels[constants.LabelRebootNeeded] = constants.True
 	}
 
 	wait.PollUntil(defaultPollInterval, func() (bool, error) {
-		if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
+		if err := k8sutil.SetNodeAnnotationsLabels(k.nc, k.node, anno, labels); err != nil {
 			glog.Errorf("Failed to set annotation %q: %v", constants.AnnotationStatus, err)
 			return false, nil
 		}
diff --git a/pkg/constants/constants.go b/pkg/constants/constants.go
index cda742a9..f32ff502 100644
--- a/pkg/constants/constants.go
+++ b/pkg/constants/constants.go
@@ -12,6 +12,7 @@ const (
 
 	// Key set to "true" by the update-agent when a reboot is requested.
 	AnnotationRebootNeeded = Prefix + "reboot-needed"
+	LabelRebootNeeded      = Prefix + "reboot-needed"
 
 	// Key set to "true" by the update-agent when node-drain and reboot is
 	// initiated.
diff --git a/pkg/k8sutil/metadata.go b/pkg/k8sutil/metadata.go
index 6ab264c3..b403d4bc 100644
--- a/pkg/k8sutil/metadata.go
+++ b/pkg/k8sutil/metadata.go
@@ -80,8 +80,8 @@ func SetNodeAnnotations(nc v1core.NodeInterface, node string, m map[string]strin
 	})
 }
 
-// SetNodeAnnotationsLabels sets all keys in a and l to their respective values in
-// node's annotations and labels, likewise
+// SetNodeAnnotationsLabels sets all keys in a and l to their values in
+// node's annotations and labels, respectively
 func SetNodeAnnotationsLabels(nc v1core.NodeInterface, node string, a, l map[string]string) error {
 	return UpdateNodeRetry(nc, node, func(n *v1api.Node) {
 		for k, v := range a {