Handle stale cache in CosmosClient

Add a retry for unintentional distributed transactions

Author: AndriySvyryd

test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs

@@ -317,6 +317,10 @@ public override Task CleanAsync(DbContext context, bool createTables = true)
     private async Task CleanAsyncImpl(DbContext context, bool createTables)
     {
         var created = await EnsureCreatedAsync(context).ConfigureAwait(false);
+
+        // Containers are deleted and recreated below only when the database already existed. A freshly-created
+        // database has brand-new containers whose metadata cannot be stale.
+        var containersRecreated = !created;
         try
         {
             if (!created)
@@ -334,12 +338,23 @@ private async Task CleanAsyncImpl(DbContext context, bool createTables)
                 created = await context.Database.EnsureCreatedAsync().ConfigureAwait(false);
                 if (!created)
                 {
+                    if (containersRecreated)
+                    {
+                        await RefreshContainerMetadataAsync(context).ConfigureAwait(false);
+                    }
+
                     await SeedAsync(context).ConfigureAwait(false);
                 }
             }
             else
             {
                 await CreateContainersAsync(context).ConfigureAwait(false);
+
+                if (containersRecreated)
+                {
+                    await RefreshContainerMetadataAsync(context).ConfigureAwait(false);
+                }
+
                 await SeedAsync(context).ConfigureAwait(false);
             }
         }
@@ -357,6 +372,68 @@ private async Task CleanAsyncImpl(DbContext context, bool createTables)
         }
     }
 
+    // Deleting and recreating a container gives it a new resource id (_rid), but the shared CosmosClient caches each
+    // container's _rid (and the partition key ranges keyed by it) by container name. The first operation on the
+    // recreated container - whether it is the reseed below or the next test's first query - can then fail with
+    // "NotFound (404) ... GetTargetPartitionKeyRanges ... failed due to stale cache", and the query pipeline does not
+    // transparently refresh and retry. Prime the client's caches here, right after recreating the containers and before
+    // any test touches them. Each attempt re-reads the container by name (refreshing the collection cache with the new
+    // _rid) and then runs a query - the exact path that otherwise fails. Because the recreate can take a moment to
+    // become consistent on the emulator gateway, retry with a short delay until a query succeeds. This is the single
+    // place that handles the stale-metadata problem for the emulator; it is fully best-effort and must never fail the
+    // clean, so all errors are ultimately swallowed.
+    private async Task RefreshContainerMetadataAsync(DbContext context)
+    {
+        if (CosmosTestEnvironment.UseTokenCredential)
+        {
+            return;
+        }
+
+        const int maxAttempts = 10;
+
+        try
+        {
+            var cosmosClient = context.Database.GetCosmosClient();
+            var database = cosmosClient.GetDatabase(Name);
+            var model = context.GetService<IDesignTimeModel>().Model;
+            var countQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c");
+
+            foreach (var containerProperties in GetContainersToCreate(model))
+            {
+                var container = database.GetContainer(containerProperties.Id);
+
+                for (var attempt = 1; attempt <= maxAttempts; attempt++)
+                {
+                    try
+                    {
+                        await container.ReadContainerAsync().ConfigureAwait(false);
+
+                        using var iterator = container.GetItemQueryIterator<int>(countQuery);
+                        while (iterator.HasMoreResults)
+                        {
+                            await iterator.ReadNextAsync().ConfigureAwait(false);
+                        }
+
+                        break;
+                    }
+                    catch when (attempt < maxAttempts)
+                    {
+                        // The recreate has not become consistent yet; wait for the gateway to catch up and retry.
+                        await Task.Delay(200).ConfigureAwait(false);
+                    }
+                    catch
+                    {
+                        // Best-effort priming; never let it fail the clean.
+                    }
+                }
+            }
+        }
+        catch
+        {
+            // Never let cache priming fail the clean.
+        }
+    }
+
     private async Task CreateContainersAsync(DbContext context)
     {
         var databaseAccount = await GetDBAccount().ConfigureAwait(false);

test/EFCore.Relational.Specification.Tests/TransactionTestBase.cs

@@ -336,41 +336,85 @@ public virtual async Task SaveChanges_uses_ambient_transaction_with_connectionSt
             return;
         }
 
