Add await attempt timeout

Author: stephanos

.github/workflows/run-tests.yml

@@ -23,6 +23,7 @@ env:
   PR_BASE_COMMIT: ${{ github.event.pull_request.base.sha }}
   DOCKER_COMPOSE_FILE: ./develop/github/docker-compose.yml
   TEMPORAL_VERSION_CHECK_DISABLED: 1
+  TEMPORAL_AWAIT_ATTEMPT_TIMEOUT: 15s
   MAX_TEST_ATTEMPTS: 3
   SHARD_COUNT: 3 # NOTE: must match shard count in optimize-test-sharding.yml
 

common/testing/await/config.go

@@ -0,0 +1,40 @@
+package await
+
+import (
+	"os"
+	"time"
+
+	"go.temporal.io/server/common/debug"
+)
+
+const attemptTimeoutEnvVar = "TEMPORAL_AWAIT_ATTEMPT_TIMEOUT"
+
+type config struct {
+	totalTimeout   time.Duration
+	pollInterval   time.Duration
+	attemptTimeout time.Duration
+	timeoutMsg     string
+}
+
+func newConfig() config {
+	return config{
+		attemptTimeout: envDuration(attemptTimeoutEnvVar, 10*time.Second) * debug.TimeoutMultiplier,
+	}
+}
+
+func legacyConfig(timeout, pollInterval time.Duration, timeoutMsg string) config {
+	cfg := newConfig()
+	cfg.totalTimeout = timeout
+	cfg.pollInterval = pollInterval
+	cfg.timeoutMsg = timeoutMsg
+	return cfg
+}
+
+func envDuration(name string, fallback time.Duration) time.Duration {
+	if s := os.Getenv(name); s != "" {
+		if d, err := time.ParseDuration(s); err == nil && d > 0 {
+			return d
+		}
+	}
+	return fallback
+}

common/testing/await/config_test.go

@@ -0,0 +1,16 @@
+package await
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.temporal.io/server/common/debug"
+)
+
+func TestConfig_OverrideAttemptTimeout(t *testing.T) {
+	t.Setenv(attemptTimeoutEnvVar, "250ms")
+
+	cfg := newConfig()
+	require.Equal(t, 250*time.Millisecond*debug.TimeoutMultiplier, cfg.attemptTimeout)
+}

common/testing/await/report.go

@@ -20,13 +20,39 @@ type attemptFailure struct {
 	errors  []string
 }
 
-// reportTimeout reports the timeout failure plus collected attempt errors.
-func reportTimeout(tb testing.TB, failures []attemptFailure, funcName, timeoutMsg string, effectiveTimeout time.Duration, polls int) {
-	reportAttemptErrors(tb, failures)
+type timeoutReport struct {
+	effectiveTimeout time.Duration
+	attempts         int
+	attemptTimeouts  int
+	failures         []attemptFailure
+}
+
+func (r *timeoutReport) nextPoll() {
+	r.attempts++
+}
+
+func (r *timeoutReport) recordErrors(errors []string) {
+	if len(errors) > 0 {
+		r.failures = append(r.failures, attemptFailure{attempt: r.attempts, errors: errors})
+	}
+}
+
+func (r *timeoutReport) recordAttemptTimeout() {
+	r.attemptTimeouts++
+}
+
+func (r timeoutReport) reportAttemptErrors(tb testing.TB) {
+	reportAttemptErrors(tb, r.failures)
+}
+
+func (r timeoutReport) reportTimeout(tb testing.TB, funcName, timeoutMsg string) {
+	r.reportAttemptErrors(tb)
 	if timeoutMsg != "" {
-		tb.Fatalf("%s: %s (not satisfied after %v, %d polls)", funcName, timeoutMsg, effectiveTimeout, polls)
+		tb.Fatalf("%s: %s (not satisfied after %v)\ndetails:\n  attempts         = %d\n  attempt timeouts = %d",
+			funcName, timeoutMsg, r.effectiveTimeout, r.attempts, r.attemptTimeouts)
 	} else {
-		tb.Fatalf("%s: condition not satisfied after %v (%d polls)", funcName, effectiveTimeout, polls)
+		tb.Fatalf("%s: condition not satisfied after %v\ndetails:\n  attempts         = %d\n  attempt timeouts = %d",
+			funcName, r.effectiveTimeout, r.attempts, r.attemptTimeouts)
 	}
 }
 

common/testing/await/report_test.go

@@ -0,0 +1,52 @@
+package await
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestReportTimeout(t *testing.T) {
+	tb := newReportRecordingTB()
+
+	timeoutReport{
+		effectiveTimeout: time.Second,
+		attempts:         3,
+		attemptTimeouts:  2,
+	}.reportTimeout(tb, "Require", "")
+
+	require.Equal(t, strings.Join([]string{
+		"Require: condition not satisfied after 1s",
+		"details:",
+		"  attempts         = 3",
+		"  attempt timeouts = 2",
+	}, "\n"), tb.fatals())
+}
+
+type reportRecordingTB struct {
+	testing.TB
+	mu            sync.Mutex
+	fatalMessages []string
+}
+
+func newReportRecordingTB() *reportRecordingTB {
+	return &reportRecordingTB{}
+}
+
+func (r *reportRecordingTB) Helper() {}
+
+func (r *reportRecordingTB) Fatalf(format string, args ...any) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.fatalMessages = append(r.fatalMessages, fmt.Sprintf(format, args...))
+}
+
+func (r *reportRecordingTB) fatals() string {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return strings.Join(r.fatalMessages, "\n")
+}

