perf(messaging): add ref-counted message pooling

[ReubenBond]

Summary

• Adds a thread-local MessagePool and ref-counted message ownership APIs for safe message reuse.
• Moves message creation/deserialization onto pooled Message instances and resets messages when they return to the pool.
• Updates core/runtime send, receive, drop, and callback lifecycle paths to release message references explicitly.
• Updates activation repartitioning to record message addresses instead of retaining Message instances, and adds focused message pool tests.

Validation

• git diff --check
• conflict-marker scan
• dotnet build src\Orleans.Core\Orleans.Core.csproj -m
• dotnet build src\Orleans.Runtime\Orleans.Runtime.csproj -m
• dotnet test test\Orleans.Core.Tests\Orleans.Core.Tests.csproj --filter MessagePool (20 passed)

Dependencies / notes

• Excludes callback pooling/dictionary, invokable pooling, benchmarks, SEDA, and the transport rewrite.
• NonSilo.Tests.csproj was not present on the current base, so message pool tests live under Orleans.Core.Tests.
• Additional end-to-end networking lifecycle/leak stress coverage would still be valuable.

Microsoft Reviewers: Open in CodeFlow

[Copilot]

Pull request overview

This PR introduces ref-counted ownership tracking for Message instances and a thread-local MessagePool to enable safe reuse of messages across the runtime’s send/receive lifecycle, while updating related runtime paths (networking, callbacks, activation repartitioning) to explicitly release message references.

Changes:

• Add MessagePool plus ref-counted Message.Acquire()/Release() ownership APIs (with DEBUG leak tracking and reset-on-return semantics).
• Switch message creation/deserialization paths to use pooled Message instances and update many lifecycle paths to call Release()/ReleaseDropped().
• Update activation repartitioning sampling to record addresses (instead of retaining Message instances) and add focused message pool tests.

Show a summary per file

File	Description
test/Orleans.Placement.Tests/ActivationRepartitioningTests/TestMessageFilter.cs	Updates test filter API to accept grain ids instead of Message.
test/Orleans.Core.Tests/Messaging/MessagePoolTests.cs	Adds unit tests for pooling/ref-count behaviors and DEBUG leak tracking.
src/Orleans.Runtime/Placement/Repartitioning/RepartitionerMessageFilter.cs	Updates filter API to operate on GrainId pairs rather than Message.
src/Orleans.Runtime/Placement/Repartitioning/ActivationRepartitioner.cs	Changes pending buffer type to store recorded message metadata.
src/Orleans.Runtime/Placement/Repartitioning/ActivationRepartitioner.MessageSink.cs	Records message addressing data instead of retaining Message objects.
src/Orleans.Runtime/Networking/SiloConnection.cs	Releases dropped/expired/rejected messages and marks ownership transfers.
src/Orleans.Runtime/Networking/GatewayInboundConnection.cs	Releases dropped/expired/rejected gateway messages.
src/Orleans.Runtime/Messaging/MessageCenter.cs	Releases blocked/expired outgoing messages and adjusts observer invocation ordering.
src/Orleans.Runtime/Messaging/Gateway.cs	Releases messages rejected due to client drop.
src/Orleans.Runtime/Core/InsideRuntimeClient.cs	Releases/marks responses in callback/no-callback/status-response paths.
src/Orleans.Runtime/Core/HostedClient.cs	Releases expired messages at dispatch.
src/Orleans.Runtime/Catalog/StatelessWorkerGrainContext.cs	Releases dropped messages when context creation fails.
src/Orleans.Runtime/Catalog/ActivationData.cs	Releases dropped messages and releases completed requests.
src/Orleans.Core/Runtime/OutsideRuntimeClient.cs	Releases/marks responses in callback/no-callback/status-response paths.
src/Orleans.Core/Runtime/InvokableObjectManager.cs	Releases messages dropped during observer dispatch/invocation/response.
src/Orleans.Core/Runtime/CallbackData.cs	Acquires request message while awaiting completion; releases on completion/timeout/cancel/fail.
src/Orleans.Core/Networking/Connection.cs	Releases send-pipeline references after flush and on send failures.
src/Orleans.Core/Messaging/MessageSerializer.cs	Deserializes into pooled Message instances.
src/Orleans.Core/Messaging/MessagePool.cs	Adds thread-local message pooling and optional DEBUG leak tracking.
src/Orleans.Core/Messaging/MessageFactory.cs	Creates request/response messages from the pool instead of allocating new ones.
src/Orleans.Core/Messaging/Message.cs	Adds ref-count ownership tracking, drop-release helper, and reset-on-return support.
src/Orleans.Core/Messaging/ClientMessageCenter.cs	Releases dropped/rejected client messages.

Copilot's findings

Comments suppressed due to low confidence (1)

src/Orleans.Runtime/Networking/SiloConnection.cs:256

