[System.ClientModel] Add support for configurable max cache size in ClientCache (Azure#49298)

ShivangiReja · web-flow · commit c110e5485ced · 2025-04-09T09:47:12.000-07:00
[System.ClientModel] Configure maxCacheSize in ClientCache
diff --git a/sdk/ai/Azure.AI.Inference/src/Azure.AI.Inference.csproj b/sdk/ai/Azure.AI.Inference/src/Azure.AI.Inference.csproj
@@ -17,7 +17,10 @@
   </ItemGroup>
   <ItemGroup>
     <PackageReference Include="Azure.Core" />
-    <PackageReference Include="System.ClientModel" VersionOverride="1.4.0-beta.1" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\..\core\System.ClientModel\src\System.ClientModel.csproj" />
   </ItemGroup>
 
 </Project>
diff --git a/sdk/ai/Azure.AI.Projects/api/Azure.AI.Projects.net8.0.cs b/sdk/ai/Azure.AI.Projects/api/Azure.AI.Projects.net8.0.cs
@@ -460,11 +460,11 @@ public static partial class AIClientModelFactory
     }
     public partial class AIProjectClient : System.ClientModel.Primitives.ConnectionProvider
     {
-        protected AIProjectClient() { }
-        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) { }
-        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }
-        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) { }
-        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }
+        protected AIProjectClient() : base (default(int)) { }
+        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) : base (default(int)) { }
+        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }
+        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) : base (default(int)) { }
+        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }
         public virtual Azure.Core.Pipeline.HttpPipeline Pipeline { get { throw null; } }
         public virtual Azure.AI.Projects.AgentsClient GetAgentsClient(string apiVersion = "2024-07-01-preview") { throw null; }
         public override System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections() { throw null; }
@@ -476,6 +476,7 @@ public AIProjectClient(System.Uri endpoint, string subscriptionId, string resour
     public partial class AIProjectClientOptions : Azure.Core.ClientOptions
     {
         public AIProjectClientOptions(Azure.AI.Projects.AIProjectClientOptions.ServiceVersion version = Azure.AI.Projects.AIProjectClientOptions.ServiceVersion.V2024_07_01_Preview) { }
+        public int ClientCacheSize { get { throw null; } set { } }
         public enum ServiceVersion
         {
             V2024_07_01_Preview = 1,
diff --git a/sdk/ai/Azure.AI.Projects/api/Azure.AI.Projects.netstandard2.0.cs b/sdk/ai/Azure.AI.Projects/api/Azure.AI.Projects.netstandard2.0.cs
@@ -460,11 +460,11 @@ public static partial class AIClientModelFactory
     }
     public partial class AIProjectClient : System.ClientModel.Primitives.ConnectionProvider
     {
-        protected AIProjectClient() { }
-        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) { }
-        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }
-        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) { }
-        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }
+        protected AIProjectClient() : base (default(int)) { }
+        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) : base (default(int)) { }
+        public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }
+        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) : base (default(int)) { }
+        public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }
         public virtual Azure.Core.Pipeline.HttpPipeline Pipeline { get { throw null; } }
         public virtual Azure.AI.Projects.AgentsClient GetAgentsClient(string apiVersion = "2024-07-01-preview") { throw null; }
         public override System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections() { throw null; }
