fix(runtime): preserve cancellation semantics

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: ReubenBond

src/Orleans.Core.Abstractions/Runtime/GrainReference.cs

@@ -94,8 +94,7 @@ public void EnsureActivator<TInvokable>(int index)
             => _ = GetOrCreateActivator<TInvokable>(index);
 
         /// <summary>
-        /// Gets an activator from the cached array at the specified index.
-        /// Callers must ensure <see cref="EnsureActivator{TInvokable}(int)"/> has been called first.
+        /// Gets or creates an activator at the specified index.
         /// </summary>
         /// <typeparam name="TInvokable">The invokable type.</typeparam>
         /// <param name="index">The index of the activator in the array.</param>

src/Orleans.Core/Runtime/CallbackData.cs

@@ -14,6 +14,7 @@ internal sealed partial class CallbackData
         private int completed;
         private StatusResponse? lastKnownStatus;
         private ValueStopwatch stopwatch;
+        private CancellationToken _cancellationToken;
         private CancellationTokenRegistration _cancellationTokenRegistration;
 
         public CallbackData(
@@ -40,12 +41,21 @@ public void SubscribeForCancellation(CancellationToken cancellationToken)
                 return;
             }
 
-            cancellationToken.ThrowIfCancellationRequested();
-            _cancellationTokenRegistration = cancellationToken.UnsafeRegister(static arg =>
+            _cancellationToken = cancellationToken;
+            var registration = cancellationToken.UnsafeRegister(static arg =>
             {
                 var callbackData = (CallbackData)arg!;
                 callbackData.OnCancellation();
             }, this);
+
+            if (IsCompleted)
+            {
+                registration.Dispose();
+            }
+            else
+            {
+                _cancellationTokenRegistration = registration;
+            }
         }
 
         private void SignalCancellation()
@@ -110,9 +120,9 @@ private void OnCancellation()
             SignalCancellation();
             shared.Unregister(Message);
             _applicationRequestInstruments.OnAppRequestsEnd((long)stopwatch.Elapsed.TotalMilliseconds);
-            _applicationRequestInstruments.OnAppRequestsTimedOut(GetTargetGrainType());
+            _applicationRequestInstruments.OnAppRequestsCanceled(GetTargetGrainType());
             OrleansCallBackDataEvent.Instance.OnCanceled(Message);
-            context.Complete(Response.FromException(new OperationCanceledException(_cancellationTokenRegistration.Token)));
+            context.Complete(Response.FromException(new OperationCanceledException(_cancellationToken)));
             _cancellationTokenRegistration.Dispose();
         }