From 2f0eb543c74e9716ef11f23de118880035ce0906 Mon Sep 17 00:00:00 2001
From: Rahul Rangith <rahul.rangith@datadoghq.com>
Date: Fri, 6 Jan 2023 17:22:26 -0500
Subject: [PATCH] New function for creating missing pvcs

---
 .../statefulset/stateful_pod_control.go          | 16 ++++++++++++++++
 .../statefulset/stateful_set_control.go          |  9 +--------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/pkg/controller/statefulset/stateful_pod_control.go b/pkg/controller/statefulset/stateful_pod_control.go
index 2b0546b4632..8e2dcf8d795 100644
--- a/pkg/controller/statefulset/stateful_pod_control.go
+++ b/pkg/controller/statefulset/stateful_pod_control.go
@@ -315,6 +315,22 @@ func (spc *StatefulPodControl) recordClaimEvent(verb string, set *apps.StatefulS
 	}
 }
 
+// createMissingPersistentVolumeClaims creates all of the required PersistentVolumeClaims for pod, and updates its retention policy
+func (spc *StatefulPodControl) createMissingPersistentVolumeClaims(set *apps.StatefulSet, pod *v1.Pod) error {
+	if err := spc.createPersistentVolumeClaims(set, pod); err != nil {
+		return err
+	}
+
+	if utilfeature.DefaultFeatureGate.Enabled(features.StatefulSetAutoDeletePVC) {
+		// Set PVC policy as much as is possible at this point.
+		if err := spc.UpdatePodClaimForRetentionPolicy(set, pod); err != nil {
+			spc.recordPodEvent("update", set, pod, err)
+			return err
+		}
+	}
+	return nil
+}
+
 // createPersistentVolumeClaims creates all of the required PersistentVolumeClaims for pod, which must be a member of
 // set. If all of the claims for Pod are successfully created, the returned error is nil. If creation fails, this method
 // may be called again until no error is returned, indicating the PersistentVolumeClaims for pod are consistent with
diff --git a/pkg/controller/statefulset/stateful_set_control.go b/pkg/controller/statefulset/stateful_set_control.go
index ea14c300796..276e41ef09f 100644
--- a/pkg/controller/statefulset/stateful_set_control.go
+++ b/pkg/controller/statefulset/stateful_set_control.go
@@ -453,16 +453,9 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 				set.Namespace,
 				set.Name,
 				replicas[i].Name)
-			if err := ssc.podControl.createPersistentVolumeClaims(set, replicas[i]); err != nil {
+			if err := ssc.podControl.createMissingPersistentVolumeClaims(set, replicas[i]); err != nil {
 				return &status, err
 			}
-			if utilfeature.DefaultFeatureGate.Enabled(features.StatefulSetAutoDeletePVC) {
-				// Set PVC policy as much as is possible at this point.
-				if err := ssc.podControl.UpdatePodClaimForRetentionPolicy(set, replicas[i]); err != nil {
-					ssc.podControl.recordPodEvent("update", set, replicas[i], err)
-					return &status, err
-				}
-			}
 		}
 		// If we find a Pod that is currently terminating, we must wait until graceful deletion
 		// completes before we continue to make progress.