fix(shared): make SetOnce.WaitWithContext deterministic once value is set (#3241)

When both the Done channel and ctx.Done() are ready, select picks a case
at random, so a wait entered with an already cancelled context could
report ctx.Err() for work that had already completed, randomly turning
finished work into a spurious failure under deadline pressure.

Check the Done channel first and document the contract: a set result
always wins; ctx.Err() is returned only when the context becomes done
before the result is set. This is deliberate -- the context bounds the
waiting, not the delivery of an already-existing result, and fail-fast
callers can still check ctx.Err() themselves.

Drop the stale caller-migration TODO (from #257): the choice between
Wait and WaitWithContext is purely about cancellation, and blocking Wait
remains correct on teardown paths where cancellation must not skip work.

Author: arkamar

packages/shared/pkg/utils/set_once.go

@@ -101,8 +101,21 @@ func (s *SetOnce[T]) Result() (T, error) {
 	return s.res.value, s.res.err
 }
 
-// WaitWithContext TODO: Use waitWithContext in all places instead of Wait.
+// WaitWithContext returns the value or error set by SetValue or SetError,
+// waiting until the result is set or ctx is done, whichever comes first.
+//
+// An already set result always wins: once SetValue or SetError has been
+// called, WaitWithContext returns that result even if ctx is already
+// cancelled. ctx.Err() is returned only when ctx becomes done before the
+// result is set.
 func (s *SetOnce[T]) WaitWithContext(ctx context.Context) (T, error) {
+	// An already set result always wins over an already cancelled context.
+	select {
+	case <-s.Done:
+		return s.Result()
+	default:
+	}
+
 	select {
 	case <-s.Done:
 		return s.Result()

packages/shared/pkg/utils/set_once_test.go

@@ -111,6 +111,24 @@ func TestSetOnceWaitWithContextCanceled(t *testing.T) {
 	wg.Wait()
 }
 
+func TestSetOnceWaitWithContextCanceledValueSet(t *testing.T) {
+	t.Parallel()
+	setOnce := NewSetOnce[int]()
+
+	require.NoError(t, setOnce.SetValue(1))
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	// A set value must always win over a cancelled context. Repeat to catch
+	// the random select choice when both channels are ready.
+	for range 200 {
+		value, err := setOnce.WaitWithContext(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, 1, value)
+	}
+}
+
 func TestSetOnceSetResultConcurrent(t *testing.T) {
 	t.Parallel()
 	setOnce := NewSetOnce[int]()