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 @@
+