common/testing/await/require_ctx.go

@@ -55,23 +55,21 @@ func hardDeadlockTimeout() time.Duration {
 // test failure. Use t.Context() inside the callback to honor the timeout.
 func Require(ctx context.Context, tb testing.TB, condition func(*T), timeout, pollInterval time.Duration) {
 	tb.Helper()
-	run(ctx, tb, condition, timeout, pollInterval, "", "Require", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(timeout, pollInterval, ""), "Require", requireMisuseHint, true)
 }
 
 // Requiref is like [Require] but adds a formatted message to the timeout
 // failure.
 func Requiref(ctx context.Context, tb testing.TB, condition func(*T), timeout, pollInterval time.Duration, msg string, args ...any) {
 	tb.Helper()
-	run(ctx, tb, condition, timeout, pollInterval, fmt.Sprintf(msg, args...), "Requiref", requireMisuseHint, true)
+	run(ctx, tb, condition, legacyConfig(timeout, pollInterval, fmt.Sprintf(msg, args...)), "Requiref", requireMisuseHint, true)
 }
 
 func run(
 	parentCtx context.Context,
 	tb testing.TB,
 	condition func(*T),
-	timeout,
-	pollInterval time.Duration,
-	timeoutMsg string,
+	cfg config,
 	funcName string,
 	misuseHint string,
 	cancellable bool,
@@ -89,7 +87,7 @@ func run(
 		return
 	}
 
-	deadline := time.Now().Add(timeout)
+	deadline := time.Now().Add(cfg.totalTimeout)
 
 	// Cap at the parent context's deadline if it's earlier than our timeout.
 	if parentDeadline, hasDeadline := parentCtx.Deadline(); hasDeadline && parentDeadline.Before(deadline) {
@@ -108,22 +106,21 @@ func run(
 	awaitCtx, awaitCancel := context.WithDeadline(parentCtx, deadline)
 	defer awaitCancel()
 
-	var failures []attemptFailure
-	polls := 0
+	report := timeoutReport{effectiveTimeout: effectiveTimeout}
 
 	for {
 		// Parent context was canceled while we were sleeping (not our deadline).
 		if err := awaitCtx.Err(); err != nil && !deadlineReached(deadline) {
-			reportAttemptErrors(tb, failures)
+			report.reportAttemptErrors(tb)
 			tb.Fatalf("%s: context canceled before condition was satisfied: %v", funcName, err)
 			return
 		}
 
-		polls++
+		report.nextPoll()
 
-		// Fresh context per attempt, scoped to the run-level ctx. runAttempt
-		// owns the soft timeout and the corresponding cancel.
-		attemptCtx, attemptCancel := context.WithCancel(awaitCtx)
+		// Per-attempt context: bounded by the configured attempt timeout and
+		// further capped by the overall awaitCtx.
+		attemptCtx, attemptCancel := context.WithTimeout(awaitCtx, cfg.attemptTimeout)
 		t := &T{tb: tb, ctx: attemptCtx}
 
 		// Run attempt.
@@ -133,18 +130,24 @@ func run(
 			panic(res.panicVal) // propagate to caller
 		}
 		if res.deadlocked {
-			reportAttemptErrors(tb, failures)
+			report.reportAttemptErrors(tb)
 			if cancellable {
-				tb.Fatalf("%s: condition still running %v past context cancellation — does it honor t.Context()? (%d polls)",
-					funcName, hardDeadlockTimeout(), polls)
+				tb.Fatalf("%s: condition still running %v past context cancellation — does it honor t.Context()? (%d attempts)",
+					funcName, hardDeadlockTimeout(), report.attempts)
 			} else {
-				tb.Fatalf("%s: condition still running %v past deadline (%d polls)",
-					funcName, hardDeadlockTimeout(), polls)
+				tb.Fatalf("%s: condition still running %v past deadline (%d attempts)",
+					funcName, hardDeadlockTimeout(), report.attempts)
 			}
 			return
 		}
-		if len(t.errors) > 0 {
-			failures = append(failures, attemptFailure{attempt: polls, errors: t.errors})
+		report.recordErrors(t.errors)
+
+		// Attempt-timeout expiry: attemptCtx is done but awaitCtx is not.
+		// Record nothing special - the attempt's recorded errors (if any)
+		// already describe what went wrong; otherwise we just retry.
+		attemptHitOwnTimeout := attemptCtx.Err() == context.DeadlineExceeded && awaitCtx.Err() == nil
+		if attemptHitOwnTimeout {
+			report.recordAttemptTimeout()
 		}
 
 		// Check misuse where the real test failed instead of just the attempt.
@@ -155,24 +158,24 @@ func run(
 
 		// Parent context was canceled during the attempt (not our deadline).
 		if err := awaitCtx.Err(); err != nil && !deadlineReached(deadline) {
-			reportAttemptErrors(tb, failures)
+			report.reportAttemptErrors(tb)
 			tb.Fatalf("%s: context canceled before condition was satisfied: %v", funcName, err)
 			return
 		}
 
 		// Our deadline expired.
 		if deadlineReached(deadline) {
-			reportTimeout(tb, failures, funcName, timeoutMsg, effectiveTimeout, polls)
+			report.reportTimeout(tb, funcName, cfg.timeoutMsg)
 			return
 		}
 
 		// Success: attempt completed without failures.
-		if !res.stopped && !t.Failed() {
+		if !res.stopped && !t.Failed() && !attemptHitOwnTimeout {
 			return
 		}
 
 		// Wait for pollInterval, or context is canceled or deadline is reached.
-		sleep(awaitCtx, deadline, pollInterval)
+		sleep(awaitCtx, deadline, cfg.pollInterval)
 	}
 }
 

common/testing/await/require_ctx_test.go

@@ -367,6 +367,37 @@ func TestRequire_FailureScenarios(t *testing.T) {
 	})
 }
 
+func TestRequire_AttemptTimeoutRetriesUntilAwaitTimeout(t *testing.T) {
+	const (
+		attemptTimeout = 50 * time.Millisecond
+		awaitTimeout   = time.Second
+		pollInterval   = 100 * time.Millisecond
+	)
+	t.Setenv("TEMPORAL_AWAIT_ATTEMPT_TIMEOUT", attemptTimeout.String())
+
+	ctx := testcontext.New(t)
+	var attempts atomic.Int32
+	var firstAttemptRemaining time.Duration
+
+	tb := newRecordingTB()
+	tb.run(func() {
+		await.Require(ctx, tb, func(t *await.T) {
+			if attempts.Add(1) == 1 {
+				deadline, ok := t.Context().Deadline()
+				require.True(t, ok)
+				firstAttemptRemaining = time.Until(deadline)
+			}
+			<-t.Context().Done()
+		}, awaitTimeout, pollInterval)
+	})
+
+	require.True(t, tb.Failed())
+	require.Contains(t, tb.fatals(), "not satisfied after")
+	require.Positive(t, firstAttemptRemaining)
+	require.LessOrEqual(t, firstAttemptRemaining, attemptTimeout)
+	require.Greater(t, attempts.Load(), int32(1))
+}
+
 func TestRequire_SoftDeadlockLogsAndCancels(t *testing.T) {
 	// not using T.Parallel() so it can use t.Setenv to override the deadlock timeouts
 	t.Setenv("TEMPORAL_AWAIT_SOFT_DEADLOCK_TIMEOUT", "50ms")

common/testing/await/require_true.go

@@ -21,7 +21,7 @@ func RequireTrue(tb testing.TB, condition func() bool, timeout, pollInterval tim
 		if !condition() {
 			t.Fail()
 		}
-	}, timeout, pollInterval, "", "RequireTrue", requireTrueMisuseHint, false)
+	}, legacyConfig(timeout, pollInterval, ""), "RequireTrue", requireTrueMisuseHint, false)
 }
 
 // RequireTruef is like [RequireTrue] but accepts a format string that is included
@@ -32,5 +32,5 @@ func RequireTruef(tb testing.TB, condition func() bool, timeout, pollInterval ti
 		if !condition() {
 			t.Fail()
 		}
-	}, timeout, pollInterval, fmt.Sprintf(msg, args...), "RequireTruef", requireTrueMisuseHint, false)
+	}, legacyConfig(timeout, pollInterval, fmt.Sprintf(msg, args...)), "RequireTruef", requireTrueMisuseHint, false)
 }