Harden adaptive Await polling

common/testing/await/config.go

@@ -7,11 +7,7 @@ import (
 	"go.temporal.io/server/common/debug"
 )
 
-const (
-	attemptTimeoutEnvVar = "TEMPORAL_AWAIT_ATTEMPT_TIMEOUT"
-	minPollInterval      = 500 * time.Millisecond
-	maxPollInterval      = 2 * time.Second
-)
+const attemptTimeoutEnvVar = "TEMPORAL_AWAIT_ATTEMPT_TIMEOUT"
 
 type config struct {
 	totalTimeout   time.Duration
@@ -25,7 +21,7 @@ func newConfig() config {
 	}
 }
 
-func legacyConfig(timeout, _ time.Duration, timeoutMsg string) config {
+func legacyConfig(timeout time.Duration, timeoutMsg string) config {
 	cfg := newConfig()
 	cfg.totalTimeout = timeout
 	cfg.timeoutMsg = timeoutMsg
@@ -35,11 +31,11 @@ func legacyConfig(timeout, _ time.Duration, timeoutMsg string) config {
 func nextPollInterval(attempt int) time.Duration {
 	switch attempt {
 	case 1:
-		return minPollInterval
+		return 500 * time.Millisecond
 	case 2:
 		return time.Second
 	default:
-		return maxPollInterval
+		return 2 * time.Second
 	}
 }
 

common/testing/await/require_ctx.go

@@ -48,16 +48,16 @@ const postAwaitTimeoutReserve = 10 * time.Second
 // Pass the *await.T to require.*/assert.* — failures cause a retry, not a
 // test failure. Use t.Context() inside the callback to honor the timeout.
 // The poll interval argument is retained for source compatibility and ignored.
-func Require(ctx context.Context, tb testing.TB, condition func(*T), timeout, pollInterval time.Duration) {
+func Require(ctx context.Context, tb testing.TB, condition func(*T), timeout, _ time.Duration) {
 	tb.Helper()
-	run(ctx, tb, condition, legacyConfig(timeout, pollInterval, ""), "Require", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(timeout, ""), "Require", requireMisuseHint, true)
 }
 
 // Requiref is like [Require] but adds a formatted message to the timeout
 // failure. Its poll interval argument is also ignored.
-func Requiref(ctx context.Context, tb testing.TB, condition func(*T), timeout, pollInterval time.Duration, msg string, args ...any) {
+func Requiref(ctx context.Context, tb testing.TB, condition func(*T), timeout, _ time.Duration, msg string, args ...any) {
 	tb.Helper()
-	run(ctx, tb, condition, legacyConfig(timeout, pollInterval, fmt.Sprintf(msg, args...)), "Requiref", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(timeout, fmt.Sprintf(msg, args...)), "Requiref", requireMisuseHint, true)
 }
 
 func run(
@@ -186,7 +186,11 @@ func run(
 		}
 
 		// Wait for the next poll interval, or context is canceled or deadline is reached.
-		sleep(awaitCtx, deadline, nextPollInterval(report.attempts))
+		pollInterval := min(
+			nextPollInterval(report.attempts),
+			max(time.Nanosecond, time.Until(deadline)/2),
+		)
+		sleep(awaitCtx, deadline, pollInterval)
 	}
 }
 

common/testing/await/require_ctx_test.go

@@ -219,7 +219,7 @@ func TestRequire_FailureScenarios(t *testing.T) {
 					firstAttemptRemaining = time.Until(deadline)
 				}
 				<-t.Context().Done()
-			}, 2*attemptTimeout+500*time.Millisecond, pollInterval)
+			}, 2*attemptTimeout+time.Second, pollInterval)
 		})
 
 		require.True(t, tb.Failed())

common/testing/await/require_true.go

@@ -16,23 +16,23 @@ const requireTrueMisuseHint = "do not use test assertions inside the predicate -
 //
 // Use [RequireTrue] for simple local predicates only. Do not use assertions or
 // side effects in the predicate - use [Require] for these.
-func RequireTrue(tb testing.TB, condition func() bool, timeout, pollInterval time.Duration) {
+func RequireTrue(tb testing.TB, condition func() bool, timeout, _ time.Duration) {
 	tb.Helper()
 	run(testcontext.For(tb), tb, func(t *T) {
 		if !condition() {
 			t.Fail()
 		}
-	}, legacyConfig(timeout, pollInterval, ""), "RequireTrue", requireTrueMisuseHint, false)
+	}, legacyConfig(timeout, ""), "RequireTrue", requireTrueMisuseHint, false)
 }
 
 // RequireTruef is like [RequireTrue] but accepts a format string that is included
 // in the failure message when the condition is not satisfied before the timeout.
 // Its poll interval argument is also ignored.
-func RequireTruef(tb testing.TB, condition func() bool, timeout, pollInterval time.Duration, msg string, args ...any) {
+func RequireTruef(tb testing.TB, condition func() bool, timeout, _ time.Duration, msg string, args ...any) {
 	tb.Helper()
 	run(testcontext.For(tb), tb, func(t *T) {
 		if !condition() {
 			t.Fail()
 		}
-	}, legacyConfig(timeout, pollInterval, fmt.Sprintf(msg, args...)), "RequireTruef", requireTrueMisuseHint, false)
+	}, legacyConfig(timeout, fmt.Sprintf(msg, args...)), "RequireTruef", requireTrueMisuseHint, false)
 }

common/testing/await/require_true_test.go

@@ -3,6 +3,7 @@ package await_test
 import (
 	"sync/atomic"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/stretchr/testify/require"
@@ -37,6 +38,20 @@ func TestRequireTrue_RetriesFalseUntilTrue(t *testing.T) {
 	require.Equal(t, int32(3), attempts.Load())
 }
 
+func TestRequireTrue_RetriesWithinShortTimeout(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		readyAt := time.Now().Add(10 * time.Millisecond)
+		var attempts int
+
+		await.RequireTrue(t, func() bool {
+			attempts++
+			return !time.Now().Before(readyAt)
+		}, 20*time.Millisecond, time.Nanosecond)
+
+		require.Greater(t, attempts, 1)
+	})
+}
+
 func TestRequireTrue_FailureScenarios(t *testing.T) {
 	t.Parallel()