From b8364fcac6385100190118f3c13162b80a3b55e6 Mon Sep 17 00:00:00 2001
From: Jason Hall <jasonhall@google.com>
Date: Wed, 4 Nov 2020 11:10:54 -0500
Subject: [PATCH] Use PATCH to update ready annotation

This should be faster and more reliable than the get-modify-write we use
today.
---
 pkg/pod/entrypoint.go      | 41 ++++++++++++++++++++++----------------
 pkg/pod/entrypoint_test.go | 30 +++++++++++++++++++++-------
 2 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/pkg/pod/entrypoint.go b/pkg/pod/entrypoint.go
index 4df660b0d3e..1e287c7fd28 100644
--- a/pkg/pod/entrypoint.go
+++ b/pkg/pod/entrypoint.go
@@ -18,14 +18,18 @@ package pod
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"log"
 	"path/filepath"
 	"strings"
 
 	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1beta1"
+	"gomodules.xyz/jsonpatch/v2"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes"
 )
 
@@ -165,26 +169,29 @@ func collectResultsName(results []v1beta1.TaskResult) string {
 	return strings.Join(resultNames, ",")
 }
 
-// UpdateReady updates the Pod's annotations to signal the first step to start
-// by projecting the ready annotation via the Downward API.
-func UpdateReady(ctx context.Context, kubeclient kubernetes.Interface, pod corev1.Pod) error {
-	newPod, err := kubeclient.CoreV1().Pods(pod.Namespace).Get(ctx, pod.Name, metav1.GetOptions{})
+var replaceReadyPatchBytes []byte
+
+func init() {
+	// https://stackoverflow.com/questions/55573724/create-a-patch-to-add-a-kubernetes-annotation
+	readyAnnotationPath := "/metadata/annotations/" + strings.Replace(readyAnnotation, "/", "~1", 1)
+	var err error
+	replaceReadyPatchBytes, err = json.Marshal([]jsonpatch.JsonPatchOperation{{
+		Operation: "replace",
+		Path:      readyAnnotationPath,
+		Value:     readyAnnotationValue,
+	}})
 	if err != nil {
-		return fmt.Errorf("error getting Pod %q when updating ready annotation: %w", pod.Name, err)
+		log.Fatalf("failed to marshal replace ready patch bytes: %v", err)
 	}
+}
 
-	// Update the Pod's "READY" annotation to signal the first step to
-	// start.
-	if newPod.ObjectMeta.Annotations == nil {
-		newPod.ObjectMeta.Annotations = map[string]string{}
-	}
-	if newPod.ObjectMeta.Annotations[readyAnnotation] != readyAnnotationValue {
-		newPod.ObjectMeta.Annotations[readyAnnotation] = readyAnnotationValue
-		if _, err := kubeclient.CoreV1().Pods(newPod.Namespace).Update(ctx, newPod, metav1.UpdateOptions{}); err != nil {
-			return fmt.Errorf("error adding ready annotation to Pod %q: %w", pod.Name, err)
-		}
-	}
-	return nil
+// UpdateReady updates the Pod's annotations to signal the first step to start
+// by projecting the ready annotation via the Downward API.
+func UpdateReady(ctx context.Context, kubeclient kubernetes.Interface, pod corev1.Pod) error {
+	// PATCH the Pod's annotations to replace the ready annotation with the
+	// "READY" value, to signal the first step to start.
+	_, err := kubeclient.CoreV1().Pods(pod.Namespace).Patch(ctx, pod.Name, types.JSONPatchType, replaceReadyPatchBytes, metav1.PatchOptions{})
+	return err
 }
 
 // StopSidecars updates sidecar containers in the Pod to a nop image, which
diff --git a/pkg/pod/entrypoint_test.go b/pkg/pod/entrypoint_test.go
index 8c83d3efa53..14e12e5b511 100644
--- a/pkg/pod/entrypoint_test.go
+++ b/pkg/pod/entrypoint_test.go
@@ -263,24 +263,38 @@ func TestUpdateReady(t *testing.T) {
 		desc            string
 		pod             corev1.Pod
 		wantAnnotations map[string]string
+		wantErr         bool
 	}{{
-		desc: "Pod without any annotations has it added",
+		desc: "Pod without any annotations fails",
 		pod: corev1.Pod{
 			ObjectMeta: metav1.ObjectMeta{
 				Name:        "pod",
 				Annotations: nil,
 			},
 		},
+		wantErr: true, // Nothing to replace.
+	}, {
+		desc: "Pod without ready annotation adds it",
+		pod: corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "pod",
+				Annotations: map[string]string{
+					"something": "else",
+				},
+			},
+		},
 		wantAnnotations: map[string]string{
-			readyAnnotation: readyAnnotationValue,
+			"something":     "else",
+			readyAnnotation: "READY",
 		},
 	}, {
-		desc: "Pod with existing annotations has it appended",
+		desc: "Pod with empty annotation value has it replaced",
 		pod: corev1.Pod{
 			ObjectMeta: metav1.ObjectMeta{
 				Name: "pod",
 				Annotations: map[string]string{
-					"something": "else",
+					"something":     "else",
+					readyAnnotation: "",
 				},
 			},
 		},
@@ -289,16 +303,18 @@ func TestUpdateReady(t *testing.T) {
 			readyAnnotation: readyAnnotationValue,
 		},
 	}, {
-		desc: "Pod with other annotation value has it updated",
+		desc: "Pod with other annotation value has it replaced",
 		pod: corev1.Pod{
 			ObjectMeta: metav1.ObjectMeta{
 				Name: "pod",
 				Annotations: map[string]string{
+					"something":     "else",
 					readyAnnotation: "something else",
 				},
 			},
 		},
 		wantAnnotations: map[string]string{
+			"something":     "else",
 			readyAnnotation: readyAnnotationValue,
 		},
 	}} {
@@ -307,8 +323,8 @@ func TestUpdateReady(t *testing.T) {
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
 			kubeclient := fakek8s.NewSimpleClientset(&c.pod)
-			if err := UpdateReady(ctx, kubeclient, c.pod); err != nil {
-				t.Errorf("UpdateReady: %v", err)
+			if err := UpdateReady(ctx, kubeclient, c.pod); (err != nil) != c.wantErr {
+				t.Errorf("UpdateReady (wantErr=%t): %v", c.wantErr, err)
 			}
 
 			got, err := kubeclient.CoreV1().Pods(c.pod.Namespace).Get(ctx, c.pod.Name, metav1.GetOptions{})