From 4fc0093388e58edb3a29c00ebbf850e144670e3a Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Tue, 16 Sep 2025 21:53:50 +0400
Subject: [PATCH] fix: fix TestCloserStack_Timeout to wait for all asyncClosers
 (#19837)

fixes https://github.com/coder/internal/issues/966


TestCloserStack_Timeout creates `asyncCloser`s which allow control over the exact timing and order of their close method returning. They also, as a final backstop will throw an error if the test context ends before they are unblocked.

TestCloserStack_Timeout unblocks all `asyncCloser`s in a defer and then ends the test. This defer _unblocks_ the running close goroutines, but does not wait for them to finish. Since the test context is canceled as soon as the test completes, this creates a race condition where the close goroutines can trigger the context cancelled  arm of the `select` statement.

The fix is to both unblock and wait for all close goroutines to complete before ending the test and cancelling the context.
---
 cli/ssh_internal_test.go | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/cli/ssh_internal_test.go b/cli/ssh_internal_test.go
index a7fac11c72..3cf562ce82 100644
--- a/cli/ssh_internal_test.go
+++ b/cli/ssh_internal_test.go
@@ -158,7 +158,7 @@ func TestCloserStack_CloseAfterContext(t *testing.T) {
 	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
 	uut := newCloserStack(ctx, logger, quartz.NewMock(t))
 	ac := newAsyncCloser(testCtx, t)
-	defer ac.complete()
+	defer ac.unblock()
 	err := uut.push("async", ac)
 	require.NoError(t, err)
 	cancel()
@@ -178,8 +178,9 @@ func TestCloserStack_CloseAfterContext(t *testing.T) {
 		t.Fatal("closed before stack was finished")
 	}
 
-	ac.complete()
+	ac.unblock()
 	testutil.TryReceive(testCtx, t, closed)
+	testutil.TryReceive(testCtx, t, ac.done)
 }
 
 func TestCloserStack_Timeout(t *testing.T) {
@@ -198,7 +199,8 @@ func TestCloserStack_Timeout(t *testing.T) {
 	}
 	defer func() {
 		for _, a := range ac {
-			a.complete()
+			a.unblock()
+			testutil.TryReceive(ctx, t, a.done) // ensure we don't race with context cancellation
 		}
 	}()
 
@@ -215,7 +217,7 @@ func TestCloserStack_Timeout(t *testing.T) {
 	testutil.TryReceive(ctx, t, ac[1].started)
 
 	// middle one finishes
-	ac[1].complete()
+	ac[1].unblock()
 	// bottom starts, but also hangs
 	testutil.TryReceive(ctx, t, ac[0].started)
 
@@ -317,34 +319,37 @@ func (c *fakeCloser) Close() error {
 }
 
 type asyncCloser struct {
-	t             *testing.T
-	ctx           context.Context
-	started       chan struct{}
-	isComplete    chan struct{}
-	comepleteOnce sync.Once
+	t           *testing.T
+	ctx         context.Context
+	started     chan struct{}
+	done        chan struct{}
+	isUnblocked chan struct{}
+	unblockOnce sync.Once
 }
 
 func (c *asyncCloser) Close() error {
 	close(c.started)
+	defer close(c.done)
 	select {
 	case <-c.ctx.Done():
 		c.t.Error("timed out")
 		return c.ctx.Err()
-	case <-c.isComplete:
+	case <-c.isUnblocked:
 		return nil
 	}
 }
 
-func (c *asyncCloser) complete() {
-	c.comepleteOnce.Do(func() { close(c.isComplete) })
+func (c *asyncCloser) unblock() {
+	c.unblockOnce.Do(func() { close(c.isUnblocked) })
 }
 
 func newAsyncCloser(ctx context.Context, t *testing.T) *asyncCloser {
 	return &asyncCloser{
-		t:          t,
-		ctx:        ctx,
-		isComplete: make(chan struct{}),
-		started:    make(chan struct{}),
+		t:           t,
+		ctx:         ctx,
+		isUnblocked: make(chan struct{}),
+		started:     make(chan struct{}),
+		done:        make(chan struct{}),
 	}
 }