Automated cherry pick of #118686: always execute condition for wait.PollUntilContextTimeout

staging/src/k8s.io/apimachinery/pkg/util/wait/loop.go

@@ -27,9 +27,11 @@ import (
 // the provided timer until the provided context is cancelled, the condition returns
 // true, or the condition returns an error. If sliding is true, the period is computed
 // after condition runs. If it is false then period includes the runtime for condition.
-// If immediate is false the first delay happens before any call to condition. The
-// returned error is the error returned by the last condition or the context error if
-// the context was terminated.
+// If immediate is false the first delay happens before any call to condition, if
+// immediate is true the condition will be invoked before waiting and guarantees that
+// the condition is invoked at least once, regardless of whether the context has been
+// cancelled. The returned error is the error returned by the last condition or the
+// context error if the context was terminated.
 //
 // This is the common loop construct for all polling in the wait package.
 func loopConditionUntilContext(ctx context.Context, t Timer, immediate, sliding bool, condition ConditionWithContextFunc) error {
@@ -38,8 +40,17 @@ func loopConditionUntilContext(ctx context.Context, t Timer, immediate, sliding
 	var timeCh <-chan time.Time
 	doneCh := ctx.Done()
 
+	// if immediate is true the condition is
+	// guaranteed to be executed at least once,
 	// if we haven't requested immediate execution, delay once
-	if !immediate {
+	if immediate {
+		if ok, err := func() (bool, error) {
+			defer runtime.HandleCrash()
+			return condition(ctx)
+		}(); err != nil || ok {
+			return err
+		}
+	} else {
 		timeCh = t.C()
 		select {
 		case <-doneCh:

staging/src/k8s.io/apimachinery/pkg/util/wait/loop_test.go

@@ -144,6 +144,33 @@ func Test_loopConditionUntilContext_semantic(t *testing.T) {
 			attemptsExpected: 0,
 			errExpected:      context.Canceled,
 		},
+		{
+			name:    "context already canceled condition success and immediate 1 attempt expected",
+			context: cancelledContext,
+			callback: func(_ int) (bool, error) {
+				return true, nil
+			},
+			immediate:        true,
+			attemptsExpected: 1,
+		},
+		{
+			name:    "context already canceled condition fail and immediate 1 attempt expected",
+			context: cancelledContext,
+			callback: func(_ int) (bool, error) {
+				return false, conditionErr
+			},
+			immediate:        true,
+			attemptsExpected: 1,
+			errExpected:      conditionErr,
+		},
+		{
+			name:             "context already canceled and immediate 1 attempt expected",
+			context:          cancelledContext,
+			callback:         defaultCallback,
+			immediate:        true,
+			attemptsExpected: 1,
+			errExpected:      context.Canceled,
+		},
 		{
 			name:               "context cancelled after 5 attempts",
 			context:            defaultContext,
@@ -152,6 +179,23 @@ func Test_loopConditionUntilContext_semantic(t *testing.T) {
 			attemptsExpected:   5,
 			errExpected:        context.Canceled,
 		},
+		{
+			name:               "context cancelled and immediate after 5 attempts",
+			context:            defaultContext,
+			callback:           defaultCallback,
+			immediate:          true,
+			cancelContextAfter: 5,
+			attemptsExpected:   5,
+			errExpected:        context.Canceled,
+		},
+		{
+			name:             "context at deadline and immediate 1 attempt expected",
+			context:          deadlinedContext,
+			callback:         defaultCallback,
+			immediate:        true,
+			attemptsExpected: 1,
+			errExpected:      context.DeadlineExceeded,
+		},
 		{
 			name:             "context at deadline no attempts expected",
 			context:          deadlinedContext,