• In the expired-send path, msg.ReleaseDropped(...) can return the message to the pool and reset it. The subsequent msg.IsPing() check/logging can become a use-after-release (and may log incorrect data or race with reuse). Capture whether it’s a ping (and any info needed for logging) before releasing, and don’t touch msg after calling ReleaseDropped.

            if (msg.IsExpired)
            {
                this.MessagingTrace.OnDropExpiredMessage(msg, MessagingInstruments.Phase.Send);
                msg.ReleaseDropped("ExpiredAtSend");

                if (msg.IsPing())
                {
                    LogWarningDroppingExpiredPingMessage(this.Log, msg);
                }

                return false;

• Files reviewed: 22/22 changed files
• Comments generated: 5

[Copilot]

Review details

Suppressed comments (2)

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs:424

• Dispose() no longer faults outstanding callbacks before disposing the timer. If Dispose is called without a prior StopAsync, any pending requests can hang and pooled request Messages held by CallbackData will remain acquired.

        {
            if (this.disposing) return;
            this.disposing = true;

            Utils.SafeExecute(() => this.callbackTimer.Dispose());

            Utils.SafeExecute(() => MessageCenter?.Dispose());

src/Orleans.Core/Messaging/Message.cs:312

• CacheInvalidationHeader can now grow beyond MaxCacheInvalidationHeaderEntries. Since MessageSerializer caps deserialization to MaxCacheInvalidationHeaderEntries, allowing unbounded growth here can cause invalidation entries to be dropped and increases message size.

                lock (_cacheInvalidationHeader)
                {
                    _cacheInvalidationHeader.Add(grainAddressCacheUpdate);
                }

• Files reviewed: 23/23 changed files
• Comments generated: 7
• Review effort level: Lite

[Copilot]

Review details

Suppressed comments (4)

Previously missed (3) — in code that hasn't changed since the last review.

src/Orleans.Runtime/Messaging/MessageCenter.cs:340

• ProcessRequestToInvalidActivation's rejectMessages path no longer adds the invalid address to the message's CacheInvalidationHeader. This makes the behavior inconsistent with ProcessRequestsToInvalidActivation (which does add invalidation entries when rejecting) and risks leaving stale cache entries if this method is used to reject invalid-activation requests again.

            // IMPORTANT: do not do anything on activation context anymore, since this activation is invalid already.
            if (rejectMessages)
            {
                this.RejectMessage(message, Message.RejectionTypes.Transient, exc, failedOperation);
            }

src/Orleans.Runtime/Core/InsideRuntimeClient.cs:67

• This repo registers a catch-all keyed TimeProvider and documents the convention that subsystems should resolve their clock via [FromKeyedServices(TimeProviderNames.X)] (see src/Orleans.Runtime/Hosting/DefaultSiloServices.cs:68-70). Changing this constructor to use the unkeyed TimeProvider prevents overriding the messaging clock independently.

            TimeProvider timeProvider,

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs:75

• This repo registers a catch-all keyed TimeProvider and documents the convention that subsystems should resolve their clock via [FromKeyedServices(TimeProviderNames.X)] (see src/Orleans.Core/Core/DefaultClientServices.cs:53-55). Using the unkeyed TimeProvider here prevents consumers from overriding the messaging clock independently.

            TimeProvider timeProvider,

src/Orleans.Core/Messaging/Message.cs:301

• AddToCacheInvalidationHeader no longer de-duplicates entries by GrainId. This can fill the header with duplicate updates, causing later distinct invalidations to be dropped once MaxCacheInvalidationHeaderEntries is reached (potentially leaving stale cache entries). Consider restoring the previous GrainId de-duplication under the existing lock.

                    lock (_cacheInvalidationHeader)
                    {
                        if (_cacheInvalidationHeader.Count < MaxCacheInvalidationHeaderEntries)
                        {
                            _cacheInvalidationHeader.Add(grainAddressCacheUpdate);

• Files reviewed: 25/25 changed files
• Comments generated: 0 new
• Review effort level: Lite

[ReubenBond]

The retried .NET CI workflow confirms PR-specific regressions across platforms and both TFMs. MessageTests.AddToCacheInvalidationHeader_DeduplicatesByGrainId consistently receives two cache updates instead of one, and SiloDisposalLeakTests.KilledSilo_WithInFlightOutboundCall_DisposesPromptly times out after 90 seconds on Linux, Windows, and macOS. These failures recur on retry and need code fixes rather than another rerun. Run: https://github.com/dotnet/orleans/actions/runs/32131337094

[Copilot]

Review details

Suppressed comments (1)

src/Orleans.Core/Networking/Connection.cs:497

• HandleReceiveMessageFailure can leak pooled Messages when deserialization fails for OneWay messages (HasDirection=true, Direction!=Request/Response) or when headers cannot be decoded enough to set Direction (HasDirection=false). In those cases the method currently returns true without transferring ownership or releasing the message, so it remains checked out from MessagePool.

        /// <summary>
        /// Handles a message receive failure.
        /// </summary>
        /// <returns><see langword="true"/> if the exception should not be caught and <see langword="false"/> if it should be caught.</returns>
        private bool HandleReceiveMessageFailure(Message message, Exception exception)
        {
            LogErrorExceptionReadingMessage(this.Log, exception, message, this.RemoteEndPoint, this.LocalEndPoint);

            // If deserialization completely failed, rethrow the exception so that it can be handled at another level.
            if (message is null || exception is InvalidMessageFrameException)
            {
                // Returning false here informs the caller that the exception should not be caught.
                return false;
            }

            // The message body was not successfully decoded, but the headers were.
            MessagingMetrics.OnRejectedMessage(message);

            if (message.HasDirection)
            {
                if (message.Direction == Message.Directions.Request)
                {
                    // Send a fast fail to the caller.
                    var response = this.MessageFactory.CreateResponseMessage(message);
                    response.Result = Message.ResponseTypes.Error;
                    response.BodyObject = Response.FromException(exception);

                    // Send the error response and continue processing the next message.
                    this.Send(response);
                    message.ReleaseDropped("ReceiveMessageDeserializationFailure");
                }
                else if (message.Direction == Message.Directions.Response)
                {
                    // If the message was a response, propagate the exception to the intended recipient.
                    message.Result = Message.ResponseTypes.Error;
                    message.BodyObject = Response.FromException(exception);
                    this.OnReceivedMessage(message);
                }
            }

            // The exception has been handled by propagating it onwards.
            return true;

• Files reviewed: 25/25 changed files
• Comments generated: 8
• Review effort level: Lite

[Copilot]

Review details

• Files reviewed: 25/25 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

Review details

Suppressed comments (1)

src/Orleans.Core/Networking/Connection.cs:412

• In ProcessOutgoing, if output.FlushAsync() returns IsCompleted/IsCanceled, the loop breaks before releasing messages in the current inflight batch. Since inflight messages hold the send pipeline’s ref-count, this can strand pooled Message instances (refcount never decremented) during connection shutdown/backpressure scenarios.

                    var flushResult = await output.FlushAsync();
                    if (flushResult.IsCompleted || flushResult.IsCanceled)
                    {
                        break;
                    }

                    // Release the send pipeline's reference after bytes have been flushed.
                    foreach (var msg in inflight)
                    {
                        msg.MarkTransferred("Connection.ProcessOutgoing:Sent");
                        msg.Release();
                    }

                    inflight.Clear();

• Files reviewed: 25/25 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

Review details

• Files reviewed: 26/26 changed files
• Comments generated: 1
• Review effort level: Lite

[github-actions]

Code coverage

77.54% line coverage - 97,821 / 126,157 lines

Coverage details

• Source files: 1,661
• Scope: every CI test matrix job, including providers, CodeGen, .NET 8/10, Linux, Windows, and macOS
• Sources: maintained repository code under src/ loaded by the complete CI test scope
• Commit: e0de3dab27
• View the test run and merged coverage artifact
• View the coverage reporting run

[Copilot]

Review details

Suppressed comments (1)

src/Orleans.Core/Messaging/MessagePool.cs:100

• The XML doc for MessagePool.Return says it "Returns a message to the pool after resetting it", but the implementation just calls message.Release(), which only resets/returns the instance when the ref-count reaches 0. This can mislead callers into thinking Return is unconditional. Consider updating the doc to reflect ref-counted semantics (eg "releases the caller's reference"), or renaming/removing Return to avoid suggesting it bypasses ownership tracking.

    /// <summary>
    /// Returns a message to the pool after resetting it.
    /// </summary>
    public static void Return(Message message) => message.Release();

• Files reviewed: 27/27 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

Review details

• Files reviewed: 32/32 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

Review details

• Files reviewed: 33/33 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

Review details

• Files reviewed: 33/33 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

src/Orleans.Runtime/Messaging/MessageCenter.cs:444

• In the forwarding-failure path, message.ReleaseDropped("ForwardingFailed") sets the debug LastTransferTag to Dropped:ForwardingFailed, but the subsequent unconditional message.MarkTransferred("MessageCenter.TryForwardRequest") overwrites that tag before the final Release(). This reduces the usefulness of the drop reason when diagnosing use-after-release/leak assertions.

Consider only calling MarkTransferred("MessageCenter.TryForwardRequest") on the success path, and on the failure path just releasing the remaining reference (eg message.Release() after ReleaseDropped(...)) so the last tag stays Dropped:ForwardingFailed.

                    message.ReleaseDropped("ForwardingFailed");
                }

                message.MarkTransferred("MessageCenter.TryForwardRequest");
                message.Release();
            }

• Files reviewed: 33/33 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

Copilot review overview

Review tier: Lite
Findings: None