From 5e3923ed1be2bbc8fcea8786e8b9c3669627d962 Mon Sep 17 00:00:00 2001
From: Antonio Ojea <aojea@google.com>
Date: Thu, 15 Jun 2023 10:24:11 +0000
Subject: [PATCH] always execute condition for wait.PollUntilContextTimeout
 with immediate=true

The API guarantees that the condition will be executed at least once
independently of the context.

Change-Id: I9ac8ed681d03f38e270c7726b482c05ca6a568ce

Kubernetes-commit: 59cd1d0b3bb378f40a639e21b615f4df1d4a5a14
---
 pkg/util/wait/loop.go      | 19 ++++++++++++----
 pkg/util/wait/loop_test.go | 44 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/pkg/util/wait/loop.go b/pkg/util/wait/loop.go
index 51864d70f..0dd13c626 100644
--- a/pkg/util/wait/loop.go
+++ b/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:
diff --git a/pkg/util/wait/loop_test.go b/pkg/util/wait/loop_test.go
index c5849250a..992d3d04d 100644
--- a/pkg/util/wait/loop_test.go
+++ b/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,