From b5dfc7b2aadb6c2fb8ba5783fecb116e00ebff8b Mon Sep 17 00:00:00 2001
From: Anthony Yeh <enisoc@google.com>
Date: Thu, 23 Feb 2017 12:01:02 -0800
Subject: [PATCH] StatefulSet: Add ControllerRef on all created Pods.

---
 .../statefulset/stateful_set_utils.go         | 15 ++++++++++-
 .../statefulset/stateful_set_utils_test.go    | 25 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/pkg/controller/statefulset/stateful_set_utils.go b/pkg/controller/statefulset/stateful_set_utils.go
index 59811c1e6d4..231ef945675 100644
--- a/pkg/controller/statefulset/stateful_set_utils.go
+++ b/pkg/controller/statefulset/stateful_set_utils.go
@@ -21,6 +21,7 @@ import (
 	"regexp"
 	"strconv"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/kubernetes/pkg/api/v1"
 	podapi "k8s.io/kubernetes/pkg/api/v1/pod"
 	apps "k8s.io/kubernetes/pkg/apis/apps/v1beta1"
@@ -230,9 +231,21 @@ func isHealthy(pod *v1.Pod) bool {
 	return isRunningAndReady(pod) && !isTerminated(pod)
 }
 
+// newControllerRef returns an ControllerRef pointing to a given StatefulSet.
+func newControllerRef(set *apps.StatefulSet) *metav1.OwnerReference {
+	isController := true
+	return &metav1.OwnerReference{
+		APIVersion: apps.SchemeGroupVersion.String(),
+		Kind:       "StatefulSet",
+		Name:       set.Name,
+		UID:        set.UID,
+		Controller: &isController,
+	}
+}
+
 // newStatefulSetPod returns a new Pod conforming to the set's Spec with an identity generated from ordinal.
 func newStatefulSetPod(set *apps.StatefulSet, ordinal int) *v1.Pod {
-	pod, _ := controller.GetPodFromTemplate(&set.Spec.Template, set, nil)
+	pod, _ := controller.GetPodFromTemplate(&set.Spec.Template, set, newControllerRef(set))
 	pod.Name = getPodName(set, ordinal)
 	updateIdentity(set, pod)
 	updateStorage(set, pod)
diff --git a/pkg/controller/statefulset/stateful_set_utils_test.go b/pkg/controller/statefulset/stateful_set_utils_test.go
index b6e80473985..c53cbb10a6d 100644
--- a/pkg/controller/statefulset/stateful_set_utils_test.go
+++ b/pkg/controller/statefulset/stateful_set_utils_test.go
@@ -31,6 +31,7 @@ import (
 	"k8s.io/kubernetes/pkg/api/v1"
 	podapi "k8s.io/kubernetes/pkg/api/v1/pod"
 	apps "k8s.io/kubernetes/pkg/apis/apps/v1beta1"
+	"k8s.io/kubernetes/pkg/controller"
 )
 
 func TestGetParentNameAndOrdinal(t *testing.T) {
@@ -271,6 +272,30 @@ func TestOverlappingStatefulSets(t *testing.T) {
 	}
 }
 
+func TestNewPodControllerRef(t *testing.T) {
+	set := newStatefulSet(1)
+	pod := newStatefulSetPod(set, 0)
+	controllerRef := controller.GetControllerOf(pod)
+	if controllerRef == nil {
+		t.Fatalf("No ControllerRef found on new pod")
+	}
+	if got, want := controllerRef.APIVersion, apps.SchemeGroupVersion.String(); got != want {
+		t.Errorf("controllerRef.APIVersion = %q, want %q", got, want)
+	}
+	if got, want := controllerRef.Kind, "StatefulSet"; got != want {
+		t.Errorf("controllerRef.Kind = %q, want %q", got, want)
+	}
+	if got, want := controllerRef.Name, set.Name; got != want {
+		t.Errorf("controllerRef.Name = %q, want %q", got, want)
+	}
+	if got, want := controllerRef.UID, set.UID; got != want {
+		t.Errorf("controllerRef.UID = %q, want %q", got, want)
+	}
+	if got, want := *controllerRef.Controller, true; got != want {
+		t.Errorf("controllerRef.Controller = %v, want %v", got, want)
+	}
+}
+
 func newPVC(name string) v1.PersistentVolumeClaim {
 	return v1.PersistentVolumeClaim{
 		ObjectMeta: metav1.ObjectMeta{