diff --git a/src/NServiceBus.AcceptanceTesting/AcceptanceTestingPersistence/Outbox/AcceptanceTestingOutboxStorage.cs b/src/NServiceBus.AcceptanceTesting/AcceptanceTestingPersistence/Outbox/AcceptanceTestingOutboxStorage.cs index 5218574ad88..d50b41fd8fb 100644 --- a/src/NServiceBus.AcceptanceTesting/AcceptanceTestingPersistence/Outbox/AcceptanceTestingOutboxStorage.cs +++ b/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 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 BeginTransaction(ContextBag context, CancellationToken cancellationToken = default) => Task.FromResult(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(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 storage = new(); static readonly Task NoOutboxMessageTask = Task.FromResult(default(OutboxMessage)); diff --git a/src/NServiceBus.Core.Tests/ApprovalFiles/APIApprovals.ApproveNServiceBus.approved.txt b/src/NServiceBus.Core.Tests/ApprovalFiles/APIApprovals.ApproveNServiceBus.approved.txt index 6a55b5f799b..e5e367c3025 100644 --- a/src/NServiceBus.Core.Tests/ApprovalFiles/APIApprovals.ApproveNServiceBus.approved.txt +++ b/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 + { + 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 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 values) { } public static System.Guid Create(string value) { } } + public class DictionaryPool + where TKey : notnull + { + public DictionaryPool(int maxPoolSize = -1, int maxRetainedCapacityPerItem = 1024) { } + public static NServiceBus.Utils.DictionaryPool Shared { get; } + public System.Collections.Generic.Dictionary Rent(int minimumCapacity = 0) { } + public void Return(System.Collections.Generic.Dictionary dictionary, bool clearDictionary = true) { } + } + public static class ReadOnlyDictionaryExtensions + { + public static void CopyTo(this System.Collections.Generic.IReadOnlyDictionary source, System.Collections.Generic.Dictionary destination) + where TKey : notnull { } + } } \ No newline at end of file diff --git a/src/NServiceBus.Core.Tests/Transports/Learning/HeaderSerializerTests.cs b/src/NServiceBus.Core.Tests/Transports/Learning/HeaderSerializerTests.cs index 1ac75b358c1..d3c68b5979a 100644 --- a/src/NServiceBus.Core.Tests/Transports/Learning/HeaderSerializerTests.cs +++ b/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> 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 + { + { "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(maxPoolSize: 4); + var input = new Dictionary + { + { "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(); + 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 + { + { "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 + { + { "\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[] 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 { { "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()); + + [Test] + public void Deserialize_with_pool_returns_dict_that_can_be_returned_to_pool() + { + var pool = new DictionaryPool(maxPoolSize: 4); + var input = new Dictionary + { + { "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(maxPoolSize: 4); + var badJson = "not valid json"u8.ToArray(); + + Assert.That(() => HeaderSerializer.Deserialize(badJson, pool), Throws.InstanceOf()); + + // 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); + } } \ No newline at end of file diff --git a/src/NServiceBus.Core.Tests/Utils/DictionaryPoolTests.cs b/src/NServiceBus.Core.Tests/Utils/DictionaryPoolTests.cs new file mode 100644 index 00000000000..db579bd32b8 --- /dev/null +++ b/src/NServiceBus.Core.Tests/Utils/DictionaryPoolTests.cs @@ -0,0 +1,202 @@ +#nullable enable + +namespace NServiceBus.Core.Tests.Utils; + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using NUnit.Framework; +using NServiceBus.Utils; + +public class DictionaryPoolTests +{ + [Test] + public void Rent_returns_empty_dictionary() + { + var pool = new DictionaryPool(maxPoolSize: 4); + var dict = pool.Rent(); + Assert.That(dict, Is.Not.Null); + Assert.That(dict.Count, Is.EqualTo(0)); + } + + [Test] + public void Returned_dictionary_is_reused_on_next_rent() + { + var pool = new DictionaryPool(maxPoolSize: 4); + var dict = pool.Rent(); + dict["a"] = "1"; + pool.Return(dict); + + var reused = pool.Rent(); + Assert.That(reused, Is.SameAs(dict)); + Assert.That(reused.Count, Is.EqualTo(0)); + } + + [Test] + public void Rent_with_minimum_capacity_prevents_resize() + { + var pool = new DictionaryPool(maxPoolSize: 4); + var dict = pool.Rent(minimumCapacity: 100); + Assert.That(dict.Count, Is.EqualTo(0)); + + for (int i = 0; i < 100; i++) + { + dict[$"key{i}"] = $"value{i}"; + } + Assert.That(dict.Count, Is.EqualTo(100)); + } + + [Test] + public void Return_preserves_capacity_for_reuse_without_resize() + { + var pool = new DictionaryPool(maxPoolSize: 4); + + // Fill a dictionary with a realistic header count so its internal + // Entry[]/buckets grow to accommodate ~50 entries. + var first = pool.Rent(); + for (int i = 0; i < 50; i++) + { + first[$"header-{i}"] = $"value-{i}"; + } + var capacityAfterFill = first.Capacity; + Assert.That(capacityAfterFill, Is.GreaterThanOrEqualTo(50), + "Sanity: filling 50 entries should grow capacity to at least 50."); + + // Return → Clear() preserves the internal arrays (Capacity). + pool.Return(first); + + // Rent again — should get the same instance with its capacity intact. + var reused = pool.Rent(); + Assert.That(reused, Is.SameAs(first)); + Assert.That(reused.Count, Is.EqualTo(0), "Cleared dictionary must be empty."); + Assert.That(reused.Capacity, Is.EqualTo(capacityAfterFill), + "Clear() must preserve Capacity so the next rent avoids resize."); + + // Refill the same number of entries — capacity must not change, + // proving no internal reallocation occurred. + for (int i = 0; i < 50; i++) + { + reused[$"header-{i}"] = $"value-{i}"; + } + Assert.That(reused.Capacity, Is.EqualTo(capacityAfterFill), + "Refilling with the same count must not trigger a resize."); + } + + [Test] + public void Oversized_return_trims_capacity() + { + var pool = new DictionaryPool(maxPoolSize: 4, maxRetainedCapacityPerItem: 5); + var dict = pool.Rent(); + for (int i = 0; i < 100; i++) + { + dict[$"key{i}"] = $"value{i}"; + } + pool.Return(dict); // Count=100 > 5 → Clear + TrimExcess + + var reused = pool.Rent(); + reused["x"] = "1"; + Assert.That(reused.Count, Is.EqualTo(1)); + } + + [Test] + public void Soft_cap_drops_excess_returns_beyond_capacity() + { + var pool = new DictionaryPool(maxPoolSize: 2); + + // Rent three distinct dictionaries without returning any, + // so the pool is empty and each Rent allocates fresh. + var d1 = pool.Rent(); + var d2 = pool.Rent(); + var d3 = pool.Rent(); + + // Return all three. The first two bring the pool to its cap (2); + // the third must be dropped, not retained. + pool.Return(d1); // count → 1 + pool.Return(d2); // count → 2 (cap) + pool.Return(d3); // exceeds cap → dropped, count stays 2 + + // Rent twice — we should get d1 and d2 back (in LIFO order via ConcurrentBag). + var r1 = pool.Rent(); + var r2 = pool.Rent(); + + Assert.That(r1, Is.SameAs(d2).Or.SameAs(d1), "First rent should return a pooled dictionary."); + Assert.That(r2, Is.SameAs(d2).Or.SameAs(d1), "Second rent should return a pooled dictionary."); + Assert.That(r1, Is.Not.SameAs(r2), "The two rents should return different dictionaries."); + + // The third rent must NOT return d3 — it was dropped by the cap. + var r3 = pool.Rent(); + Assert.That(r3, Is.Not.SameAs(d3), "Dropped dictionary must not be returned from the pool."); + Assert.That(r3.Count, Is.EqualTo(0), "Fresh rent must be empty."); + } + + [Test] + public void Return_null_throws_argument_null_exception() + { + var pool = new DictionaryPool(maxPoolSize: 4); + Assert.Throws(() => pool.Return(null!)); + } + + [Test] + public void Return_without_clear_preserves_data_for_caller_that_cleared() + { + var pool = new DictionaryPool(maxPoolSize: 4); + var dict = pool.Rent(); + dict["a"] = "1"; + + // Caller is responsible for clearing when passing clearDictionary: false. + dict.Clear(); + pool.Return(dict, clearDictionary: false); + + var reused = pool.Rent(); + Assert.That(reused, Is.SameAs(dict)); + Assert.That(reused.Count, Is.EqualTo(0)); + } + + [Test] + public void Concurrent_rent_return_does_not_corrupt() + { + var pool = new DictionaryPool(maxPoolSize: 64); + var iterations = 5000; + var failures = new System.Collections.Concurrent.ConcurrentQueue(); + + Parallel.For(0, iterations, i => + { + var dict = pool.Rent(); + + // Every rented dictionary must start empty — a stale entry + // means Clear was skipped or a cross-thread race leaked data. + if (dict.Count != 0) + { + failures.Enqueue($"Iteration {i}: rented dict had Count={dict.Count} (expected 0)"); + return; + } + + // Write a unique key per iteration so collisions are detectable. + dict[$"key-{i}"] = $"value-{i}"; + + pool.Return(dict); + }); + + Assert.That(failures, Is.Empty, $"{failures.Count} rent-return cycles saw stale data: {string.Join("; ", failures.Take(5))}"); + + // The pool's approximate count must never exceed the cap. + Assert.That(pool.Count, Is.LessThanOrEqualTo(64), + $"Pool count {pool.Count} exceeded maxPoolSize of 64."); + + // Drain the pool by renting until it's empty. Every returned dictionary + // must be empty (Return must have cleared it). + var drained = 0; + while (pool.Count > 0) + { + var dict = pool.Rent(); + drained++; + if (dict.Count != 0) + { + Assert.Fail($"Drained dictionary #{drained} had Count={dict.Count} (expected 0 after Return)."); + } + } + + Assert.That(drained, Is.LessThanOrEqualTo(64), $"Pool retained {drained} dictionaries, exceeding maxPoolSize of 64."); + } +} \ No newline at end of file diff --git a/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs b/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs index a7ec2fc198b..a5ed0988821 100644 --- a/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs +++ b/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs @@ -3,10 +3,10 @@ namespace NServiceBus; using System; -using System.Collections.Generic; using System.Threading.Tasks; using Pipeline; using Transport; +using NServiceBus.Utils; class InvokeAuditPipelineBehavior : IForkConnector { @@ -22,7 +22,9 @@ public async Task Invoke(IIncomingPhysicalMessageContext context, Func(context.Message.Headers), context.Message.Body); + var auditHeaders = HeaderPool.Shared.Rent(context.Message.Headers.Count); + context.Message.Headers.CopyTo(auditHeaders); + var processedMessage = new OutgoingMessage(context.Message.MessageId, auditHeaders, context.Message.Body); var auditContext = this.CreateAuditContext(processedMessage, auditAddress, timeToBeReceived, context); diff --git a/src/NServiceBus.Core/Pipeline/Outgoing/ImmediateDispatchTerminator.cs b/src/NServiceBus.Core/Pipeline/Outgoing/ImmediateDispatchTerminator.cs index c6ac3a760bd..fe469715640 100644 --- a/src/NServiceBus.Core/Pipeline/Outgoing/ImmediateDispatchTerminator.cs +++ b/src/NServiceBus.Core/Pipeline/Outgoing/ImmediateDispatchTerminator.cs @@ -7,19 +7,22 @@ namespace NServiceBus; using Pipeline; using Transport; -class ImmediateDispatchTerminator : PipelineTerminator +class ImmediateDispatchTerminator(IMessageDispatcher dispatcher) : PipelineTerminator { - public ImmediateDispatchTerminator(IMessageDispatcher dispatcher) - { - this.dispatcher = dispatcher; - } - - protected override Task Terminate(IDispatchContext context) + protected override async Task Terminate(IDispatchContext context) { var transaction = context.Extensions.GetOrCreate(); var operations = context.Operations as TransportOperation[] ?? context.Operations.ToArray(); - return dispatcher.Dispatch(new TransportOperations(operations), transaction, context.CancellationToken); + try + { + await dispatcher.Dispatch(new TransportOperations(operations), transaction, context.CancellationToken).ConfigureAwait(false); + } + finally + { + foreach (var operation in operations) + { + HeaderPool.Shared.Return(operation.Message.Headers); + } + } } - - readonly IMessageDispatcher dispatcher; } \ No newline at end of file diff --git a/src/NServiceBus.Core/Pipeline/Outgoing/RoutingContextExtensions.cs b/src/NServiceBus.Core/Pipeline/Outgoing/RoutingContextExtensions.cs index 21e5d67d522..f88eeb624d6 100644 --- a/src/NServiceBus.Core/Pipeline/Outgoing/RoutingContextExtensions.cs +++ b/src/NServiceBus.Core/Pipeline/Outgoing/RoutingContextExtensions.cs @@ -6,12 +6,22 @@ namespace NServiceBus; using Pipeline; using Routing; using Transport; +using NServiceBus.Utils; static class RoutingContextExtensions { public static TransportOperation ToTransportOperation(this IRoutingContext context, RoutingStrategy strategy, DispatchConsistency dispatchConsistency, bool copySharedMutableMessageState) { - var headers = copySharedMutableMessageState ? new Dictionary(context.Message.Headers) : context.Message.Headers; + Dictionary headers; + if (copySharedMutableMessageState) + { + headers = HeaderPool.Shared.Rent(context.Message.Headers.Count); + context.Message.Headers.CopyTo(headers); + } + else + { + headers = context.Message.Headers; + } var dispatchProperties = context.Extensions.TryGet(out var properties) ? copySharedMutableMessageState ? new DispatchProperties(properties) : properties : []; diff --git a/src/NServiceBus.Core/Pipeline/Outgoing/RoutingToDispatchConnector.cs b/src/NServiceBus.Core/Pipeline/Outgoing/RoutingToDispatchConnector.cs index bc4be17db8b..048db29bb70 100644 --- a/src/NServiceBus.Core/Pipeline/Outgoing/RoutingToDispatchConnector.cs +++ b/src/NServiceBus.Core/Pipeline/Outgoing/RoutingToDispatchConnector.cs @@ -62,6 +62,14 @@ public override Task Invoke(IRoutingContext context, Func(out var pendingOperations)) { pendingOperations.AddRange(operations); diff --git a/src/NServiceBus.Core/Transports/HeaderPool.cs b/src/NServiceBus.Core/Transports/HeaderPool.cs new file mode 100644 index 00000000000..15e47da87c4 --- /dev/null +++ b/src/NServiceBus.Core/Transports/HeaderPool.cs @@ -0,0 +1,34 @@ +#nullable enable + +namespace NServiceBus.Transport; + +using NServiceBus.Utils; + +/// +/// A pool of Dictionary<string, string> instances specialized for +/// message headers, with defaults tuned for typical header counts. +/// Inherits from . +/// +/// +/// Use for the process-wide instance. The retained-capacity +/// threshold defaults to 64, well above typical header counts, so normal usage +/// always takes the no-realloc path on reuse. +/// +public class HeaderPool : DictionaryPool +{ + /// A shared, process-wide header pool instance. + public static new HeaderPool Shared { get; } = new(); + + /// + /// Soft cap on the number of dictionaries retained. Defaults to a generous + /// multiple of processor count. + /// + /// + /// If a returned dictionary's entry count exceeds this, it is trimmed before + /// being pooled. Defaults to 64, well above typical header counts. + /// + public HeaderPool(int maxPoolSize = -1, int maxRetainedCapacityPerItem = 64) + : base(maxPoolSize, maxRetainedCapacityPerItem) + { + } +} \ No newline at end of file diff --git a/src/NServiceBus.Core/Transports/Learning/HeaderSerializer.cs b/src/NServiceBus.Core/Transports/Learning/HeaderSerializer.cs index 5b04c4913ce..30aa56cbdbe 100644 --- a/src/NServiceBus.Core/Transports/Learning/HeaderSerializer.cs +++ b/src/NServiceBus.Core/Transports/Learning/HeaderSerializer.cs @@ -1,18 +1,103 @@ +#nullable enable + namespace NServiceBus; +using System; using System.Collections.Generic; +using System.Text; using System.Text.Json; using System.Text.Json.Serialization; +using NServiceBus.Utils; static partial class HeaderSerializer { public static string Serialize(Dictionary dictionary) => JsonSerializer.Serialize(dictionary, HeaderSerializationContext.Default.DictionaryStringString); + // TODO: Remove since this is only used in tests? public static Dictionary Deserialize(string value) => - JsonSerializer.Deserialize(value, HeaderSerializationContext.Default.DictionaryStringString); + Deserialize(Encoding.UTF8.GetBytes(value), pool: null); + + /// + /// Deserializes header JSON (UTF-8) into a dictionary, renting from + /// when provided to avoid allocating a fresh dictionary. + /// When is null a new dictionary is allocated. + /// + /// + /// Uses a manual parse rather than + /// JsonSerializer.Deserialize<Dictionary<string,string>>, because + /// System.Text.Json always allocates a fresh dictionary and cannot populate a + /// rented/pooled one, so the pool would cut zero GC pressure otherwise. + /// + public static Dictionary Deserialize(ReadOnlySpan utf8Json, DictionaryPool? pool = null) + { + var dict = pool?.Rent() ?? []; + try + { + Populate(utf8Json, dict); + } + catch + { + // Return the rented dictionary (which clears it) so a parse failure + // doesn't waste the pooled instance. + pool?.Return(dict); + throw; + } + + return dict; + } + + static void Populate(ReadOnlySpan utf8Json, Dictionary dict) + { + var reader = new Utf8JsonReader(utf8Json); + + if (!reader.Read()) + { + throw new JsonException("Invalid JSON: no data."); + } + + if (reader.TokenType != JsonTokenType.StartObject) + { + throw new JsonException($"Invalid JSON: expected start of object, got {reader.TokenType}."); + } + + while (reader.Read()) + { + if (reader.TokenType == JsonTokenType.EndObject) + { + return; + } + + if (reader.TokenType != JsonTokenType.PropertyName) + { + throw new JsonException($"Invalid JSON: expected property name, got {reader.TokenType}."); + } + + var key = reader.GetString()!; + + if (!reader.Read()) + { + throw new JsonException("Invalid JSON: expected value."); + } + + if (reader.TokenType == JsonTokenType.String) + { + dict[key] = reader.GetString()!; + } + else if (reader.TokenType == JsonTokenType.Null) + { + dict[key] = null!; + } + else + { + throw new JsonException($"Invalid JSON: expected string value, got {reader.TokenType}."); + } + } + + throw new JsonException("Invalid JSON: unexpected end of input."); + } [JsonSourceGenerationOptions(WriteIndented = true, IndentSize = 2, NewLine = "\r\n")] [JsonSerializable(typeof(Dictionary))] - sealed partial class HeaderSerializationContext : JsonSerializerContext; + internal sealed partial class HeaderSerializationContext : JsonSerializerContext; } \ No newline at end of file diff --git a/src/NServiceBus.Core/Transports/Learning/LearningTransportMessagePump.cs b/src/NServiceBus.Core/Transports/Learning/LearningTransportMessagePump.cs index a6a8068aa24..fb6c514d2fe 100644 --- a/src/NServiceBus.Core/Transports/Learning/LearningTransportMessagePump.cs +++ b/src/NServiceBus.Core/Transports/Learning/LearningTransportMessagePump.cs @@ -11,6 +11,7 @@ using Logging; using Transport; + class LearningTransportMessagePump : IMessageReceiver { public LearningTransportMessagePump(string id, @@ -306,70 +307,50 @@ async Task ProcessFileAndComplete(ILearningTransportTransaction transaction, str async Task ProcessFile(ILearningTransportTransaction transaction, string messageId, CancellationToken messageProcessingCancellationToken) { - var message = await AsyncFile.ReadText(transaction.FileToProcess, messageProcessingCancellationToken).ConfigureAwait(false); + var messageBytes = await AsyncFile.ReadBytes(transaction.FileToProcess, messageProcessingCancellationToken).ConfigureAwait(false); var bodyPath = Path.Combine(bodyDir, $"{messageId}{BodyFileSuffix}"); - var headers = HeaderSerializer.Deserialize(message); + var headers = HeaderSerializer.Deserialize(messageBytes, HeaderPool.Shared); + Dictionary errorHeaders = null; var fileCreatedAt = File.GetCreationTimeUtc(transaction.FileToProcess); - if (headers.Remove(LearningTransportHeaders.TimeToBeReceived, out var ttbrString)) + try { - var ttbr = TimeSpan.Parse(ttbrString); - - var utcNow = DateTime.UtcNow; - - if (fileCreatedAt + ttbr < utcNow) + if (headers.Remove(LearningTransportHeaders.TimeToBeReceived, out var ttbrString)) { - await transaction.Commit(messageProcessingCancellationToken).ConfigureAwait(false); - log.InfoFormat("Dropping message '{0}' as the specified TimeToBeReceived of '{1}' expired since sending the message at '{2:O}'. Current UTC time is '{3:O}'", messageId, ttbrString, fileCreatedAt, utcNow); - return; - } - } - - var body = await AsyncFile.ReadBytes(bodyPath, messageProcessingCancellationToken).ConfigureAwait(false); - - var transportTransaction = new TransportTransaction(); - - if (transactionMode == TransportTransactionMode.SendsAtomicWithReceive) - { - transportTransaction.Set(transaction); - } + var ttbr = TimeSpan.Parse(ttbrString); - var processingContext = new ContextBag(); - var receiveProperties = new ReceiveProperties(new Dictionary - { - ["LearningTransport.FileCreatedAt"] = fileCreatedAt.ToString("O") - }); + var utcNow = DateTime.UtcNow; - var messageContext = new MessageContext(messageId, headers, body, receiveProperties, transportTransaction, ReceiveAddress, processingContext); + if (fileCreatedAt + ttbr < utcNow) + { + await transaction.Commit(messageProcessingCancellationToken).ConfigureAwait(false); + log.InfoFormat("Dropping message '{0}' as the specified TimeToBeReceived of '{1}' expired since sending the message at '{2:O}'. Current UTC time is '{3:O}'", messageId, ttbrString, fileCreatedAt, utcNow); + return; + } + } - try - { - await onMessage(messageContext, messageProcessingCancellationToken).ConfigureAwait(false); - } - catch (Exception ex) when (ex.IsCausedBy(messageProcessingCancellationToken)) - { - log.Debug("Message processing canceled. Rolling back transaction.", ex); - transaction.Rollback(); - throw; - } - catch (Exception exception) - { - transaction.ClearPendingOutgoingOperations(); + var body = await AsyncFile.ReadBytes(bodyPath, messageProcessingCancellationToken).ConfigureAwait(false); - var processingFailures = retryCounts.AddOrUpdate(messageId, id => 1, (id, currentCount) => currentCount + 1); + var transportTransaction = new TransportTransaction(); - headers = HeaderSerializer.Deserialize(message); - headers.Remove(LearningTransportHeaders.TimeToBeReceived); + if (transactionMode == TransportTransactionMode.SendsAtomicWithReceive) + { + transportTransaction.Set(transaction); + } - var errorContext = new ErrorContext(exception, headers, messageId, body, receiveProperties, transportTransaction, processingFailures, ReceiveAddress, processingContext); + var processingContext = new ContextBag(); + var receiveProperties = new ReceiveProperties(new Dictionary + { + ["LearningTransport.FileCreatedAt"] = fileCreatedAt.ToString("O") + }); - ErrorHandleResult result; + var messageContext = new MessageContext(messageId, headers, body, receiveProperties, transportTransaction, ReceiveAddress, processingContext); try { - result = await onError(errorContext, messageProcessingCancellationToken).ConfigureAwait(false); + await onMessage(messageContext, messageProcessingCancellationToken).ConfigureAwait(false); } catch (Exception ex) when (ex.IsCausedBy(messageProcessingCancellationToken)) { @@ -377,20 +358,52 @@ async Task ProcessFile(ILearningTransportTransaction transaction, string message transaction.Rollback(); throw; } - catch (Exception ex) + catch (Exception exception) { - criticalErrorAction($"Failed to execute recoverability policy for message with native ID: `{messageContext.NativeMessageId}`", ex, messageProcessingCancellationToken); - result = ErrorHandleResult.RetryRequired; + transaction.ClearPendingOutgoingOperations(); + + var processingFailures = retryCounts.AddOrUpdate(messageId, id => 1, (id, currentCount) => currentCount + 1); + + errorHeaders = HeaderSerializer.Deserialize(messageBytes, HeaderPool.Shared); + errorHeaders.Remove(LearningTransportHeaders.TimeToBeReceived); + + var errorContext = new ErrorContext(exception, errorHeaders, messageId, body, receiveProperties, transportTransaction, processingFailures, ReceiveAddress, processingContext); + + ErrorHandleResult result; + + try + { + result = await onError(errorContext, messageProcessingCancellationToken).ConfigureAwait(false); + } + catch (Exception ex) when (ex.IsCausedBy(messageProcessingCancellationToken)) + { + log.Debug("Message processing canceled. Rolling back transaction.", ex); + transaction.Rollback(); + throw; + } + catch (Exception ex) + { + criticalErrorAction($"Failed to execute recoverability policy for message with native ID: `{messageContext.NativeMessageId}`", ex, messageProcessingCancellationToken); + result = ErrorHandleResult.RetryRequired; + } + + if (result == ErrorHandleResult.RetryRequired) + { + transaction.Rollback(); + return; + } } - if (result == ErrorHandleResult.RetryRequired) + await transaction.Commit(messageProcessingCancellationToken).ConfigureAwait(false); + } + finally + { + if (errorHeaders is not null) { - transaction.Rollback(); - return; + HeaderPool.Shared.Return(errorHeaders); } + HeaderPool.Shared.Return(headers); } - - await transaction.Commit(messageProcessingCancellationToken).ConfigureAwait(false); } CancellationTokenSource messagePumpCancellationTokenSource; diff --git a/src/NServiceBus.Core/Unicast/IncomingMessageOperations.cs b/src/NServiceBus.Core/Unicast/IncomingMessageOperations.cs index d67d491dbb4..e9cc0ccd58d 100644 --- a/src/NServiceBus.Core/Unicast/IncomingMessageOperations.cs +++ b/src/NServiceBus.Core/Unicast/IncomingMessageOperations.cs @@ -1,18 +1,23 @@ namespace NServiceBus; +using System.Collections.Generic; using System.Threading.Tasks; using Pipeline; using Routing; using Transport; +using NServiceBus.Utils; static class IncomingMessageOperations { public static Task ForwardCurrentMessageTo(IIncomingContext context, string destination) { var messageBeingProcessed = context.Extensions.Get(); + var headers = HeaderPool.Shared.Rent(messageBeingProcessed.Headers.Count); + messageBeingProcessed.Headers.CopyTo(headers); + var outgoingMessage = new OutgoingMessage( messageBeingProcessed.MessageId, - messageBeingProcessed.Headers, + headers, messageBeingProcessed.Body); var routingContext = new RoutingContext(outgoingMessage, new UnicastRoutingStrategy(destination), context); diff --git a/src/NServiceBus.Core/Unicast/MessageOperations.cs b/src/NServiceBus.Core/Unicast/MessageOperations.cs index 99e770a0883..65b3f095031 100644 --- a/src/NServiceBus.Core/Unicast/MessageOperations.cs +++ b/src/NServiceBus.Core/Unicast/MessageOperations.cs @@ -5,6 +5,7 @@ namespace NServiceBus; using System.Diagnostics.CodeAnalysis; using System.Threading.Tasks; using Extensibility; +using Utils; using MessageInterfaces; using Pipeline; using Transport; @@ -52,10 +53,7 @@ public Task Publish(IBehaviorContext context, object message, PublishOptions opt async Task Publish(IBehaviorContext context, Type messageType, object message, PublishOptions options) { var messageId = options.UserDefinedMessageId ?? CombGuid.Generate().ToString(); - var headers = new Dictionary(options.OutgoingHeaders) - { - [Headers.MessageId] = messageId - }; + var headers = RentOutgoingHeaders(options.OutgoingHeaders, messageId); var publishContext = new OutgoingPublishContext( new OutgoingLogicalMessage(messageType, message), @@ -71,10 +69,7 @@ async Task Publish(IBehaviorContext context, Type messageType, object message, P await publishPipeline.Invoke(publishContext, activity).ConfigureAwait(false); } - public Task Subscribe(IBehaviorContext context, Type eventType, SubscribeOptions options) - { - return Subscribe(context, new Type[] { eventType }, options); - } + public Task Subscribe(IBehaviorContext context, Type eventType, SubscribeOptions options) => Subscribe(context, [eventType], options); public async Task Subscribe(IBehaviorContext context, Type[] eventTypes, SubscribeOptions options) { @@ -104,10 +99,7 @@ public async Task Unsubscribe(IBehaviorContext context, Type eventType, Unsubscr await unsubscribePipeline.Invoke(unsubscribeContext, activity).ConfigureAwait(false); } - public Task Send<[DynamicallyAccessedMembers(IMessageCreator.CreatorMembersRequired)] T>(IBehaviorContext context, Action messageConstructor, SendOptions options) - { - return SendMessage(context, typeof(T), messageMapper.CreateInstance(messageConstructor), options); - } + public Task Send<[DynamicallyAccessedMembers(IMessageCreator.CreatorMembersRequired)] T>(IBehaviorContext context, Action messageConstructor, SendOptions options) => SendMessage(context, typeof(T), messageMapper.CreateInstance(messageConstructor), options); public Task Send(IBehaviorContext context, object message, SendOptions options) { @@ -119,10 +111,7 @@ public Task Send(IBehaviorContext context, object message, SendOptions options) async Task SendMessage(IBehaviorContext context, Type messageType, object message, SendOptions options) { var messageId = options.UserDefinedMessageId ?? CombGuid.Generate().ToString(); - var headers = new Dictionary(options.OutgoingHeaders) - { - [Headers.MessageId] = messageId - }; + var headers = RentOutgoingHeaders(options.OutgoingHeaders, messageId); var outgoingContext = new OutgoingSendContext( new OutgoingLogicalMessage(messageType, message), @@ -145,18 +134,12 @@ public Task Reply(IBehaviorContext context, object message, ReplyOptions options return ReplyMessage(context, messageType, message, options); } - public Task Reply<[DynamicallyAccessedMembers(IMessageCreator.CreatorMembersRequired)] T>(IBehaviorContext context, Action messageConstructor, ReplyOptions options) - { - return ReplyMessage(context, typeof(T), messageMapper.CreateInstance(messageConstructor), options); - } + public Task Reply<[DynamicallyAccessedMembers(IMessageCreator.CreatorMembersRequired)] T>(IBehaviorContext context, Action messageConstructor, ReplyOptions options) => ReplyMessage(context, typeof(T), messageMapper.CreateInstance(messageConstructor), options); async Task ReplyMessage(IBehaviorContext context, Type messageType, object message, ReplyOptions options) { var messageId = options.UserDefinedMessageId ?? CombGuid.Generate().ToString(); - var headers = new Dictionary(options.OutgoingHeaders) - { - [Headers.MessageId] = messageId - }; + var headers = RentOutgoingHeaders(options.OutgoingHeaders, messageId); var outgoingContext = new OutgoingReplyContext( new OutgoingLogicalMessage(messageType, message), @@ -172,9 +155,15 @@ async Task ReplyMessage(IBehaviorContext context, Type messageType, object messa await replyPipeline.Invoke(outgoingContext, activity).ConfigureAwait(false); } - static void MergeDispatchProperties(ContextBag context, DispatchProperties dispatchProperties) - { + static void MergeDispatchProperties(ContextBag context, DispatchProperties dispatchProperties) => // we can't add the constraints directly to the SendOptions ContextBag as the options can be reused context.Set(new DispatchProperties(dispatchProperties)); + + static Dictionary RentOutgoingHeaders(Dictionary outgoingHeaders, string messageId) + { + var headers = HeaderPool.Shared.Rent(outgoingHeaders.Count); + outgoingHeaders.CopyTo(headers); + headers[Headers.MessageId] = messageId; + return headers; } } \ No newline at end of file diff --git a/src/NServiceBus.Core/Utils/DictionaryPool.cs b/src/NServiceBus.Core/Utils/DictionaryPool.cs new file mode 100644 index 00000000000..33c74654d09 --- /dev/null +++ b/src/NServiceBus.Core/Utils/DictionaryPool.cs @@ -0,0 +1,140 @@ +#nullable enable + +namespace NServiceBus.Utils; + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Threading; + +/// +/// A thread-safe pool of reusable instances, +/// modeled on 's Rent/Return pattern. +/// +/// +/// +/// Backed by a , a lock-free (CAS) linked stack: +/// (TryPop) and (Push) are wait-free +/// and never take a monitor lock. Header dictionaries are rented and returned across +/// await boundaries (rent on the pump/handler thread, return on a dispatch +/// continuation thread), so there is no thread-local locality to exploit; a single +/// lock-free structure minimizes the unavoidable cross-thread synchronization cost. +/// +/// +/// A soft capacity cap (enforced via an atomic counter) bounds retention: returns +/// beyond the cap are dropped and left to the GC, so the pool never grows unbounded +/// under bursty or pathological load. The cap is just an upper bound on retention; +/// actual retention self-limits to the number of in-flight dictionaries. +/// +/// +/// calls by default, +/// which preserves the internal bucket/entry arrays (Capacity). For workloads where +/// successive usages need roughly the same number of entries (e.g. message headers, +/// which are uniform across a given endpoint), this means the preserved capacity +/// serves the next rent without any internal resize. TrimExcess +/// only fires when a returned dictionary's entry count exceeds +/// maxRetainedCapacityPerItem, a safety valve so one anomalously large usage +/// doesn't permanently inflate the pool's footprint. This is not the common path. +/// +/// +public class DictionaryPool where TKey : notnull +{ + /// A shared, process-wide pool instance, analogous to ArrayPool<T>.Shared. + public static DictionaryPool Shared { get; } = new(); + + readonly ConcurrentStack> stack = []; + readonly int maxPoolSize; + readonly int maxRetainedCapacityPerItem; + int count; // approximate size, maintained via Interlocked + + /// + /// Approximate number of dictionaries currently retained in the pool. + /// Thread-safe and safe to read at any time, but may lag behind the actual + /// contents under concurrent access. Intended for diagnostics and testing. + /// + internal int Count => Interlocked.CompareExchange(ref count, 0, 0); + + /// + /// Soft cap on the number of dictionaries retained. Returns beyond this limit + /// are dropped rather than growing the pool unbounded. Defaults to a generous + /// multiple of processor count so it's never the bottleneck under normal load. + /// + /// + /// If a returned dictionary's entry count exceeds this, it is trimmed + /// (Clear + TrimExcess) before being pooled, so one unusually large + /// usage doesn't permanently inflate the pool's memory footprint. Defaults to + /// 1024, well above typical header counts, so normal usage always takes the + /// no-realloc path. + /// + public DictionaryPool(int maxPoolSize = -1, int maxRetainedCapacityPerItem = 1024) + { + this.maxPoolSize = maxPoolSize > 0 ? maxPoolSize : Math.Max(Environment.ProcessorCount * 4, 64); + this.maxRetainedCapacityPerItem = maxRetainedCapacityPerItem; + } + + /// + /// Rents a dictionary from the pool, or allocates a new one if the pool is + /// currently empty. The returned dictionary is always empty. + /// + /// + /// Optional capacity hint, mirroring ArrayPool<T>.Rent(minimumLength). + /// When provided, the returned dictionary is pre-sized so you can fill it + /// without triggering internal resizes. + /// + public Dictionary Rent(int minimumCapacity = 0) + { + Dictionary item; + if (stack.TryPop(out var taken)) + { + Interlocked.Decrement(ref count); + item = taken; + } + else + { + item = []; + } + + if (minimumCapacity > 0) + { + item.EnsureCapacity(minimumCapacity); + } + + return item; + } + + /// + /// Returns a previously-rented dictionary to the pool for reuse. + /// + /// The dictionary previously obtained from . + /// + /// Whether to clear the dictionary's contents before pooling it. Defaults to true. + /// Only pass false if you've already cleared it yourself; otherwise the next + /// caller sees stale data. + /// + public void Return(Dictionary dictionary, bool clearDictionary = true) + { + ArgumentNullException.ThrowIfNull(dictionary); + + bool tooLarge = dictionary.Count > maxRetainedCapacityPerItem; + + if (clearDictionary || tooLarge) + { + dictionary.Clear(); + } + + if (tooLarge) + { + // Release the oversized backing arrays so one outlier usage + // doesn't permanently inflate the pool's memory footprint. + dictionary.TrimExcess(); + } + + if (Interlocked.Increment(ref count) > maxPoolSize) + { + Interlocked.Decrement(ref count); + return; // pool is full: drop it and let the GC reclaim it + } + + stack.Push(dictionary); + } +} \ No newline at end of file diff --git a/src/NServiceBus.Core/Utils/ReadOnlyDictionaryExtensions.cs b/src/NServiceBus.Core/Utils/ReadOnlyDictionaryExtensions.cs new file mode 100644 index 00000000000..a1e0683b0e8 --- /dev/null +++ b/src/NServiceBus.Core/Utils/ReadOnlyDictionaryExtensions.cs @@ -0,0 +1,23 @@ +#nullable enable + +namespace NServiceBus.Utils; + +using System.Collections.Generic; + +/// +/// Extension methods for dictionary operations. +/// +public static class ReadOnlyDictionaryExtensions +{ + /// + /// Copies all entries from into , + /// overwriting any existing keys. The destination dictionary is not cleared first. + /// + public static void CopyTo(this IReadOnlyDictionary source, Dictionary destination) where TKey : notnull + { + foreach (var kvp in source) + { + destination[kvp.Key] = kvp.Value; + } + } +} \ No newline at end of file