@@ -476,6 +476,7 @@ public AIProjectClient(System.Uri endpoint, string subscriptionId, string resour
     public partial class AIProjectClientOptions : Azure.Core.ClientOptions
     {
         public AIProjectClientOptions(Azure.AI.Projects.AIProjectClientOptions.ServiceVersion version = Azure.AI.Projects.AIProjectClientOptions.ServiceVersion.V2024_07_01_Preview) { }
+        public int ClientCacheSize { get { throw null; } set { } }
         public enum ServiceVersion
         {
             V2024_07_01_Preview = 1,
diff --git a/sdk/ai/Azure.AI.Projects/src/Azure.AI.Projects.csproj b/sdk/ai/Azure.AI.Projects/src/Azure.AI.Projects.csproj
@@ -14,10 +14,13 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Core" />
-    <PackageReference Include="System.ClientModel" VersionOverride="1.4.0-beta.1" />
     <PackageReference Include="Azure.Identity" />
   </ItemGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\..\..\core\System.ClientModel\src\System.ClientModel.csproj" />
+  </ItemGroup>
+
   <!-- Shared source from Azure.Core -->
   <ItemGroup>
     <Compile Include="$(AzureCoreSharedSources)ArrayBufferWriter.cs" LinkBase="Shared" />
diff --git a/sdk/ai/Azure.AI.Projects/src/Custom/AIProjectClient.cs b/sdk/ai/Azure.AI.Projects/src/Custom/AIProjectClient.cs
@@ -54,6 +54,7 @@ public AIProjectClient(string connectionString, TokenCredential credential, AIPr
         /// <exception cref="ArgumentNullException"> <paramref name="endpoint"/>, <paramref name="subscriptionId"/>, <paramref name="resourceGroupName"/>, <paramref name="projectName"/> or <paramref name="credential"/> is null. </exception>
         /// <exception cref="ArgumentException"> <paramref name="subscriptionId"/>, <paramref name="resourceGroupName"/> or <paramref name="projectName"/> is an empty string, and was expected to be non-empty. </exception>
         public AIProjectClient(Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, TokenCredential credential, AIProjectClientOptions options)
+            : base(options.ClientCacheSize)
         {
             Argument.AssertNotNull(endpoint, nameof(endpoint));
             Argument.AssertNotNullOrEmpty(subscriptionId, nameof(subscriptionId));
diff --git a/sdk/ai/Azure.AI.Projects/src/Custom/AIProjectClientOptions.cs b/sdk/ai/Azure.AI.Projects/src/Custom/AIProjectClientOptions.cs
@@ -0,0 +1,14 @@
+﻿// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Core;
+
+namespace Azure.AI.Projects
+{
+    /// <summary> Client options for AIProjectClient. </summary>
+    public partial class AIProjectClientOptions : ClientOptions
+    {
+        /// <summary> The size of the client cache. </summary>
+        public int ClientCacheSize { get; set; }
+    }
+}
diff --git a/sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs b/sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs
@@ -90,7 +90,7 @@ public enum ClientAuthenticationMethod
     }
     public partial class ClientCache
     {
-        public ClientCache() { }
+        public ClientCache(int maxSize = 100) { }
         public T GetClient<T>(System.Func<T> createClient, string? id) where T : class { throw null; }
     }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
@@ -182,7 +182,7 @@ public void AddRange(System.Collections.Generic.IEnumerable<System.ClientModel.P
     }
     public abstract partial class ConnectionProvider
     {
-        protected ConnectionProvider() { }
+        protected ConnectionProvider(int maxCacheSize = 100) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public System.ClientModel.Primitives.ClientCache Subclients { get { throw null; } }
         public abstract System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections();
diff --git a/sdk/core/System.ClientModel/api/System.ClientModel.net8.0.cs b/sdk/core/System.ClientModel/api/System.ClientModel.net8.0.cs
@@ -90,7 +90,7 @@ public enum ClientAuthenticationMethod
     }
     public partial class ClientCache
     {
-        public ClientCache() { }
+        public ClientCache(int maxSize = 100) { }
         public T GetClient<T>(System.Func<T> createClient, string? id) where T : class { throw null; }
     }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
@@ -182,7 +182,7 @@ public void AddRange(System.Collections.Generic.IEnumerable<System.ClientModel.P
     }
     public abstract partial class ConnectionProvider
     {
-        protected ConnectionProvider() { }
+        protected ConnectionProvider(int maxCacheSize = 100) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public System.ClientModel.Primitives.ClientCache Subclients { get { throw null; } }
         public abstract System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections();
diff --git a/sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs b/sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs
@@ -90,7 +90,7 @@ public enum ClientAuthenticationMethod
     }
     public partial class ClientCache
     {
-        public ClientCache() { }
+        public ClientCache(int maxSize = 100) { }
         public T GetClient<T>(System.Func<T> createClient, string? id) where T : class { throw null; }
     }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
@@ -182,7 +182,7 @@ public void AddRange(System.Collections.Generic.IEnumerable<System.ClientModel.P
     }
     public abstract partial class ConnectionProvider
     {
-        protected ConnectionProvider() { }
+        protected ConnectionProvider(int maxCacheSize = 100) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public System.ClientModel.Primitives.ClientCache Subclients { get { throw null; } }
         public abstract System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections();
diff --git a/sdk/core/System.ClientModel/src/Convenience/ClientCache.cs b/sdk/core/System.ClientModel/src/Convenience/ClientCache.cs
@@ -17,7 +17,16 @@ public class ClientCache
     private readonly Dictionary<(Type, string), ClientEntry> _clients = new();
     private readonly ReaderWriterLockSlim _lock = new(LockRecursionPolicy.SupportsRecursion);
 
-    private const int MaxCacheSize = 100;
+    private readonly int _maxSize;
+
+    /// <summary>
+    /// Initializes the ClientCache with a configurable cache size.
+    /// </summary>
+    /// <param name="maxSize">The maximum number of clients to store in the cache.</param>
+    public ClientCache(int maxSize = 100)
+    {
+        _maxSize = maxSize;
+    }
 
     /// <summary>
     /// Retrieves a client from the cache or creates a new one if it doesn't exist.
@@ -46,7 +55,7 @@ public T GetClient<T>(Func<T> createClient, string? id) where T : class
             _clients[key] = new ClientEntry(created, Stopwatch.GetTimestamp());
 
             // After insertion, if cache exceeds the limit, perform cleanup.
-            if (_clients.Count > MaxCacheSize)
+            if (_clients.Count > _maxSize)
             {
                 Cleanup();
             }
@@ -63,7 +72,7 @@ public T GetClient<T>(Func<T> createClient, string? id) where T : class
     /// </summary>
     private void Cleanup()
     {
-        int excess = _clients.Count - MaxCacheSize;
+        int excess = _clients.Count - _maxSize;
         if (excess <= 0)
         {
             return;
diff --git a/sdk/core/System.ClientModel/src/Convenience/ConnectionProvider.cs b/sdk/core/System.ClientModel/src/Convenience/ConnectionProvider.cs
@@ -12,6 +12,17 @@ namespace System.ClientModel.Primitives;
 /// </summary>
 public abstract class ConnectionProvider
 {
+    private readonly ClientCache _subclients;
+
+    /// <summary>
+    /// Initializes a new instance of the ConnectionProvider class.
+    /// </summary>
+    /// <param name="maxCacheSize">The maximum number of subclients to store in the cache.</param>
+    protected ConnectionProvider(int maxCacheSize = 100)
+    {
+        _subclients = new ClientCache(maxCacheSize);
+    }
+
     /// <summary>
     /// Retrieves the connection settings associated with a given connection ID.
     /// </summary>
@@ -26,8 +37,8 @@ public abstract class ConnectionProvider
     public abstract IEnumerable<ClientConnection> GetAllConnections();
 
     /// <summary>
-    /// Gets a cache for subclients to optimize performance by reusing client instances.
+    /// Gets the cache for subclients to optimize performance by reusing client instances.
     /// </summary>
     [EditorBrowsable(EditorBrowsableState.Never)]
-    public ClientCache Subclients { get; } = new ClientCache();
+    public ClientCache Subclients => _subclients;
 }
diff --git a/sdk/core/System.ClientModel/tests/Convenience/ClientCacheTests.cs b/sdk/core/System.ClientModel/tests/Convenience/ClientCacheTests.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Linq;
 using System.Reflection;
 using NUnit.Framework;
 
@@ -159,6 +160,56 @@ public void CacheShouldHandleDifferentClientIdsSeparately()
         Assert.IsTrue(clients!.ContainsKey((typeof(object), "client1")), "Client1 should be in the cache.");
         Assert.IsTrue(clients!.ContainsKey((typeof(object), "client2")), "Client2 should be in the cache.");
     }
+
+    [Test]
+    public void ClientCacheShouldRespectMaxCacheSize()
+    {
+        var cache = new ClientCache(maxSize: 3);
+
+        // Insert 3 clients
+        cache.GetClient(() => new DummyClient("A"), "A");
+        cache.GetClient(() => new DummyClient("B"), "B");
+        cache.GetClient(() => new DummyClient("C"), "C");
+
+        // Access client A to mark it as most recently used
+        cache.GetClient<DummyClient>(() => throw new Exception("Should not recreate A"), "A");
+
+        // Add a new client D to trigger eviction
+        cache.GetClient(() => new DummyClient("D"), "D");
+
+        // A should still be in the cache (was recently used), so this should not call the factory
+        var wasRecreated = false;
+        cache.GetClient(() =>
+        {
+            wasRecreated = true;
+            return new DummyClient("A");
+        }, "A");
+
+        Assert.False(wasRecreated, "Client A was unexpectedly recreated");
+
+        var clientsField = typeof(ClientCache).GetField("_clients", BindingFlags.NonPublic | BindingFlags.Instance);
+        var clients = clientsField?.GetValue(cache) as Dictionary<(Type, string), ClientEntry>;
+
+        Assert.IsNotNull(clients, "_clients dictionary should not be null");
+
+        var keys = clients!.Keys.Select(k => k.Item2).ToList();
+
+        // Verify that the cache contains the expected clients
+        CollectionAssert.AreEquivalent(new[] { "A", "C", "D" }, keys, "Cache should contain A, C, and D");
+
+        // Ensure B was evicted as it was the least recently used
+        Assert.IsFalse(keys.Contains("B"), "Client B should have been evicted.");
+    }
+}
+
+internal class DummyClient
+{
+    public string Id { get; }
+
+    public DummyClient(string id)
+    {
+        Id = id;
+    }
 }
 
 // Helper class to simulate a disposable client
diff --git a/sdk/openai/Azure.AI.OpenAI/src/Azure.AI.OpenAI.csproj b/sdk/openai/Azure.AI.OpenAI/src/Azure.AI.OpenAI.csproj
@@ -24,8 +24,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="System.ClientModel" />
-    <PackageReference Update="System.ClientModel" VersionOverride="1.4.0-beta.1" />
+    <ProjectReference Include="..\..\..\core\System.ClientModel\src\System.ClientModel.csproj" />
   </ItemGroup>
 
   <Choose>
diff --git a/sdk/search/Azure.Search.Documents/src/Azure.Search.Documents.csproj b/sdk/search/Azure.Search.Documents/src/Azure.Search.Documents.csproj
@@ -28,8 +28,11 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Core" />
-    <PackageReference Include="System.ClientModel" VersionOverride="1.4.0-beta.1" />
     <PackageReference Include="System.Threading.Channels" />
   </ItemGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\..\..\core\System.ClientModel\src\System.ClientModel.csproj" />
+  </ItemGroup>
+
 </Project>

Original file line number	Diff line number	Diff line change
`@@ -460,11 +460,11 @@ public static partial class AIClientModelFactory`
`460`	`460`	`}`
`461`	`461`	`public partial class AIProjectClient : System.ClientModel.Primitives.ConnectionProvider`
`462`	`462`	`{`
`463`		`- protected AIProjectClient() { }`
`464`		`- public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) { }`
`465`		`- public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }`
`466`		`- public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) { }`
`467`		`- public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) { }`
	`463`	`+ protected AIProjectClient() : base (default(int)) { }`
	`464`	`+ public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential = null) : base (default(int)) { }`
	`465`	`+ public AIProjectClient(string connectionString, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }`
	`466`	`+ public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential) : base (default(int)) { }`
	`467`	`+ public AIProjectClient(System.Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, Azure.Core.TokenCredential credential, Azure.AI.Projects.AIProjectClientOptions options) : base (default(int)) { }`
`468`	`468`	`public virtual Azure.Core.Pipeline.HttpPipeline Pipeline { get { throw null; } }`
`469`	`469`	`public virtual Azure.AI.Projects.AgentsClient GetAgentsClient(string apiVersion = "2024-07-01-preview") { throw null; }`
`470`	`470`	`public override System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections() { throw null; }`
`@@ -476,6 +476,7 @@ public AIProjectClient(System.Uri endpoint, string subscriptionId, string resour`
`476`	`476`	`public partial class AIProjectClientOptions : Azure.Core.ClientOptions`
`477`	`477`	`{`
`478`	`478`	`public AIProjectClientOptions(Azure.AI.Projects.AIProjectClientOptions.ServiceVersion version = Azure.AI.Projects.AIProjectClientOptions.ServiceVersion.V2024_07_01_Preview) { }`
	`479`	`+ public int ClientCacheSize { get { throw null; } set { } }`
`479`	`480`	`public enum ServiceVersion`
`480`	`481`	`{`
`481`	`482`	`V2024_07_01_Preview = 1,`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,7 @@ public AIProjectClient(string connectionString, TokenCredential credential, AIPr`
`54`	`54`	`/// <exception cref="ArgumentNullException"> <paramref name="endpoint"/>, <paramref name="subscriptionId"/>, <paramref name="resourceGroupName"/>, <paramref name="projectName"/> or <paramref name="credential"/> is null. </exception>`
`55`	`55`	`/// <exception cref="ArgumentException"> <paramref name="subscriptionId"/>, <paramref name="resourceGroupName"/> or <paramref name="projectName"/> is an empty string, and was expected to be non-empty. </exception>`
`56`	`56`	`public AIProjectClient(Uri endpoint, string subscriptionId, string resourceGroupName, string projectName, TokenCredential credential, AIProjectClientOptions options)`
	`57`	`+ : base(options.ClientCacheSize)`
`57`	`58`	`{`
`58`	`59`	`Argument.AssertNotNull(endpoint, nameof(endpoint));`
`59`	`60`	`Argument.AssertNotNullOrEmpty(subscriptionId, nameof(subscriptionId));`
Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ public enum ClientAuthenticationMethod`
`90`	`90`	`}`
`91`	`91`	`public partial class ClientCache`
`92`	`92`	`{`
`93`		`- public ClientCache() { }`
	`93`	`+ public ClientCache(int maxSize = 100) { }`
`94`	`94`	`public T GetClient<T>(System.Func<T> createClient, string? id) where T : class { throw null; }`
`95`	`95`	`}`
`96`	`96`	`[System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]`
`@@ -182,7 +182,7 @@ public void AddRange(System.Collections.Generic.IEnumerable<System.ClientModel.P`
`182`	`182`	`}`
`183`	`183`	`public abstract partial class ConnectionProvider`
`184`	`184`	`{`
`185`		`- protected ConnectionProvider() { }`
	`185`	`+ protected ConnectionProvider(int maxCacheSize = 100) { }`
`186`	`186`	`[System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]`
`187`	`187`	`public System.ClientModel.Primitives.ClientCache Subclients { get { throw null; } }`
`188`	`188`	`public abstract System.Collections.Generic.IEnumerable<System.ClientModel.Primitives.ClientConnection> GetAllConnections();`