Header pooling

src/NServiceBus.AcceptanceTesting/AcceptanceTestingPersistence/Outbox/AcceptanceTestingOutboxStorage.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -17,7 +18,11 @@ public Task<OutboxMessage> Get(string messageId, ContextBag context, Cancellatio
             return NoOutboxMessageTask!;
         }
 
-        return Task.FromResult(new OutboxMessage(messageId, storedMessage.TransportOperations));
+        // Return copies so callers (the dispatch pipeline) own their header dictionaries. The dispatch
+        // pipeline pools and clears outgoing header dictionaries after dispatch, so the outbox must hand out
+        // fresh dictionaries on every Get. Sharing the stored references would let dispatch clear/pool the
+        // stored dictionaries, corrupting storage and leaking aliased dictionaries into the header pool.
+        return Task.FromResult(new OutboxMessage(messageId, storedMessage.TransportOperations.Select(CopyOperation).ToArray()));
     }
 
     public Task<IOutboxTransaction> BeginTransaction(ContextBag context, CancellationToken cancellationToken = default) => Task.FromResult<IOutboxTransaction>(new AcceptanceTestingOutboxTransaction());
@@ -27,7 +32,7 @@ public Task Store(OutboxMessage message, IOutboxTransaction transaction, Context
         var tx = (AcceptanceTestingOutboxTransaction)transaction;
         tx.Enlist(() =>
         {
-            if (!storage.TryAdd(message.MessageId, new StoredMessage(message.MessageId, message.TransportOperations.Select(o => o.DeepCopy()).ToArray())))
+            if (!storage.TryAdd(message.MessageId, new StoredMessage(message.MessageId, message.TransportOperations.Select(CopyOperation).ToArray())))
             {
                 throw new Exception($"Outbox message with id '{message.MessageId}' is already present in storage.");
             }
@@ -58,6 +63,25 @@ public void RemoveEntriesOlderThan(DateTime dateTime)
         }
     }
 
+    // Copies headers/options/body into fresh instances so stored data is independent of the live pipeline
+    // dictionaries (which are pooled and cleared after dispatch). Mirrors NonDurableOutboxStorage.CopyOperation.
+    static TransportOperation CopyOperation(TransportOperation operation)
+    {
+        var headers = operation.Headers != null
+            ? new Dictionary<string, string>(operation.Headers)
+            : [];
+
+        var options = operation.Options != null
+            ? new NServiceBus.Transport.DispatchProperties(operation.Options)
+            : [];
+
+        var body = operation.Body.IsEmpty
+            ? []
+            : operation.Body.ToArray();
+
+        return new TransportOperation(operation.MessageId, options, body, headers);
+    }
+
     readonly ConcurrentDictionary<string, StoredMessage> storage = new();
     static readonly Task<OutboxMessage?> NoOutboxMessageTask = Task.FromResult(default(OutboxMessage));
 

src/NServiceBus.Core.Tests/ApprovalFiles/APIApprovals.ApproveNServiceBus.approved.txt

@@ -2484,6 +2484,11 @@ namespace NServiceBus.Transport
         Handled = 0,
         RetryRequired = 1,
     }
+    public class HeaderPool : NServiceBus.Utils.DictionaryPool<string, string>
+    {
+        public HeaderPool(int maxPoolSize = -1, int maxRetainedCapacityPerItem = 64) { }
+        public new static NServiceBus.Transport.HeaderPool Shared { get; }
+    }
     public class HostSettings
     {
         public HostSettings(string name, string hostDisplayName, NServiceBus.StartupDiagnosticEntries startupDiagnostic, System.Action<string, System.Exception, System.Threading.CancellationToken> criticalErrorAction, bool setupInfrastructure, NServiceBus.Settings.IReadOnlySettings? coreSettings = null) { }
@@ -2788,4 +2793,17 @@ namespace NServiceBus.Utils
         public static System.Guid Create([System.Runtime.CompilerServices.ParamCollection] [System.Runtime.CompilerServices.ScopedRef] System.ReadOnlySpan<string> values) { }
         public static System.Guid Create(string value) { }
     }