-        DbConnection connection;
-        using (var context = CreateContextWithConnectionString())
+        DbConnection connection = null;
+
+        await RetryOnDistributedTransactionNotSupportedAsync(async () =>
         {
-            using (TestUtilities.TestStore.CreateTransactionScope())
+            using (var context = CreateContextWithConnectionString())
             {
-                context.Database.AutoTransactionBehavior = autoTransactionBehavior;
+                using (TestUtilities.TestStore.CreateTransactionScope())
+                {
+                    context.Database.AutoTransactionBehavior = autoTransactionBehavior;
 
-                connection = context.Database.GetDbConnection();
-                Assert.Equal(ConnectionState.Closed, connection.State);
+                    connection = context.Database.GetDbConnection();
+                    Assert.Equal(ConnectionState.Closed, connection.State);
 
-                await context.AddAsync(
-                    new TransactionCustomer { Id = 77, Name = "Bobble" });
+                    await context.AddAsync(
+                        new TransactionCustomer { Id = 77, Name = "Bobble" });
 
-                context.Entry(context.Set<TransactionCustomer>().OrderBy(c => c.Id).Last()).State = EntityState.Added;
+                    context.Entry(context.Set<TransactionCustomer>().OrderBy(c => c.Id).Last()).State = EntityState.Added;
 
-                if (async)
-                {
-                    await Assert.ThrowsAsync<DbUpdateException>(() => context.SaveChangesAsync());
-                }
-                else
-                {
-                    Assert.Throws<DbUpdateException>(() => context.SaveChanges());
-                }
+                    // SaveChanges is expected to fail with DbUpdateException (duplicate key).
+                    try
+                    {
+                        if (async)
+                        {
+                            await context.SaveChangesAsync();
+                        }
+                        else
+                        {
+                            context.SaveChanges();
+                        }
+
+                        Assert.Fail("Expected a DbUpdateException to be thrown.");
+                    }
+                    catch (PlatformNotSupportedException)
+                    {
+                        // When the ambient transaction is promoted to a distributed one and throw PlatformNotSupportedException let it
+                        // propagate so the whole operation is retried.
+                        throw;
+                    }
+                    catch (Exception exception)
+                    {
+                        Assert.IsType<DbUpdateException>(exception);
+                    }
 
-                Assert.Equal(ConnectionState.Closed, connection.State);
+                    Assert.Equal(ConnectionState.Closed, connection.State);
 
-                context.Database.AutoTransactionBehavior = AutoTransactionBehavior.WhenNeeded;
+                    context.Database.AutoTransactionBehavior = AutoTransactionBehavior.WhenNeeded;
+                }
             }
-        }
+        });
 
         Assert.Equal(ConnectionState.Closed, connection.State);
 
         AssertStoreInitialState();
     }
 
+    // Sometimes an ambient/enlisted transaction ends up spanning more than one physical connection due to concurrency issues
+    // and is promoted to a distributed transaction, which throws PlatformNotSupportedException.
+    // Retry the operation a few times (with exponential back-off) so a run that keeps to a single connection (and
+    // therefore reaches the behavior under test) is used.
+    private static async Task RetryOnDistributedTransactionNotSupportedAsync(Func<Task> test)
+    {
+        const int maxAttempts = 3;
+        var delay = TimeSpan.FromSeconds(1);
+
+        for (var attempt = 1; ; attempt++)
+        {
+            try
+            {
+                await test();
+                return;
+            }
+            catch (PlatformNotSupportedException) when (attempt < maxAttempts)
+            {
+                await Task.Delay(delay);
+                delay *= 2;
+            }
+        }
+    }
+
     [Theory, InlineData(true), InlineData(false)]
     public virtual void SaveChanges_throws_for_suppressed_ambient_transactions(bool connectionString)
     {