PodDisruptionBudget should use ControllerRef

pkg/controller/disruption/disruption.go

@@ -35,6 +35,8 @@ import (
 	"k8s.io/kubernetes/pkg/api"
 	"k8s.io/kubernetes/pkg/api/v1"
 	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
+	apps "k8s.io/kubernetes/pkg/apis/apps/v1beta1"
+	"k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
 	policy "k8s.io/kubernetes/pkg/apis/policy/v1beta1"
 	"k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
 	policyclientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset/typed/policy/v1beta1"
@@ -173,94 +175,119 @@ func (dc *DisruptionController) finders() []podControllerFinder {
 		dc.getPodStatefulSets}
 }
 
+var (
+	controllerKindRS  = v1beta1.SchemeGroupVersion.WithKind("ReplicaSet")
+	controllerKindSS  = apps.SchemeGroupVersion.WithKind("StatefulSet")
+	controllerKindRC  = v1.SchemeGroupVersion.WithKind("ReplicationController")
+	controllerKindDep = v1beta1.SchemeGroupVersion.WithKind("Deployment")
+)
+
 // getPodReplicaSets finds replicasets which have no matching deployments.
 func (dc *DisruptionController) getPodReplicaSets(pod *v1.Pod) ([]controllerAndScale, error) {
-	cas := []controllerAndScale{}
-	rss, err := dc.rsLister.GetPodReplicaSets(pod)
-	// GetPodReplicaSets returns an error only if no ReplicaSets are found.  We
-	// don't return that as an error to the caller.
+	var casSlice []controllerAndScale
+	controllerRef := controller.GetControllerOf(pod)
+	if controllerRef == nil {
+		return nil, nil
+	}
+	if controllerRef.Kind != controllerKindRS.Kind {
+		return nil, nil
+	}
+	rs, err := dc.rsLister.ReplicaSets(pod.Namespace).Get(controllerRef.Name)
 	if err != nil {
-		return cas, nil
+		// The only possible error is NotFound, which is ok here.
+		return nil, nil
 	}
-	controllerScale := map[types.UID]int32{}
-	for _, rs := range rss {
-		// GetDeploymentsForReplicaSet returns an error only if no matching
-		// deployments are found.
-		_, err := dc.dLister.GetDeploymentsForReplicaSet(rs)
-		if err == nil { // A deployment was found, so this finder will not count this RS.
-			continue
-		}
-		controllerScale[rs.UID] = *(rs.Spec.Replicas)
+	if rs.UID != controllerRef.UID {
+		return nil, nil
 	}
-
-	for uid, scale := range controllerScale {
-		cas = append(cas, controllerAndScale{UID: uid, scale: scale})
+	controllerRef = controller.GetControllerOf(rs)
+	if controllerRef != nil && controllerRef.Kind == controllerKindDep.Kind {
+		// Skip RS if it's controlled by a Deployment.
+		return nil, nil
 	}
-
-	return cas, nil
+	casSlice = append(casSlice, controllerAndScale{rs.UID, *(rs.Spec.Replicas)})
+	return casSlice, nil
 }
 
 // getPodStatefulSet returns the statefulset managing the given pod.
 func (dc *DisruptionController) getPodStatefulSets(pod *v1.Pod) ([]controllerAndScale, error) {
-	cas := []controllerAndScale{}
-	ss, err := dc.ssLister.GetPodStatefulSets(pod)
-
-	// GetPodStatefulSets returns an error only if no StatefulSets are found. We
-	// don't return that as an error to the caller.
+	var casSlice []controllerAndScale
+	controllerRef := controller.GetControllerOf(pod)
+	if controllerRef == nil {
+		return nil, nil
+	}
+	if controllerRef.Kind != controllerKindSS.Kind {
+		return nil, nil
+	}
+	ss, err := dc.ssLister.StatefulSets(pod.Namespace).Get(controllerRef.Name)
 	if err != nil {
-		return cas, nil
+		// The only possible error is NotFound, which is ok here.
+		return nil, nil
+	}
+	if ss.UID != controllerRef.UID {
+		return nil, nil
 	}
 
-	controllerScale := map[types.UID]int32{}
-	for _, s := range ss {
-		controllerScale[s.UID] = *(s.Spec.Replicas)
-	}
-
-	for uid, scale := range controllerScale {
-		cas = append(cas, controllerAndScale{UID: uid, scale: scale})
-	}
-
-	return cas, nil
+	casSlice = append(casSlice, controllerAndScale{ss.UID, *(ss.Spec.Replicas)})
+	return casSlice, nil
 }
 
 // getPodDeployments finds deployments for any replicasets which are being managed by deployments.
 func (dc *DisruptionController) getPodDeployments(pod *v1.Pod) ([]controllerAndScale, error) {
-	cas := []controllerAndScale{}
-	rss, err := dc.rsLister.GetPodReplicaSets(pod)
-	// GetPodReplicaSets returns an error only if no ReplicaSets are found.  We
-	// don't return that as an error to the caller.
+	var casSlice []controllerAndScale
+	controllerRef := controller.GetControllerOf(pod)
+	if controllerRef == nil {
+		return nil, nil
+	}
+	if controllerRef.Kind != controllerKindRS.Kind {
+		return nil, nil
+	}
+	rs, err := dc.rsLister.ReplicaSets(pod.Namespace).Get(controllerRef.Name)
 	if err != nil {
-		return cas, nil
+		// The only possible error is NotFound, which is ok here.
+		return nil, nil
 	}
-	controllerScale := map[types.UID]int32{}
-	for _, rs := range rss {
-		ds, err := dc.dLister.GetDeploymentsForReplicaSet(rs)
-		// GetDeploymentsForReplicaSet returns an error only if no matching
-		// deployments are found.  In that case we skip this ReplicaSet.
-		if err != nil {
-			continue
-		}
-		for _, d := range ds {
-			controllerScale[d.UID] = *(d.Spec.Replicas)
-		}
+	if rs.UID != controllerRef.UID {
+		return nil, nil
 	}
-
-	for uid, scale := range controllerScale {
-		cas = append(cas, controllerAndScale{UID: uid, scale: scale})
+	controllerRef = controller.GetControllerOf(rs)
+	if controllerRef == nil {
+		return nil, nil
 	}
-
-	return cas, nil
+	if controllerRef.Kind != controllerKindDep.Kind {
+		return nil, nil
+	}
+	deployment, err := dc.dLister.Deployments(rs.Namespace).Get(controllerRef.Name)
+	if err != nil {
+		// The only possible error is NotFound, which is ok here.
+		return nil, nil
+	}
+	if deployment.UID != controllerRef.UID {
+		return nil, nil
+	}
+	casSlice = append(casSlice, controllerAndScale{deployment.UID, *(deployment.Spec.Replicas)})
+	return casSlice, nil
 }
 
 func (dc *DisruptionController) getPodReplicationControllers(pod *v1.Pod) ([]controllerAndScale, error) {
-	cas := []controllerAndScale{}
-	rcs, err := dc.rcLister.GetPodControllers(pod)
-	if err == nil {
-		for _, rc := range rcs {
-			cas = append(cas, controllerAndScale{UID: rc.UID, scale: *(rc.Spec.Replicas)})
-		}
+	var casSlice []controllerAndScale
+	controllerRef := controller.GetControllerOf(pod)
+	if controllerRef == nil {
+		return nil, nil
 	}
-	return cas, nil
+	if controllerRef.Kind != controllerKindRC.Kind {
+		return nil, nil
+	}
+	rc, err := dc.rcLister.ReplicationControllers(pod.Namespace).Get(controllerRef.Name)
+	if err != nil {
+		// The only possible error is NotFound, which is ok here.
+		return nil, nil
+	}
+	if rc.UID != controllerRef.UID {
+		return nil, nil
+	}
+	casSlice = append(casSlice, controllerAndScale{rc.UID, *(rc.Spec.Replicas)})
+	return casSlice, nil
 }
 
 func (dc *DisruptionController) Run(stopCh <-chan struct{}) {