+    public class DictionaryPool<TKey, TValue>
+        where TKey :  notnull
+    {
+        public DictionaryPool(int maxPoolSize = -1, int maxRetainedCapacityPerItem = 1024) { }
+        public static NServiceBus.Utils.DictionaryPool<TKey, TValue> Shared { get; }
+        public System.Collections.Generic.Dictionary<TKey, TValue> Rent(int minimumCapacity = 0) { }
+        public void Return(System.Collections.Generic.Dictionary<TKey, TValue> dictionary, bool clearDictionary = true) { }
+    }
+    public static class ReadOnlyDictionaryExtensions
+    {
+        public static void CopyTo<TKey, TValue>(this System.Collections.Generic.IReadOnlyDictionary<TKey, TValue> source, System.Collections.Generic.Dictionary<TKey, TValue> destination)
+            where TKey :  notnull { }
+    }
 }

src/NServiceBus.Core.Tests/Transports/Learning/HeaderSerializerTests.cs

@@ -1,11 +1,20 @@
-namespace NServiceBus.Core.Tests.Transports.Learning;
+#nullable enable
+
+namespace NServiceBus.Core.Tests.Transports.Learning;
 
 using System.Collections.Generic;
+using System.Text;
+using System.Text.Json;
+using System.Text.Json.Serialization.Metadata;
 using NUnit.Framework;
 using Particular.Approvals;
