Wire contexts to Apps controllers

pkg/controller/statefulset/stateful_pod_control.go

@@ -41,7 +41,7 @@ import (
 type StatefulPodControlInterface interface {
 	// CreateStatefulPod create a Pod in a StatefulSet. Any PVCs necessary for the Pod are created prior to creating
 	// the Pod. If the returned error is nil the Pod and its PVCs have been created.
-	CreateStatefulPod(set *apps.StatefulSet, pod *v1.Pod) error
+	CreateStatefulPod(ctx context.Context, set *apps.StatefulSet, pod *v1.Pod) error
 	// UpdateStatefulPod Updates a Pod in a StatefulSet. If the Pod already has the correct identity and stable
 	// storage this method is a no-op. If the Pod must be mutated to conform to the Set, it is mutated and updated.
 	// pod is an in-out parameter, and any updates made to the pod are reflected as mutations to this parameter. If
@@ -72,14 +72,14 @@ type realStatefulPodControl struct {
 	recorder  record.EventRecorder
 }
 
-func (spc *realStatefulPodControl) CreateStatefulPod(set *apps.StatefulSet, pod *v1.Pod) error {
+func (spc *realStatefulPodControl) CreateStatefulPod(ctx context.Context, set *apps.StatefulSet, pod *v1.Pod) error {
 	// Create the Pod's PVCs prior to creating the Pod
 	if err := spc.createPersistentVolumeClaims(set, pod); err != nil {
 		spc.recordPodEvent("create", set, pod, err)
 		return err
 	}
 	// If we created the PVCs attempt to create the Pod
-	_, err := spc.client.CoreV1().Pods(set.Namespace).Create(context.TODO(), pod, metav1.CreateOptions{})
+	_, err := spc.client.CoreV1().Pods(set.Namespace).Create(ctx, pod, metav1.CreateOptions{})
 	// sink already exists errors
 	if apierrors.IsAlreadyExists(err) {
 		return err

pkg/controller/statefulset/stateful_pod_control_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package statefulset
 
 import (
+	"context"
 	"errors"
 	"strings"
 	"testing"
@@ -56,7 +57,7 @@ func TestStatefulPodControlCreatesPods(t *testing.T) {
 		create := action.(core.CreateAction)
 		return true, create.GetObject(), nil
 	})
-	if err := control.CreateStatefulPod(set, pod); err != nil {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); err != nil {
 		t.Errorf("StatefulPodControl failed to create Pod error: %s", err)
 	}
 	events := collectEvents(recorder.Events)
@@ -90,7 +91,7 @@ func TestStatefulPodControlCreatePodExists(t *testing.T) {
 	fakeClient.AddReactor("create", "pods", func(action core.Action) (bool, runtime.Object, error) {
 		return true, pod, apierrors.NewAlreadyExists(action.GetResource().GroupResource(), pod.Name)
 	})
-	if err := control.CreateStatefulPod(set, pod); !apierrors.IsAlreadyExists(err) {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); !apierrors.IsAlreadyExists(err) {
 		t.Errorf("Failed to create Pod error: %s", err)
 	}
 	events := collectEvents(recorder.Events)
@@ -117,7 +118,7 @@ func TestStatefulPodControlCreatePodPvcCreateFailure(t *testing.T) {
 		create := action.(core.CreateAction)
 		return true, create.GetObject(), nil
 	})
-	if err := control.CreateStatefulPod(set, pod); err == nil {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); err == nil {
 		t.Error("Failed to produce error on PVC creation failure")
 	}
 	events := collectEvents(recorder.Events)
@@ -153,7 +154,7 @@ func TestStatefulPodControlCreatePodPvcDeleting(t *testing.T) {
 		create := action.(core.CreateAction)
 		return true, create.GetObject(), nil
 	})
-	if err := control.CreateStatefulPod(set, pod); err == nil {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); err == nil {
 		t.Error("Failed to produce error on deleting PVC")
 	}
 	events := collectEvents(recorder.Events)
@@ -191,7 +192,7 @@ func TestStatefulPodControlCreatePodPvcGetFailure(t *testing.T) {
 		create := action.(core.CreateAction)
 		return true, create.GetObject(), nil
 	})
-	if err := control.CreateStatefulPod(set, pod); err == nil {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); err == nil {
 		t.Error("Failed to produce error on PVC creation failure")
 	}
 	events := collectEvents(recorder.Events)
@@ -220,7 +221,7 @@ func TestStatefulPodControlCreatePodFailed(t *testing.T) {
 	fakeClient.AddReactor("create", "pods", func(action core.Action) (bool, runtime.Object, error) {
 		return true, nil, apierrors.NewInternalError(errors.New("API server down"))
 	})
-	if err := control.CreateStatefulPod(set, pod); err == nil {
+	if err := control.CreateStatefulPod(context.TODO(), set, pod); err == nil {
 		t.Error("Failed to produce error on Pod creation failure")
 	}
 	events := collectEvents(recorder.Events)

pkg/controller/statefulset/stateful_set.go

@@ -141,22 +141,22 @@ func NewStatefulSetController(
 }
 
 // Run runs the statefulset controller.
-func (ssc *StatefulSetController) Run(workers int, stopCh <-chan struct{}) {
+func (ssc *StatefulSetController) Run(ctx context.Context, workers int) {
 	defer utilruntime.HandleCrash()
 	defer ssc.queue.ShutDown()
 
 	klog.Infof("Starting stateful set controller")
 	defer klog.Infof("Shutting down statefulset controller")
 
-	if !cache.WaitForNamedCacheSync("stateful set", stopCh, ssc.podListerSynced, ssc.setListerSynced, ssc.pvcListerSynced, ssc.revListerSynced) {
+	if !cache.WaitForNamedCacheSync("stateful set", ctx.Done(), ssc.podListerSynced, ssc.setListerSynced, ssc.pvcListerSynced, ssc.revListerSynced) {
 		return
 	}
 
 	for i := 0; i < workers; i++ {
-		go wait.Until(ssc.worker, time.Second, stopCh)
+		go wait.UntilWithContext(ctx, ssc.worker, time.Second)
 	}
 
-	<-stopCh
+	<-ctx.Done()
 }
 
 // addPod adds the statefulset for the pod to the sync queue
@@ -287,7 +287,7 @@ func (ssc *StatefulSetController) deletePod(obj interface{}) {
 //
 // NOTE: Returned Pods are pointers to objects from the cache.
 //       If you need to modify one, you need to copy it first.
-func (ssc *StatefulSetController) getPodsForStatefulSet(set *apps.StatefulSet, selector labels.Selector) ([]*v1.Pod, error) {
+func (ssc *StatefulSetController) getPodsForStatefulSet(ctx context.Context, set *apps.StatefulSet, selector labels.Selector) ([]*v1.Pod, error) {
 	// List all pods to include the pods that don't match the selector anymore but
 	// has a ControllerRef pointing to this StatefulSet.
 	pods, err := ssc.podLister.Pods(set.Namespace).List(labels.Everything())
@@ -300,15 +300,15 @@ func (ssc *StatefulSetController) getPodsForStatefulSet(set *apps.StatefulSet, s
 		return isMemberOf(set, pod)
 	}
 
-	cm := controller.NewPodControllerRefManager(ssc.podControl, set, selector, controllerKind, ssc.canAdoptFunc(set))
-	return cm.ClaimPods(pods, filter)
+	cm := controller.NewPodControllerRefManager(ssc.podControl, set, selector, controllerKind, ssc.canAdoptFunc(ctx, set))
+	return cm.ClaimPods(ctx, pods, filter)
 }
 
 // If any adoptions are attempted, we should first recheck for deletion with
 // an uncached quorum read sometime after listing Pods/ControllerRevisions (see #42639).
-func (ssc *StatefulSetController) canAdoptFunc(set *apps.StatefulSet) func() error {
-	return controller.RecheckDeletionTimestamp(func() (metav1.Object, error) {
-		fresh, err := ssc.kubeClient.AppsV1().StatefulSets(set.Namespace).Get(context.TODO(), set.Name, metav1.GetOptions{})
+func (ssc *StatefulSetController) canAdoptFunc(ctx context.Context, set *apps.StatefulSet) func(ctx2 context.Context) error {
+	return controller.RecheckDeletionTimestamp(func(ctx context.Context) (metav1.Object, error) {
+		fresh, err := ssc.kubeClient.AppsV1().StatefulSets(set.Namespace).Get(ctx, set.Name, metav1.GetOptions{})
 		if err != nil {
 			return nil, err
 		}
@@ -320,7 +320,7 @@ func (ssc *StatefulSetController) canAdoptFunc(set *apps.StatefulSet) func() err
 }
 
 // adoptOrphanRevisions adopts any orphaned ControllerRevisions matched by set's Selector.
-func (ssc *StatefulSetController) adoptOrphanRevisions(set *apps.StatefulSet) error {
+func (ssc *StatefulSetController) adoptOrphanRevisions(ctx context.Context, set *apps.StatefulSet) error {
 	revisions, err := ssc.control.ListRevisions(set)
 	if err != nil {
 		return err
@@ -332,7 +332,7 @@ func (ssc *StatefulSetController) adoptOrphanRevisions(set *apps.StatefulSet) er
 		}
 	}
 	if len(orphanRevisions) > 0 {
-		canAdoptErr := ssc.canAdoptFunc(set)()
+		canAdoptErr := ssc.canAdoptFunc(ctx, set)(ctx)
 		if canAdoptErr != nil {
 			return fmt.Errorf("can't adopt ControllerRevisions: %v", canAdoptErr)
 		}
@@ -403,13 +403,13 @@ func (ssc *StatefulSetController) enqueueSSAfter(ss *apps.StatefulSet, duration
 
 // processNextWorkItem dequeues items, processes them, and marks them done. It enforces that the syncHandler is never
 // invoked concurrently with the same key.
-func (ssc *StatefulSetController) processNextWorkItem() bool {
+func (ssc *StatefulSetController) processNextWorkItem(ctx context.Context) bool {
 	key, quit := ssc.queue.Get()
 	if quit {
 		return false
 	}
 	defer ssc.queue.Done(key)
-	if err := ssc.sync(key.(string)); err != nil {
+	if err := ssc.sync(ctx, key.(string)); err != nil {
 		utilruntime.HandleError(fmt.Errorf("error syncing StatefulSet %v, requeuing: %v", key.(string), err))
 		ssc.queue.AddRateLimited(key)
 	} else {
@@ -419,13 +419,13 @@ func (ssc *StatefulSetController) processNextWorkItem() bool {
 }
 
 // worker runs a worker goroutine that invokes processNextWorkItem until the controller's queue is closed
-func (ssc *StatefulSetController) worker() {
-	for ssc.processNextWorkItem() {
+func (ssc *StatefulSetController) worker(ctx context.Context) {
+	for ssc.processNextWorkItem(ctx) {
 	}
 }
 
 // sync syncs the given statefulset.
-func (ssc *StatefulSetController) sync(key string) error {
+func (ssc *StatefulSetController) sync(ctx context.Context, key string) error {
 	startTime := time.Now()
 	defer func() {
 		klog.V(4).Infof("Finished syncing statefulset %q (%v)", key, time.Since(startTime))
@@ -452,25 +452,25 @@ func (ssc *StatefulSetController) sync(key string) error {
 		return nil
 	}
 
-	if err := ssc.adoptOrphanRevisions(set); err != nil {
+	if err := ssc.adoptOrphanRevisions(ctx, set); err != nil {
 		return err
 	}
 
-	pods, err := ssc.getPodsForStatefulSet(set, selector)
+	pods, err := ssc.getPodsForStatefulSet(ctx, set, selector)
 	if err != nil {
 		return err
 	}
 
-	return ssc.syncStatefulSet(set, pods)
+	return ssc.syncStatefulSet(ctx, set, pods)
 }
 
 // syncStatefulSet syncs a tuple of (statefulset, []*v1.Pod).
-func (ssc *StatefulSetController) syncStatefulSet(set *apps.StatefulSet, pods []*v1.Pod) error {
+func (ssc *StatefulSetController) syncStatefulSet(ctx context.Context, set *apps.StatefulSet, pods []*v1.Pod) error {
 	klog.V(4).Infof("Syncing StatefulSet %v/%v with %d pods", set.Namespace, set.Name, len(pods))
 	var status *apps.StatefulSetStatus
 	var err error
 	// TODO: investigate where we mutate the set during the update as it is not obvious.
-	status, err = ssc.control.UpdateStatefulSet(set.DeepCopy(), pods)
+	status, err = ssc.control.UpdateStatefulSet(ctx, set.DeepCopy(), pods)
 	if err != nil {
 		return err
 	}

pkg/controller/statefulset/stateful_set_control.go

@@ -17,6 +17,7 @@ limitations under the License.
 package statefulset
 
 import (
+	"context"
 	"sort"
 
 	apps "k8s.io/api/apps/v1"
@@ -38,7 +39,7 @@ type StatefulSetControlInterface interface {
 	// If an implementation returns a non-nil error, the invocation will be retried using a rate-limited strategy.
 	// Implementors should sink any errors that they do not wish to trigger a retry, and they may feel free to
 	// exit exceptionally at any point provided they wish the update to be re-run at a later point in time.
-	UpdateStatefulSet(set *apps.StatefulSet, pods []*v1.Pod) (*apps.StatefulSetStatus, error)
+	UpdateStatefulSet(ctx context.Context, set *apps.StatefulSet, pods []*v1.Pod) (*apps.StatefulSetStatus, error)
 	// ListRevisions returns a array of the ControllerRevisions that represent the revisions of set. If the returned
 	// error is nil, the returns slice of ControllerRevisions is valid.
 	ListRevisions(set *apps.StatefulSet) ([]*apps.ControllerRevision, error)
@@ -73,7 +74,7 @@ type defaultStatefulSetControl struct {
 // strategy allows these constraints to be relaxed - pods will be created and deleted eagerly and
 // in no particular order. Clients using the burst strategy should be careful to ensure they
 // understand the consistency implications of having unpredictable numbers of pods available.
-func (ssc *defaultStatefulSetControl) UpdateStatefulSet(set *apps.StatefulSet, pods []*v1.Pod) (*apps.StatefulSetStatus, error) {
+func (ssc *defaultStatefulSetControl) UpdateStatefulSet(ctx context.Context, set *apps.StatefulSet, pods []*v1.Pod) (*apps.StatefulSetStatus, error) {
 	// list all revisions and sort them
 	revisions, err := ssc.ListRevisions(set)
 	if err != nil {
@@ -81,7 +82,7 @@ func (ssc *defaultStatefulSetControl) UpdateStatefulSet(set *apps.StatefulSet, p
 	}
 	history.SortControllerRevisions(revisions)
 
-	currentRevision, updateRevision, status, err := ssc.performUpdate(set, pods, revisions)
+	currentRevision, updateRevision, status, err := ssc.performUpdate(ctx, set, pods, revisions)
 	if err != nil {
 		return nil, utilerrors.NewAggregate([]error{err, ssc.truncateHistory(set, pods, revisions, currentRevision, updateRevision)})
 	}
@@ -91,7 +92,7 @@ func (ssc *defaultStatefulSetControl) UpdateStatefulSet(set *apps.StatefulSet, p
 }
 
 func (ssc *defaultStatefulSetControl) performUpdate(
-	set *apps.StatefulSet, pods []*v1.Pod, revisions []*apps.ControllerRevision) (*apps.ControllerRevision, *apps.ControllerRevision, *apps.StatefulSetStatus, error) {
+	ctx context.Context, set *apps.StatefulSet, pods []*v1.Pod, revisions []*apps.ControllerRevision) (*apps.ControllerRevision, *apps.ControllerRevision, *apps.StatefulSetStatus, error) {
 	var currentStatus *apps.StatefulSetStatus
 	// get the current, and update revisions
 	currentRevision, updateRevision, collisionCount, err := ssc.getStatefulSetRevisions(set, revisions)
@@ -100,12 +101,12 @@ func (ssc *defaultStatefulSetControl) performUpdate(
 	}
 
 	// perform the main update function and get the status
-	currentStatus, err = ssc.updateStatefulSet(set, currentRevision, updateRevision, collisionCount, pods)
+	currentStatus, err = ssc.updateStatefulSet(ctx, set, currentRevision, updateRevision, collisionCount, pods)
 	if err != nil {
 		return currentRevision, updateRevision, currentStatus, err
 	}
 	// update the set's status
-	err = ssc.updateStatefulSetStatus(set, currentStatus)
+	err = ssc.updateStatefulSetStatus(ctx, set, currentStatus)
 	if err != nil {
 		return currentRevision, updateRevision, currentStatus, err
 	}
@@ -268,6 +269,7 @@ func (ssc *defaultStatefulSetControl) getStatefulSetRevisions(
 // Pods must be at Status.UpdateRevision. If the returned error is nil, the returned StatefulSetStatus is valid and the
 // update must be recorded. If the error is not nil, the method should be retried until successful.
 func (ssc *defaultStatefulSetControl) updateStatefulSet(
+	ctx context.Context,
 	set *apps.StatefulSet,
 	currentRevision *apps.ControllerRevision,
 	updateRevision *apps.ControllerRevision,
@@ -416,7 +418,7 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 		}
 		// If we find a Pod that has not been created we create the Pod
 		if !isCreated(replicas[i]) {
-			if err := ssc.podControl.CreateStatefulPod(set, replicas[i]); err != nil {
+			if err := ssc.podControl.CreateStatefulPod(ctx, set, replicas[i]); err != nil {
 				return &status, err
 			}
 			status.Replicas++
@@ -579,6 +581,7 @@ func (ssc *defaultStatefulSetControl) updateStatefulSet(
 // mutated to indicate completion. If status is semantically equivalent to set's Status no update is performed. If the
 // returned error is nil, the update is successful.
 func (ssc *defaultStatefulSetControl) updateStatefulSetStatus(
+	ctx context.Context,
 	set *apps.StatefulSet,
 	status *apps.StatefulSetStatus) error {
 	// complete any in progress rolling update if necessary
@@ -591,7 +594,7 @@ func (ssc *defaultStatefulSetControl) updateStatefulSetStatus(
 
 	// copy set and update its status
 	set = set.DeepCopy()
-	if err := ssc.statusUpdater.UpdateStatefulSetStatus(set, status); err != nil {
+	if err := ssc.statusUpdater.UpdateStatefulSetStatus(ctx, set, status); err != nil {
 		return err
 	}
 

pkg/controller/statefulset/stateful_set_control_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package statefulset
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -229,7 +230,7 @@ func ReplacesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc)
 		if err != nil {
 			t.Error(err)
 		}
-		if _, err = ssc.UpdateStatefulSet(set, pods); err != nil {
+		if _, err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 			t.Errorf("Failed to update StatefulSet : %s", err)
 		}
 		set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -239,7 +240,7 @@ func ReplacesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc)
 		if pods, err = spc.setPodRunning(set, i); err != nil {
 			t.Error(err)
 		}
-		if _, err = ssc.UpdateStatefulSet(set, pods); err != nil {
+		if _, err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 			t.Errorf("Failed to update StatefulSet : %s", err)
 		}
 		set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -254,7 +255,7 @@ func ReplacesPods(t *testing.T, set *apps.StatefulSet, invariants invariantFunc)
 	if err != nil {
 		t.Error(err)
 	}
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		t.Errorf("Failed to update StatefulSet : %s", err)
 	}
 	set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -278,7 +279,7 @@ func RecreatesFailedPod(t *testing.T, set *apps.StatefulSet, invariants invarian
 	if err != nil {
 		t.Error(err)
 	}
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		t.Errorf("Error updating StatefulSet %s", err)
 	}
 	if err := invariants(set, spc); err != nil {
@@ -290,7 +291,7 @@ func RecreatesFailedPod(t *testing.T, set *apps.StatefulSet, invariants invarian
 	}
 	pods[0].Status.Phase = v1.PodFailed
 	spc.podsIndexer.Update(pods[0])
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		t.Errorf("Error updating StatefulSet %s", err)
 	}
 	if err := invariants(set, spc); err != nil {
@@ -371,7 +372,7 @@ func UpdatePodFailure(t *testing.T, set *apps.StatefulSet, invariants invariantF
 	spc.podsIndexer.Update(pods[0])
 
 	// now it should fail
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil && isOrHasInternalError(err) {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil && isOrHasInternalError(err) {
 		t.Errorf("StatefulSetControl did not return InternalError found %s", err)
 	}
 }
@@ -417,7 +418,7 @@ func PodRecreateDeleteFailure(t *testing.T, set *apps.StatefulSet, invariants in
 	if err != nil {
 		t.Error(err)
 	}
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		t.Errorf("Error updating StatefulSet %s", err)
 	}
 	if err := invariants(set, spc); err != nil {
@@ -430,13 +431,13 @@ func PodRecreateDeleteFailure(t *testing.T, set *apps.StatefulSet, invariants in
 	pods[0].Status.Phase = v1.PodFailed
 	spc.podsIndexer.Update(pods[0])
 	spc.SetDeleteStatefulPodError(apierrors.NewInternalError(errors.New("API server failed")), 0)
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil && isOrHasInternalError(err) {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil && isOrHasInternalError(err) {
 		t.Errorf("StatefulSet failed to %s", err)
 	}
 	if err := invariants(set, spc); err != nil {
 		t.Error(err)
 	}
-	if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		t.Errorf("Error updating StatefulSet %s", err)
 	}
 	if err := invariants(set, spc); err != nil {
@@ -1282,7 +1283,7 @@ func TestStatefulSetControlLimitsHistory(t *testing.T) {
 			if err != nil {
 				t.Fatalf("%s: %s", test.name, err)
 			}
-			_, err = ssc.UpdateStatefulSet(set, pods)
+			_, err = ssc.UpdateStatefulSet(context.TODO(), set, pods)
 			if err != nil {
 				t.Fatalf("%s: %s", test.name, err)
 			}
@@ -1629,7 +1630,7 @@ func TestStatefulSetAvailability(t *testing.T) {
 		if err != nil {
 			t.Fatalf("%s: %s", test.name, err)
 		}
-		status, err := ssc.UpdateStatefulSet(set, pods)
+		status, err := ssc.UpdateStatefulSet(context.TODO(), set, pods)
 		if err != nil {
 			t.Fatalf("%s: %s", test.name, err)
 		}
@@ -1824,7 +1825,7 @@ func (spc *fakeStatefulPodControl) setPodTerminated(set *apps.StatefulSet, ordin
 	return spc.podsLister.Pods(set.Namespace).List(selector)
 }
 
-func (spc *fakeStatefulPodControl) CreateStatefulPod(set *apps.StatefulSet, pod *v1.Pod) error {
+func (spc *fakeStatefulPodControl) CreateStatefulPod(ctx context.Context, set *apps.StatefulSet, pod *v1.Pod) error {
 	defer spc.createPodTracker.inc()
 	if spc.createPodTracker.errorReady() {
 		defer spc.createPodTracker.reset()
@@ -1890,7 +1891,7 @@ func newFakeStatefulSetStatusUpdater(setInformer appsinformers.StatefulSetInform
 	}
 }
 
-func (ssu *fakeStatefulSetStatusUpdater) UpdateStatefulSetStatus(set *apps.StatefulSet, status *apps.StatefulSetStatus) error {
+func (ssu *fakeStatefulSetStatusUpdater) UpdateStatefulSetStatus(ctx context.Context, set *apps.StatefulSet, status *apps.StatefulSetStatus) error {
 	defer ssu.updateStatusTracker.inc()
 	if ssu.updateStatusTracker.errorReady() {
 		defer ssu.updateStatusTracker.reset()
@@ -2089,7 +2090,7 @@ func scaleUpStatefulSetControl(set *apps.StatefulSet,
 			}
 		}
 		// run the controller once and check invariants
-		_, err = ssc.UpdateStatefulSet(set, pods)
+		_, err = ssc.UpdateStatefulSet(context.TODO(), set, pods)
 		if err != nil {
 			return err
 		}
@@ -2117,7 +2118,7 @@ func scaleDownStatefulSetControl(set *apps.StatefulSet, ssc StatefulSetControlIn
 		}
 		sort.Sort(ascendingOrdinal(pods))
 		if ordinal := len(pods) - 1; ordinal >= 0 {
-			if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+			if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 				return err
 			}
 			set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -2127,7 +2128,7 @@ func scaleDownStatefulSetControl(set *apps.StatefulSet, ssc StatefulSetControlIn
 			if pods, err = spc.addTerminatingPod(set, ordinal); err != nil {
 				return err
 			}
-			if _, err = ssc.UpdateStatefulSet(set, pods); err != nil {
+			if _, err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 				return err
 			}
 			set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -2144,7 +2145,7 @@ func scaleDownStatefulSetControl(set *apps.StatefulSet, ssc StatefulSetControlIn
 				spc.podsIndexer.Delete(pods[len(pods)-1])
 			}
 		}
-		if _, err := ssc.UpdateStatefulSet(set, pods); err != nil {
+		if _, err := ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 			return err
 		}
 		set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)
@@ -2207,7 +2208,7 @@ func updateStatefulSetControl(set *apps.StatefulSet,
 	if err != nil {
 		return err
 	}
-	if _, err = ssc.UpdateStatefulSet(set, pods); err != nil {
+	if _, err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 		return err
 	}
 
@@ -2255,7 +2256,7 @@ func updateStatefulSetControl(set *apps.StatefulSet,
 			}
 		}
 
-		if _, err = ssc.UpdateStatefulSet(set, pods); err != nil {
+		if _, err = ssc.UpdateStatefulSet(context.TODO(), set, pods); err != nil {
 			return err
 		}
 		set, err = spc.setsLister.StatefulSets(set.Namespace).Get(set.Name)

pkg/controller/statefulset/stateful_set_status_updater.go

@@ -33,7 +33,7 @@ import (
 type StatefulSetStatusUpdaterInterface interface {
 	// UpdateStatefulSetStatus sets the set's Status to status. Implementations are required to retry on conflicts,
 	// but fail on other errors. If the returned error is nil set's Status has been successfully set to status.
-	UpdateStatefulSetStatus(set *apps.StatefulSet, status *apps.StatefulSetStatus) error
+	UpdateStatefulSetStatus(ctx context.Context, set *apps.StatefulSet, status *apps.StatefulSetStatus) error
 }
 
 // NewRealStatefulSetStatusUpdater returns a StatefulSetStatusUpdaterInterface that updates the Status of a StatefulSet,
@@ -50,11 +50,13 @@ type realStatefulSetStatusUpdater struct {
 }
 
 func (ssu *realStatefulSetStatusUpdater) UpdateStatefulSetStatus(
+	ctx context.Context,
 	set *apps.StatefulSet,
 	status *apps.StatefulSetStatus) error {
 	// don't wait due to limited number of clients, but backoff after the default number of steps
 	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
 		set.Status = *status
+		// TODO: This context.TODO should use a real context once we have RetryOnConflictWithContext
 		_, updateErr := ssu.client.AppsV1().StatefulSets(set.Namespace).UpdateStatus(context.TODO(), set, metav1.UpdateOptions{})
 		if updateErr == nil {
 			return nil

pkg/controller/statefulset/stateful_set_status_updater_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package statefulset
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -41,7 +42,7 @@ func TestStatefulSetUpdaterUpdatesSetStatus(t *testing.T) {
 		update := action.(core.UpdateAction)
 		return true, update.GetObject(), nil
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err != nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err != nil {
 		t.Errorf("Error returned on successful status update: %s", err)
 	}
 	if set.Status.Replicas != 2 {
@@ -62,7 +63,7 @@ func TestStatefulSetStatusUpdaterUpdatesObservedGeneration(t *testing.T) {
 		}
 		return true, sts, nil
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err != nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err != nil {
 		t.Errorf("Error returned on successful status update: %s", err)
 	}
 }
@@ -78,7 +79,7 @@ func TestStatefulSetStatusUpdaterUpdateReplicasFailure(t *testing.T) {
 	fakeClient.AddReactor("update", "statefulsets", func(action core.Action) (bool, runtime.Object, error) {
 		return true, nil, apierrors.NewInternalError(errors.New("API server down"))
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err == nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err == nil {
 		t.Error("Failed update did not return error")
 	}
 }
@@ -101,7 +102,7 @@ func TestStatefulSetStatusUpdaterUpdateReplicasConflict(t *testing.T) {
 		return true, update.GetObject(), nil
 
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err != nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err != nil {
 		t.Errorf("UpdateStatefulSetStatus returned an error: %s", err)
 	}
 	if set.Status.Replicas != 2 {
@@ -121,7 +122,7 @@ func TestStatefulSetStatusUpdaterUpdateReplicasConflictFailure(t *testing.T) {
 		update := action.(core.UpdateAction)
 		return true, update.GetObject(), apierrors.NewConflict(action.GetResource().GroupResource(), set.Name, errors.New("object already exists"))
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err == nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err == nil {
 		t.Error("UpdateStatefulSetStatus failed to return an error on get failure")
 	}
 }
@@ -136,7 +137,7 @@ func TestStatefulSetStatusUpdaterGetAvailableReplicas(t *testing.T) {
 		update := action.(core.UpdateAction)
 		return true, update.GetObject(), nil
 	})
-	if err := updater.UpdateStatefulSetStatus(set, &status); err != nil {
+	if err := updater.UpdateStatefulSetStatus(context.TODO(), set, &status); err != nil {
 		t.Errorf("Error returned on successful status update: %s", err)
 	}
 	if set.Status.AvailableReplicas != 3 {

pkg/controller/statefulset/stateful_set_test.go

@@ -18,6 +18,7 @@ package statefulset
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"sort"
 	"testing"
@@ -106,7 +107,7 @@ func TestStatefulSetControllerRespectsTermination(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	ssc.syncStatefulSet(set, pods)
+	ssc.syncStatefulSet(context.TODO(), set, pods)
 	selector, err := metav1.LabelSelectorAsSelector(set.Spec.Selector)
 	if err != nil {
 		t.Error(err)
@@ -555,7 +556,7 @@ func TestGetPodsForStatefulSetAdopt(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	pods, err := ssc.getPodsForStatefulSet(set, selector)
+	pods, err := ssc.getPodsForStatefulSet(context.TODO(), set, selector)
 	if err != nil {
 		t.Fatalf("getPodsForStatefulSet() error: %v", err)
 	}
@@ -592,7 +593,7 @@ func TestAdoptOrphanRevisions(t *testing.T) {
 	spc.revisionsIndexer.Add(ss1Rev1)
 	spc.revisionsIndexer.Add(ss1Rev2)
 
-	err = ssc.adoptOrphanRevisions(ss1)
+	err = ssc.adoptOrphanRevisions(context.TODO(), ss1)
 	if err != nil {
 		t.Errorf("adoptOrphanRevisions() error: %v", err)
 	}
@@ -634,7 +635,7 @@ func TestGetPodsForStatefulSetRelease(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	pods, err := ssc.getPodsForStatefulSet(set, selector)
+	pods, err := ssc.getPodsForStatefulSet(context.TODO(), set, selector)
 	if err != nil {
 		t.Fatalf("getPodsForStatefulSet() error: %v", err)
 	}
@@ -673,7 +674,7 @@ func newFakeStatefulSetController(initialObjects ...runtime.Object) (*StatefulSe
 
 func fakeWorker(ssc *StatefulSetController) {
 	if obj, done := ssc.queue.Get(); !done {
-		ssc.sync(obj.(string))
+		ssc.sync(context.TODO(), obj.(string))
 		ssc.queue.Done(obj)
 	}
 }