diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 9202b36f2a..d0c90d5cea 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -17,6 +17,7 @@ + @@ -180,4 +181,4 @@ runtime; build; native; contentfiles; analyzers; buildtransitive - \ No newline at end of file + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/MemoryCacheExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/MemoryCacheExtensions.cs index 670223ef4f..d7ee0b3e43 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/MemoryCacheExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/MemoryCacheExtensions.cs @@ -1,11 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Concurrent; using System.Diagnostics; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.Caching.Memory; +using AsyncKeyedLock; namespace Microsoft.Agents.AI.Hosting.OpenAI; @@ -19,7 +19,7 @@ namespace Microsoft.Agents.AI.Hosting.OpenAI; /// internal static class MemoryCacheExtensions { - private static readonly ConcurrentDictionary<(IMemoryCache, object), SemaphoreSlim> s_semaphores = new(); + private static readonly AsyncKeyedLocker<(IMemoryCache, object)> s_semaphores = new(); /// /// Atomically gets the value associated with this key if it exists, or generates a new entry @@ -44,28 +44,7 @@ public static async Task GetOrCreateAtomicAsync( return (T)value; } - // Get or create a semaphore for this cache key - bool isOwner = false; - var semaphoreKey = (memoryCache, key); - if (!s_semaphores.TryGetValue(semaphoreKey, out SemaphoreSlim? semaphore)) - { - SemaphoreSlim? createdSemaphore = null; - semaphore = s_semaphores.GetOrAdd(semaphoreKey, _ => createdSemaphore = new SemaphoreSlim(1)); - - // If we created the semaphore that made it into the dictionary, we're the owner - if (ReferenceEquals(createdSemaphore, semaphore)) - { - isOwner = true; - } - else - { - // Our semaphore wasn't the one stored, so dispose it - createdSemaphore?.Dispose(); - } - } - - await semaphore.WaitAsync(cancellationToken).ConfigureAwait(false); - try + using (await s_semaphores.LockAsync((memoryCache, key), cancellationToken).ConfigureAwait(false)) { // Double-check: another thread might have created the value while we were waiting if (!memoryCache.TryGetValue(key, out value)) @@ -76,20 +55,8 @@ public static async Task GetOrCreateAtomicAsync( Debug.Assert(value is not null); return (T)value; } - Debug.Assert(value is not null); return (T)value; } - finally - { - // If we were the owner of the semaphore, remove it from the dictionary - // This prevents memory leaks from accumulating semaphores for evicted cache entries - if (isOwner) - { - s_semaphores.TryRemove(semaphoreKey, out _); - } - - semaphore.Release(); - } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj index 923f8e3eb6..a77263f6fb 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj @@ -16,6 +16,7 @@ +