fb

Author: BrennanConroy

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

@@ -6,25 +6,28 @@
 using System.Diagnostics.Metrics;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 
 internal sealed class PinnedBlockMemoryPoolFactory : IMemoryPoolFactory<byte>, IHeartbeatHandler
 {
     private readonly IMeterFactory _meterFactory;
+    private readonly ILogger? _logger;
     private readonly TimeProvider _timeProvider;
     // micro-optimization: Using nuint as the value type to avoid GC write barriers; could replace with ConcurrentHashSet if that becomes available
     private readonly ConcurrentDictionary<PinnedBlockMemoryPool, nuint> _pools = new();
 
-    public PinnedBlockMemoryPoolFactory(IMeterFactory meterFactory, TimeProvider? timeProvider = null)
+    public PinnedBlockMemoryPoolFactory(IMeterFactory meterFactory, TimeProvider? timeProvider = null, ILogger<PinnedBlockMemoryPoolFactory>? logger = null)
     {
         _timeProvider = timeProvider ?? TimeProvider.System;
         _meterFactory = meterFactory;
+        _logger = logger;
     }
 
     public MemoryPool<byte> Create()
     {
-        var pool = new PinnedBlockMemoryPool(_meterFactory);
+        var pool = new PinnedBlockMemoryPool(_meterFactory, _logger);
 
         _pools.TryAdd(pool, nuint.Zero);
 

src/Servers/Kestrel/Core/test/PinnedBlockMemoryPoolFactoryTests.cs

@@ -94,13 +94,19 @@ public async Task FactoryHeartbeatWorks()
 
         static async Task VerifyPoolEviction(PinnedBlockMemoryPool pool, int previousCount)
         {
+            // Because the eviction happens on a thread pool thread, we need to wait for it to complete
+            // and the only way to do that (without adding a test hook in the pool code) is to delay.
+            // But we don't want to add an arbitrary delay, so we do a short delay with block count checks
+            // to reduce the wait time.
             var maxWait = TimeSpan.FromSeconds(5);
             while (pool.BlockCount() > previousCount - (previousCount / 30) && maxWait > TimeSpan.Zero)
             {
                 await Task.Delay(50);
                 maxWait -= TimeSpan.FromMilliseconds(50);
             }
 
+            // Assert that the block count has decreased by 3.3-10%.
+            // This relies on the current implementation of eviction logic which may change in the future.
             Assert.InRange(pool.BlockCount(), previousCount - (previousCount / 10), previousCount - (previousCount / 30));
         }
     }

src/Shared/Buffers.MemoryPool/DefaultMemoryPoolFactory.cs

@@ -18,11 +18,13 @@ internal sealed class DefaultMemoryPoolFactory : IMemoryPoolFactory<byte>, IAsyn
     private readonly ConcurrentDictionary<PinnedBlockMemoryPool, bool> _pools = new();
     private readonly PeriodicTimer _timer;
     private readonly Task _timerTask;
+    private readonly ILogger? _logger;
 
     public DefaultMemoryPoolFactory(IMeterFactory? meterFactory = null, ILogger<DefaultMemoryPoolFactory>? logger = null)
     {
         _meterFactory = meterFactory;
-        _timer = new PeriodicTimer(TimeSpan.FromSeconds(10));
+        _logger = logger;
+        _timer = new PeriodicTimer(PinnedBlockMemoryPool.DefaultEvictionDelay);
         _timerTask = Task.Run(async () =>
         {
             try
@@ -37,14 +39,14 @@ public DefaultMemoryPoolFactory(IMeterFactory? meterFactory = null, ILogger<Defa
             }
             catch (Exception ex)
             {
-                logger?.LogCritical(ex, "Error while evicting memory from pools.");
+                _logger?.LogCritical(ex, "Error while evicting memory from pools.");
             }
         });
     }
 
     public MemoryPool<byte> Create()
     {
-        var pool = new PinnedBlockMemoryPool(_meterFactory);
+        var pool = new PinnedBlockMemoryPool(_meterFactory, _logger);
 
         _pools.TryAdd(pool, true);
 

src/Shared/Buffers.MemoryPool/PinnedBlockMemoryPool.cs

@@ -4,6 +4,7 @@
 using System.Buffers;
 using System.Collections.Concurrent;
 using System.Diagnostics.Metrics;
+using Microsoft.Extensions.Logging;
 
 #nullable enable
 
@@ -19,6 +20,12 @@ internal sealed class PinnedBlockMemoryPool : MemoryPool<byte>, IThreadPoolWorkI
     /// </summary>
     private const int _blockSize = 4096;
 
+    // 10 seconds chosen arbitrarily. Trying to avoid running eviction too frequently
+    // to avoid trashing if the server gets batches of requests, but want to run often
+    // enough to avoid memory bloat if the server is idle for a while.
+    // This can be tuned later if needed.
+    public static readonly TimeSpan DefaultEvictionDelay = TimeSpan.FromSeconds(10);
+
     /// <summary>
     /// Max allocation block size for pooled blocks,
     /// larger values can be leased but they will be disposed after use rather than returned to the pool.
@@ -42,10 +49,11 @@ internal sealed class PinnedBlockMemoryPool : MemoryPool<byte>, IThreadPoolWorkI
     private bool _isDisposed; // To detect redundant calls
 
     private readonly PinnedBlockMemoryPoolMetrics? _metrics;
+    private readonly ILogger? _logger;
 
     private long _currentMemory;
     private long _evictedMemory;
-    private DateTimeOffset _nextEviction = DateTime.UtcNow.AddSeconds(10);
+    private DateTimeOffset _nextEviction = DateTime.UtcNow.Add(DefaultEvictionDelay);
 
     private uint _rentCount;
     private uint _returnCount;
@@ -60,9 +68,10 @@ internal sealed class PinnedBlockMemoryPool : MemoryPool<byte>, IThreadPoolWorkI
     /// </summary>
     private const int AnySize = -1;
 
-    public PinnedBlockMemoryPool(IMeterFactory? meterFactory = null)
+    public PinnedBlockMemoryPool(IMeterFactory? meterFactory = null, ILogger? logger = null)
     {
         _metrics = meterFactory is null ? null : new PinnedBlockMemoryPoolMetrics(meterFactory);
+        _logger = logger;
     }
 
     /// <summary>
@@ -140,7 +149,7 @@ public bool TryScheduleEviction(DateTimeOffset now)
     {
         if (now >= _nextEviction)
         {
-            _nextEviction = now.AddSeconds(10);
+            _nextEviction = now.Add(DefaultEvictionDelay);
             ThreadPool.UnsafeQueueUserWorkItem(this, preferLocal: false);
             return true;
         }
@@ -150,7 +159,15 @@ public bool TryScheduleEviction(DateTimeOffset now)
 
     void IThreadPoolWorkItem.Execute()
     {
-        PerformEviction();
+        try
+        {
+            PerformEviction();
+        }
+        catch (Exception ex)
+        {
+            // Log the exception, but don't let it crash the thread pool
+            _logger?.LogError(ex, "Error while performing eviction in PinnedBlockMemoryPool.");
+        }
     }
 
     /// <summary>
@@ -214,6 +231,11 @@ protected override void Dispose(bool disposing)
 
         lock (_disposeSync)
         {
+            if (_isDisposed)
+            {
+                return;
+            }
+
             _isDisposed = true;
 
             _onPoolDisposed?.Invoke(_onPoolDisposedState, this);