Fix flaky tests

engine/access/rpc/connection/cache.go

@@ -7,7 +7,6 @@ import (
 	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/onflow/crypto"
 	"github.com/rs/zerolog"
-	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/connectivity"
 
@@ -16,20 +15,25 @@ import (
 
 // CachedClient represents a gRPC client connection that is cached for reuse.
 type CachedClient struct {
-	conn    *grpc.ClientConn
 	address string
 	cfg     Config
 
-	cache          *Cache
-	closeRequested *atomic.Bool
-	wg             sync.WaitGroup
-	mu             sync.RWMutex
+	cache *Cache
+
+	connMu sync.RWMutex
+	conn   *grpc.ClientConn
+
+	closeRequested bool
+	// wgMu mutex is needed to protect the workgroup from being added to
+	// if we are in a closeRequested state.
+	wgMu sync.RWMutex
+	wg   sync.WaitGroup
 }
 
 // ClientConn returns the underlying gRPC client connection.
 func (cc *CachedClient) ClientConn() *grpc.ClientConn {
-	cc.mu.RLock()
-	defer cc.mu.RUnlock()
+	cc.connMu.RLock()
+	defer cc.connMu.RUnlock()
 	return cc.conn
 }
 
@@ -40,12 +44,22 @@ func (cc *CachedClient) Address() string {
 
 // CloseRequested returns true if the CachedClient has been marked for closure.
 func (cc *CachedClient) CloseRequested() bool {
-	return cc.closeRequested.Load()
+	cc.wgMu.RLock()
+	defer cc.wgMu.RUnlock()
+
+	return cc.closeRequested
 }
 
 // AddRequest increments the in-flight request counter for the CachedClient.
 // It returns a function that should be called when the request completes to decrement the counter
 func (cc *CachedClient) AddRequest() func() {
+	cc.wgMu.RLock()
+	defer cc.wgMu.RUnlock()
+
+	// if close is requested, cc.wg might already be done
+	if cc.closeRequested {
+		return func() {}
+	}
 	cc.wg.Add(1)
 	return cc.wg.Done
 }
@@ -61,15 +75,16 @@ func (cc *CachedClient) Invalidate() {
 // Close closes the CachedClient connection. It marks the connection for closure and waits asynchronously for ongoing
 // requests to complete before closing the connection.
 func (cc *CachedClient) Close() {
-	// Mark the connection for closure
-	if !cc.closeRequested.CompareAndSwap(false, true) {
-		return
-	}
+	func() {
+		cc.wgMu.Lock()
+		defer cc.wgMu.Unlock()
+		cc.closeRequested = true
+	}()
 
 	// Obtain the lock to ensure that any connection attempts have completed
-	cc.mu.RLock()
+	cc.connMu.RLock()
 	conn := cc.conn
-	cc.mu.RUnlock()
+	cc.connMu.RUnlock()
 
 	// If the initial connection attempt failed, conn will be nil
 	if conn == nil {
@@ -127,10 +142,9 @@ func (c *Cache) GetConnected(
 	connectFn func(string, Config, crypto.PublicKey, *CachedClient) (*grpc.ClientConn, error),
 ) (*CachedClient, error) {
 	client := &CachedClient{
-		address:        address,
-		cfg:            cfg,
-		closeRequested: atomic.NewBool(false),
-		cache:          c,
+		address: address,
+		cfg:     cfg,
+		cache:   c,
 	}
 
 	// Note: PeekOrAdd does not "visit" the existing entry, so we need to call Get explicitly
@@ -145,8 +159,8 @@ func (c *Cache) GetConnected(
 		c.metrics.ConnectionAddedToPool()
 	}
 
-	client.mu.Lock()
-	defer client.mu.Unlock()
+	client.connMu.Lock()
+	defer client.connMu.Unlock()
 
 	// after getting the lock, check if the connection is still active
 	if client.conn != nil && client.conn.GetState() != connectivity.Shutdown {

engine/access/rpc/connection/cache_test.go

@@ -20,34 +20,30 @@ import (
 func TestCachedClientShutdown(t *testing.T) {
 	// Test that a completely uninitialized client can be closed without panics
 	t.Run("uninitialized client", func(t *testing.T) {
-		client := &CachedClient{
-			closeRequested: atomic.NewBool(false),
-		}
+		client := &CachedClient{}
 		client.Close()
-		assert.True(t, client.closeRequested.Load())
+		assert.True(t, client.CloseRequested())
 	})
 
 	// Test closing a client with no outstanding requests
 	// Close() should return quickly
 	t.Run("with no outstanding requests", func(t *testing.T) {
 		client := &CachedClient{
-			closeRequested: atomic.NewBool(false),
-			conn:           setupGRPCServer(t),
+			conn: setupGRPCServer(t),
 		}
 
 		unittest.RequireReturnsBefore(t, func() {
 			client.Close()
 		}, 100*time.Millisecond, "client timed out closing connection")
 
-		assert.True(t, client.closeRequested.Load())
+		assert.True(t, client.CloseRequested())
 	})
 
 	// Test closing a client with outstanding requests waits for requests to complete
 	// Close() should block until the request completes
 	t.Run("with some outstanding requests", func(t *testing.T) {
 		client := &CachedClient{
-			closeRequested: atomic.NewBool(false),
-			conn:           setupGRPCServer(t),
+			conn: setupGRPCServer(t),
 		}
 		done := client.AddRequest()
 
@@ -62,17 +58,17 @@ func TestCachedClientShutdown(t *testing.T) {
 			client.Close()
 		}, 100*time.Millisecond, "client timed out closing connection")
 
-		assert.True(t, client.closeRequested.Load())
+		assert.True(t, client.CloseRequested())
 		assert.True(t, doneCalled.Load())
 	})
 
 	// Test closing a client that is already closing does not block
 	// Close() should return immediately
 	t.Run("already closing", func(t *testing.T) {
 		client := &CachedClient{
-			closeRequested: atomic.NewBool(true), // close already requested
-			conn:           setupGRPCServer(t),
+			conn: setupGRPCServer(t),
 		}
+		client.Close()
 		done := client.AddRequest()
 
 		doneCalled := atomic.NewBool(false)
@@ -89,23 +85,21 @@ func TestCachedClientShutdown(t *testing.T) {
 			client.Close()
 		}, 10*time.Millisecond, "client timed out closing connection")
 
-		assert.True(t, client.closeRequested.Load())
+		assert.True(t, client.CloseRequested())
 		assert.False(t, doneCalled.Load())
 	})
 
 	// Test closing a client that is locked during connection setup
 	// Close() should wait for the lock before shutting down
 	t.Run("connection setting up", func(t *testing.T) {
-		client := &CachedClient{
-			closeRequested: atomic.NewBool(false),
-		}
+		client := &CachedClient{}
 
 		// simulate an in-progress connection setup
-		client.mu.Lock()
+		client.connMu.Lock()
 
 		go func() {
 			// unlock after setting up the connection
-			defer client.mu.Unlock()
+			defer client.connMu.Unlock()
 
 			// pause before setting the connection to cause client.Close() to block
 			time.Sleep(100 * time.Millisecond)
@@ -117,7 +111,7 @@ func TestCachedClientShutdown(t *testing.T) {
 			client.Close()
 		}, 500*time.Millisecond, "client timed out closing connection")
 
-		assert.True(t, client.closeRequested.Load())
+		assert.True(t, client.CloseRequested())
 		assert.NotNil(t, client.conn)
 	})
 }

engine/access/rpc/connection/connection_test.go

@@ -714,7 +714,7 @@ func TestEvictingCacheClients(t *testing.T) {
 
 		// Invalidate marks the connection for closure asynchronously, so give it some time to run
 		require.Eventually(t, func() bool {
-			return cachedClient.closeRequested.Load()
+			return cachedClient.CloseRequested()
 		}, 100*time.Millisecond, 10*time.Millisecond, "client timed out closing connection")
 
 		// Call a gRPC method on the client, requests should be blocked since the connection is invalidated

ledger/complete/compactor.go

@@ -282,7 +282,14 @@ Loop:
 	}
 	c.logger.Info().Msg("Finished draining trie update channel in compactor on shutdown")
 
-	// Don't wait for checkpointing to finish because it might take too long.
+	// Don't wait for checkpointing to finish if it takes more than 10ms because it might take too long.
+	// We wait at least a little bit to make the tests a lot more stable.
+	if !checkpointSem.TryAcquire(1) {
+		select {
+		case <-checkpointResultCh:
+		case <-time.After(10 * time.Millisecond):
+		}
+	}
 }
 
 // checkpoint creates checkpoint of tries snapshot,

network/test/cohort2/unicast_authorization_test.go

@@ -463,7 +463,11 @@ func (u *UnicastAuthorizationTestSuite) TestUnicastAuthorization_ReceiverHasNoSu
 	err = senderCon.Unicast(&libp2pmessage.TestMessage{
 		Text: string("hello"),
 	}, u.receiverID.NodeID)
-	require.NoError(u.T(), err)
+	if err != nil {
+		// It can happen that the receiver resets before the sender closes,
+		// in which case the error will be "stream reset"
+		require.ErrorContains(u.T(), err, "stream reset", "expected stream-related error when receiver has no subscription")
+	}
 
 	// wait for slashing violations consumer mock to invoke run func and close ch if expected method call happens
 	unittest.RequireCloseBefore(u.T(), u.waitCh, u.channelCloseDuration, "could close ch on time")