+using NServiceBus.Utils;
 
 public class HeaderSerializerTests
 {
+    static readonly JsonTypeInfo<Dictionary<string, string>> SourceGenTypeInfo =
+        HeaderSerializer.HeaderSerializationContext.Default.DictionaryStringString;
+
     [Test]
     public void Can_round_trip_headers()
     {
@@ -26,4 +35,175 @@ public void Can_round_trip_headers()
         var deserialize = HeaderSerializer.Deserialize(serialized);
         Assert.That(deserialize, Is.EquivalentTo(input));
     }
+
+    [Test]
+    public void Can_deserialize_from_utf8_bytes()
+    {
+        var input = new Dictionary<string, string>
+        {
+            { "key 1", "value 1" },
+            { "key 2", "value 2" }
+        };
+        var serialized = HeaderSerializer.Serialize(input);
+        var bytes = Encoding.UTF8.GetBytes(serialized);
+
+        var deserialized = HeaderSerializer.Deserialize(bytes);
+        Assert.That(deserialized, Is.EquivalentTo(input));
+    }
+
+    [Test]
+    public void Deserialize_from_bytes_with_pool_reuses_dictionary()
+    {
+        var pool = new DictionaryPool<string, string>(maxPoolSize: 4);
+        var input = new Dictionary<string, string>
+        {
+            { "key 1", "value 1" },
+            { "key 2", "value 2" }
+        };
+        var serialized = HeaderSerializer.Serialize(input);
+        var bytes = Encoding.UTF8.GetBytes(serialized);
+
+        var first = HeaderSerializer.Deserialize(bytes, pool);
+        pool.Return(first);
+
+        var second = HeaderSerializer.Deserialize(bytes, pool);
+        Assert.That(second, Is.SameAs(first));
+        Assert.That(second, Is.EquivalentTo(input));
+    }
+
+    [Test]
+    public void Round_trip_empty_dictionary()
+    {
+        var input = new Dictionary<string, string>();
+        var serialized = HeaderSerializer.Serialize(input);
+        var deserialized = HeaderSerializer.Deserialize(serialized);
+        Assert.That(deserialized, Is.Empty);
+    }
+
+    [Test]
+    public void Round_trip_special_characters_in_values()
+    {
+        var input = new Dictionary<string, string>
+        {
+            { "quotes", "value with \"quoted\" text" },
+            { "backslash", "C:\\Program Files\\" },
+            { "newline", "line1\nline2\r\nline3" },
+            { "tab", "a\tb" },
+            { "unicode", "\u00e9\u00e8\u00ea" },
+            { "emoji", "\ud83d\ude00" }
+        };
+        var serialized = HeaderSerializer.Serialize(input);
+        var deserialized = HeaderSerializer.Deserialize(serialized);
+        Assert.That(deserialized, Is.EquivalentTo(input));
+    }
+
+    [Test]
+    public void Round_trip_unicode_keys_and_values()
+    {
+        var input = new Dictionary<string, string>
+        {
+            { "\u00fcnicode-key", "value-with-\u00fcmlaut" },
+            { "\u4e2d\u6587", "\u65e5\u672c\u8a9e" },
+            { "key \u00e0", "value \u00e9" }
+        };
+        var serialized = HeaderSerializer.Serialize(input);
+        var deserialized = HeaderSerializer.Deserialize(serialized);
+        Assert.That(deserialized, Is.EquivalentTo(input));
+    }
+
+    [Test]
+    public void Manual_parser_matches_source_generated_deserializer_for_serialized_output()
+    {
+        // The old Deserialize used JsonSerializer.Deserialize with the source-generated
+        // context (AOT-friendly). The new code uses a manual Utf8JsonReader. Verify
+        // they agree on every input that Serialize can produce.
+        Dictionary<string, string>[] testCases =
+        [
+            new() { { "a", "b" } },
+            new()
+            {
+                { "key 1", "value 1" },
+                { "key 2", "value 2" }
+            },
+            new()
+            {
+                { "special", "\"quotes\" and \\backslashes" },
+                { "multiline", "line1\nline2" }
+            }
+        ];
+
+        foreach (var input in testCases)
+        {
+            var serialized = HeaderSerializer.Serialize(input);
+            var bytes = Encoding.UTF8.GetBytes(serialized);
+
+            // Old path: source-generated AOT-friendly byte-span overload
+            var oldResult = JsonSerializer.Deserialize(bytes, SourceGenTypeInfo);
+
+            // New path: manual Utf8JsonReader parser
+            var newResult = HeaderSerializer.Deserialize(serialized);
+
+            Assert.That(newResult, Is.EquivalentTo(oldResult!),
+                $"Mismatch for input with {input.Count} entries. Serialized: {serialized}");
+        }
+    }
+
+    [Test]
+    public void Deserialize_handles_leading_and_trailing_whitespace()
+    {
+        var input = new Dictionary<string, string> { { "key", "value" } };
+        var serialized = HeaderSerializer.Serialize(input);
+        var padded = "  \n  " + serialized + "  \n  ";
+
+        var deserialized = HeaderSerializer.Deserialize(padded);
+        Assert.That(deserialized, Is.EquivalentTo(input));
+    }
+
+    [TestCase("")]
+    [TestCase("not json")]
+    [TestCase("[]")]
+    [TestCase("123")]
+    [TestCase("\"string\"")]
+    [TestCase("{ \"key\": 123 }")]
+    [TestCase("{ \"key\": [1, 2] }")]
+    [TestCase("{ \"key\": { \"nested\": true } }")]
+    [TestCase("{ \"key\": }")]
+    [TestCase("{ \"key\": \"value\"")]
+    public void Deserialize_throws_for_malformed_json(string malformed) => Assert.That(() => HeaderSerializer.Deserialize(malformed), Throws.InstanceOf<JsonException>());
+
+    [Test]
+    public void Deserialize_with_pool_returns_dict_that_can_be_returned_to_pool()
+    {
+        var pool = new DictionaryPool<string, string>(maxPoolSize: 4);
+        var input = new Dictionary<string, string>
+        {
+            { "key 1", "value 1" },
+            { "key 2", "value 2" }
+        };
+        var serialized = HeaderSerializer.Serialize(input);
+        var bytes = Encoding.UTF8.GetBytes(serialized);
+
+        var dict = HeaderSerializer.Deserialize(bytes, pool);
+        Assert.That(dict, Has.Count.EqualTo(2));
+
+        // Returning and re-renting should give back the same (cleared) dictionary.
+        pool.Return(dict);
+        var reused = pool.Rent();
+        Assert.That(reused, Is.SameAs(dict));
+        Assert.That(reused, Is.Empty);
+    }
+
+    [Test]
+    public void Deserialize_from_bytes_with_pool_returns_dict_on_parse_failure()
+    {
+        var pool = new DictionaryPool<string, string>(maxPoolSize: 4);
+        var badJson = "not valid json"u8.ToArray();
+
+        Assert.That(() => HeaderSerializer.Deserialize(badJson, pool), Throws.InstanceOf<JsonException>());
+
+        // The failed-parse dictionary was returned to the pool (cleared).
+        // A subsequent rent should succeed normally.
+        var dict = pool.Rent();
+        Assert.That(dict, Is.Empty);
+    }
 }