fb

Author: BrennanConroy

src/Servers/Kestrel/Core/src/Internal/PinnedBlockMemoryPoolFactory.cs

@@ -27,10 +27,10 @@ public MemoryPool<byte> Create()
 
         _pools.TryAdd(pool, pool);
 
-        pool.DisposeCallback = (self) =>
+        pool.OnPoolDisposed(static (state, self) =>
         {
-            _pools.TryRemove(self, out _);
-        };
+            ((ConcurrentDictionary<PinnedBlockMemoryPool, PinnedBlockMemoryPool>)state!).TryRemove(self, out _);
+        }, _pools);
 
         return pool;
     }

src/Shared/Buffers.MemoryPool/DefaultMemoryPoolFactory.cs

@@ -10,17 +10,18 @@ namespace System.Buffers;
 
 #nullable enable
 
-internal sealed class DefaultMemoryPoolFactory : IMemoryPoolFactory<byte>, IDisposable
+internal sealed class DefaultMemoryPoolFactory : IMemoryPoolFactory<byte>, IAsyncDisposable
 {
     private readonly IMeterFactory _meterFactory;
     private readonly ConcurrentDictionary<PinnedBlockMemoryPool, PinnedBlockMemoryPool> _pools = new();
     private readonly PeriodicTimer _timer;
+    private readonly Task _timerTask;
 
     public DefaultMemoryPoolFactory(IMeterFactory? meterFactory = null)
     {
         _meterFactory = meterFactory ?? NoopMeterFactory.Instance;
         _timer = new PeriodicTimer(TimeSpan.FromSeconds(10));
-        _ = Task.Run(async () =>
+        _timerTask = Task.Run(async () =>
         {
             try
             {
@@ -45,17 +46,18 @@ public MemoryPool<byte> Create()
 
         _pools.TryAdd(pool, pool);
 
-        pool.DisposeCallback = (self) =>
+        pool.OnPoolDisposed(static (state, self) =>
         {
-            _pools.TryRemove(self, out _);
-        };
+            ((ConcurrentDictionary<PinnedBlockMemoryPool, PinnedBlockMemoryPool>)state!).TryRemove(self, out _);
+        }, _pools);
 
         return pool;
     }
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
         _timer.Dispose();
+        await _timerTask;
     }
 
     private sealed class NoopMeterFactory : IMeterFactory

src/Shared/Buffers.MemoryPool/PinnedBlockMemoryPool.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Concurrent;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.Metrics;
 
 #nullable enable
@@ -30,11 +31,6 @@ internal sealed class PinnedBlockMemoryPool : MemoryPool<byte>, IThreadPoolWorkI
     /// </summary>
     public static int BlockSize => _blockSize;
 
-    /// <summary>
-    /// Optional callback to call when pool is being disposed.
-    /// </summary>
-    public Action<PinnedBlockMemoryPool>? DisposeCallback { get; set; }
-
     /// <summary>
     /// Thread-safe collection of blocks which are currently in the pool. A slab will pre-allocate all of the block tracking objects
     /// and add them to this collection. When memory is requested it is taken from here first, and when it is returned it is re-added.
@@ -57,6 +53,9 @@ internal sealed class PinnedBlockMemoryPool : MemoryPool<byte>, IThreadPoolWorkI
 
     private readonly object _disposeSync = new object();
 
+    private Action<object?, PinnedBlockMemoryPool>? _onPoolDisposed;
+    private object? _onPoolDisposedState;
+
     /// <summary>
     /// This default value passed in to Rent to use the default value for the pool.
     /// </summary>
@@ -72,6 +71,15 @@ public PinnedBlockMemoryPool(IMeterFactory meterFactory)
         _metrics = new(meterFactory);
     }
 
+    /// <summary>
+    /// Register a callback to call when the pool is being disposed.
+    /// </summary>
+    public void OnPoolDisposed(Action<object?, PinnedBlockMemoryPool> onPoolDisposed, object? state = null)
+    {
+        _onPoolDisposed = onPoolDisposed;
+        _onPoolDisposedState = state;
+    }
+
     public override IMemoryOwner<byte> Rent(int size = AnySize)
     {
         if (size > _blockSize)
@@ -214,7 +222,7 @@ protected override void Dispose(bool disposing)
         {
             _isDisposed = true;
 
-            DisposeCallback?.Invoke(this);
+            _onPoolDisposed?.Invoke(_onPoolDisposedState, this);
 
             if (disposing)
             {