[rc2] Fix race condition in SingletonCosmosClientWrapper.Client (#36767)

Fixes #36714

Co-authored-by: AndriySvyryd <6539701+AndriySvyryd@users.noreply.github.com>

Author: Copilot

src/EFCore.Cosmos/Metadata/Conventions/Internal/CosmosConventionSetBuilder.cs

@@ -82,7 +82,7 @@ private static IServiceScope CreateServiceScope()
         var serviceProvider = new ServiceCollection()
             .AddEntityFrameworkCosmos()
             .AddDbContext<DbContext>((p, o) =>
-                o.UseCosmos("localhost", "_", "_")
+                o.UseCosmos("https://localhost:8081", "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", "_")
                     .UseInternalServiceProvider(p))
             .BuildServiceProvider();
 

src/EFCore.Cosmos/Storage/Internal/SingletonCosmosClientWrapper.cs

@@ -16,12 +16,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 public class SingletonCosmosClientWrapper : ISingletonCosmosClientWrapper
 {
     private static readonly string UserAgent = " Microsoft.EntityFrameworkCore.Cosmos/" + ProductInfo.GetVersion();
-    private readonly CosmosClientOptions _options;
-    private readonly string? _endpoint;
-    private readonly string? _key;
-    private readonly string? _connectionString;
-    private readonly TokenCredential? _tokenCredential;
-    private CosmosClient? _client;
+    private readonly CosmosClient _client;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -31,10 +26,6 @@ public class SingletonCosmosClientWrapper : ISingletonCosmosClientWrapper
     /// </summary>
     public SingletonCosmosClientWrapper(ICosmosSingletonOptions options)
     {
-        _endpoint = options.AccountEndpoint;
-        _key = options.AccountKey;
-        _connectionString = options.ConnectionString;
-        _tokenCredential = options.TokenCredential;
         var configuration = new CosmosClientOptions { ApplicationName = UserAgent, Serializer = new JsonCosmosSerializer() };
 
         if (options.Region != null)
@@ -97,7 +88,13 @@ public SingletonCosmosClientWrapper(ICosmosSingletonOptions options)
             configuration.HttpClientFactory = options.HttpClientFactory;
         }
 
-        _options = configuration;
+        _client = options switch
+        {
+            { ConnectionString: not null and not "" } => new CosmosClient(options.ConnectionString, configuration),
+            { TokenCredential: not null } => new CosmosClient(options.AccountEndpoint, options.TokenCredential, configuration),
+            { AccountEndpoint: not null } => new CosmosClient(options.AccountEndpoint, options.AccountKey, configuration),
+            _ => throw new InvalidOperationException(CosmosStrings.ConnectionInfoMissing)
+        };
     }
 
     /// <summary>
@@ -106,14 +103,7 @@ public SingletonCosmosClientWrapper(ICosmosSingletonOptions options)
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public virtual CosmosClient Client
-        => _client ??= string.IsNullOrEmpty(_connectionString)
-            ? _tokenCredential == null
-                ? _endpoint == null
-                    ? throw new InvalidOperationException(CosmosStrings.ConnectionInfoMissing)
-                    : new CosmosClient(_endpoint, _key, _options)
-                : new CosmosClient(_endpoint, _tokenCredential, _options)
-            : new CosmosClient(_connectionString, _options);
+    public virtual CosmosClient Client => _client;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -123,7 +113,6 @@ public virtual CosmosClient Client
     /// </summary>
     public virtual void Dispose()
     {
-        _client?.Dispose();
-        _client = null;
+        _client.Dispose();
     }
 }

test/EFCore.Cosmos.FunctionalTests/ConfigPatternsCosmosTest.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Concurrent;
 using Microsoft.Azure.Cosmos;
 
 // ReSharper disable UnusedAutoPropertyAccessor.Local
@@ -124,6 +125,41 @@ public async Task Should_throw_if_specified_connection_mode_is_wrong()
         });
     }
 
+    [ConditionalFact]
+    public async Task Cosmos_client_instance_is_thread_safe()
+    {
+        await using var testDatabase = await CosmosTestStore.CreateInitializedAsync(DatabaseName);
+        var options = CreateOptions(testDatabase);
+
+        var threadCount = Environment.ProcessorCount;
+        const int iterationsPerThread = 2;
+        var clients = new ConcurrentBag<CosmosClient>();
+        var tasks = new List<Task>();
+
+        for (var i = 0; i < threadCount; i++)
+        {
+            tasks.Add(Task.Run(async () =>
+            {
+                for (var j = 0; j < iterationsPerThread; j++)
+                {
+                    await Task.Yield(); // Force context switching
+                    using var context = new CustomerContext(options);
+                    var client = context.Database.GetCosmosClient();
+                    clients.Add(client);
+                }
+            }));
+        }
+
+        await Task.WhenAll(tasks);
+
+        // All retrieved clients should be the same instance
+        var clientsArray = clients.ToArray();
+        Assert.Equal(threadCount * iterationsPerThread, clientsArray.Length);
+        
+        var uniqueClients = clientsArray.Distinct().ToArray();
+        Assert.Single(uniqueClients); // Should only have one unique client instance
+    }
+
     private DbContextOptions CreateOptions(CosmosTestStore testDatabase, Action<DbContextOptionsBuilder> configure = null)
     {
         var builder = Fixture.AddOptions(testDatabase.AddProviderOptions(new DbContextOptionsBuilder()))

test/EFCore.Cosmos.FunctionalTests/ConnectionSpecificationTest.cs

@@ -36,11 +36,10 @@ protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
     public async Task Throws_for_missing_connection_info()
     {
         using var context = new NoConnectionContext();
-        var creator = context.GetService<IDatabaseCreator>();
 
         Assert.Equal(
             CosmosStrings.ConnectionInfoMissing,
-            (await Assert.ThrowsAsync<InvalidOperationException>(() => creator.EnsureDeletedAsync())).Message);
+            (await Assert.ThrowsAsync<InvalidOperationException>(() => context.GetService<IDatabaseCreator>().EnsureDeletedAsync())).Message);
     }
 
     public class NoConnectionContext : DbContext

test/EFCore.Cosmos.Tests/Extensions/PartitionKeyBuilderExtensionsTest.cs

@@ -120,7 +120,7 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
         }
 
         protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
-            => optionsBuilder.UseCosmos("localhost", "_", "_");
+            => optionsBuilder.UseCosmos("https://localhost:8081", "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", "_");
     }
 
     public class Customer1