fixed flaky tests

ibetitsmike · ibetitsmike · commit 85c505d3d495 · 2025-08-20T12:38:35.000Z
diff --git a/agent/immortalstreams/backedpipe/backed_pipe.go b/agent/immortalstreams/backedpipe/backed_pipe.go
@@ -327,7 +327,7 @@ func (bp *BackedPipe) ForceReconnect() error {
 	// Deduplicate concurrent ForceReconnect calls so only one reconnection
 	// attempt runs at a time from this API. Use the pipe's internal context
 	// to ensure Close() cancels any in-flight attempt.
-	_, err, _ := bp.sf.Do("backedpipe-reconnect", func() (interface{}, error) {
+	_, err, _ := bp.sf.Do("force-reconnect", func() (interface{}, error) {
 		bp.mu.Lock()
 		defer bp.mu.Unlock()
 
diff --git a/agent/immortalstreams/backedpipe/backed_pipe_test.go b/agent/immortalstreams/backedpipe/backed_pipe_test.go
@@ -164,21 +164,30 @@ type blockingReconnector struct {
 	callCount   int
 	blockChan   <-chan struct{}
 	blockedChan chan struct{}
+	mu          sync.Mutex
+	signalOnce  sync.Once // Ensure we only signal once for the first actual reconnect
 }
 
 func (b *blockingReconnector) Reconnect(ctx context.Context, writerSeqNum uint64) (io.ReadWriteCloser, uint64, error) {
+	b.mu.Lock()
 	b.callCount++
+	currentCall := b.callCount
+	b.mu.Unlock()
 
-	if b.callCount == 1 {
+	if currentCall == 1 {
 		// Initial connect
 		return b.conn1, 0, nil
 	}
 
-	// Signal that we're about to block
-	select {
-	case b.blockedChan <- struct{}{}:
-	default:
-	}
+	// Signal that we're about to block, but only once for the first reconnect attempt
+	// This ensures we properly test singleflight deduplication
+	b.signalOnce.Do(func() {
+		select {
+		case b.blockedChan <- struct{}{}:
+		default:
+			// If channel is full, don't block
+		}
+	})
 
 	// For subsequent calls, block until channel is closed
 	select {
@@ -686,22 +695,57 @@ func TestBackedPipe_DuplicateReconnectionPrevention(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 1, *callCount, "should have exactly 1 call after initial connect")
 
-	// Start multiple concurrent ForceReconnect calls while the reconnect function is blocked
+	// We'll use channels to coordinate the test execution:
+	// 1. Start all goroutines but have them wait
+	// 2. Release the first one and wait for it to block
+	// 3. Release the others while the first is still blocked
+
+	const numConcurrent = 3
+	startSignals := make([]chan struct{}, numConcurrent)
+	startedSignals := make([]chan struct{}, numConcurrent)
+	for i := range startSignals {
+		startSignals[i] = make(chan struct{})
+		startedSignals[i] = make(chan struct{})
+	}
+
+	errors := make([]error, numConcurrent)
 	var wg sync.WaitGroup
-	errors := make([]error, 3)
 
-	for i := 0; i < 3; i++ {
+	// Start all goroutines
+	for i := 0; i < numConcurrent; i++ {
 		wg.Add(1)
 		go func(idx int) {
 			defer wg.Done()
+			// Wait for the signal to start
+			<-startSignals[idx]
+			// Signal that we're about to call ForceReconnect
+			close(startedSignals[idx])
 			errors[idx] = bp.ForceReconnect()
 		}(i)
 	}
 
-	// Wait for the reconnect function to signal that it's blocking
+	// Start the first ForceReconnect and wait for it to block
+	close(startSignals[0])
+	<-startedSignals[0]
+
+	// Wait for the first reconnect to actually start and block
 	testutil.RequireReceive(testCtx, t, blockedChan)
 
-	// Verify that exactly one more reconnect call was made (singleflight deduplication)
+	// Now start all the other ForceReconnect calls
+	// They should all join the same singleflight operation
+	for i := 1; i < numConcurrent; i++ {
+		close(startSignals[i])
+	}
+
+	// Wait for all additional goroutines to have started their calls
+	for i := 1; i < numConcurrent; i++ {
+		<-startedSignals[i]
+	}
+
+	// At this point, one reconnect has started and is blocked,
+	// and all other goroutines have called ForceReconnect and should be
+	// waiting on the same singleflight operation.
+	// Due to singleflight, only one reconnect should have been attempted.
 	require.Equal(t, 2, *callCount, "should have exactly 2 calls: initial connect + 1 reconnect due to singleflight")
 
 	// Release the blocking reconnect function
@@ -710,9 +754,9 @@ func TestBackedPipe_DuplicateReconnectionPrevention(t *testing.T) {
 	// Wait for all ForceReconnect calls to complete
 	wg.Wait()
 
-	// With singleflight, all concurrent calls should succeed (they share the same result)
+	// All calls should succeed (they share the same result from singleflight)
 	for i, err := range errors {
-		require.NoError(t, err, "ForceReconnect %d should succeed", i)
+		require.NoError(t, err, "ForceReconnect %d should succeed", i, err)
 	}
 
 	// Final verification: call count should still be exactly 2
diff --git a/agent/immortalstreams/backedpipe/backed_writer_test.go b/agent/immortalstreams/backedpipe/backed_writer_test.go
@@ -744,15 +744,16 @@ func TestBackedWriter_ConcurrentWriteAndReconnect(t *testing.T) {
 	// Wait for write to start
 	testutil.RequireReceive(ctx, t, writeStarted)
 
-	// Give the write a moment to actually start and get blocked on the mutex
-	time.Sleep(testutil.IntervalFast)
+	// Use a small timeout to ensure the write goroutine has a chance to get blocked
+	// on the mutex before we check if it's still blocked
+	writeCheckTimer := time.NewTimer(testutil.IntervalFast)
+	defer writeCheckTimer.Stop()
 
-	// Verify write is still blocked (reconnect is still in progress)
 	select {
 	case <-writeComplete:
 		t.Fatal("Write should be blocked during reconnect")
-	default:
-		// Good, write is still blocked
+	case <-writeCheckTimer.C:
+		// Write is still blocked after a reasonable wait
 	}
 
 	// Allow replay to complete, which will allow reconnect to finish
@@ -818,13 +819,24 @@ func TestBackedWriter_ConcurrentReconnectAndClose(t *testing.T) {
 	testutil.RequireReceive(ctx, t, reconnectStarted)
 
 	// Start Close() in a separate goroutine since it will block until Reconnect() completes
+	closeStarted := make(chan struct{}, 1)
 	closeComplete := make(chan error, 1)
 	go func() {
+		closeStarted <- struct{}{} // Signal that Close() is starting
 		closeComplete <- bw.Close()
 	}()
 
-	// Give Close() a moment to start and block on the mutex
-	time.Sleep(testutil.IntervalFast)
+	// Wait for Close() to start, then give it a moment to attempt to acquire the mutex
+	testutil.RequireReceive(ctx, t, closeStarted)
+	closeCheckTimer := time.NewTimer(testutil.IntervalFast)
+	defer closeCheckTimer.Stop()
+
+	select {
+	case <-closeComplete:
+		t.Fatal("Close should be blocked during reconnect")
+	case <-closeCheckTimer.C:
+		// Good, Close is still blocked after a reasonable wait
+	}
 
 	// Allow replay to complete so reconnection can finish
 	close(replayCanComplete)
@@ -891,8 +903,11 @@ func TestBackedWriter_MultipleWritesDuringReconnect(t *testing.T) {
 		testutil.RequireReceive(ctx, t, writesStarted)
 	}
 
-	// Give the writes a moment to actually begin execution
-	time.Sleep(testutil.IntervalFast)
+	// Use a timer to ensure all write goroutines have had a chance to start executing
+	// and potentially get blocked on the mutex before we start the reconnection
+	writesReadyTimer := time.NewTimer(testutil.IntervalFast)
+	defer writesReadyTimer.Stop()
+	<-writesReadyTimer.C
 
 	// Start reconnection with controlled replay
 	replayStarted := make(chan struct{})
