Always reset RuntimeContext to previous value after use

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

@@ -23,31 +23,22 @@ internal static class RuntimeContext
         /// Sets the current grain context.
         /// </summary>
         /// <param name="newContext">The new context.</param>
+        /// <param name="currentContext">The current context at the time of the call.</param>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal static void SetExecutionContext(IGrainContext newContext)
+        internal static void SetExecutionContext(IGrainContext newContext, out IGrainContext currentContext)
         {
+            currentContext = _threadLocalContext;
             _threadLocalContext = newContext;
         }
 
         /// <summary>
-        /// Sets the current grain context.
-        /// </summary>
-        /// <param name="newContext">The new context.</param>
-        /// <param name="existingContext">The existing context.</param>
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal static void SetExecutionContext(IGrainContext newContext, out IGrainContext existingContext)
-        {
-            existingContext = _threadLocalContext;
-            _threadLocalContext = newContext;
-        }
-
-        /// <summary>
-        /// Resets the current grain context.
+        /// Resets the current grain context to the provided original context.
         /// </summary>
+        /// <param name="originalContext">The original context.</param>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal static void ResetExecutionContext()
+        internal static void ResetExecutionContext(IGrainContext originalContext)
         {
-            _threadLocalContext = null;
+            _threadLocalContext = originalContext;
         }
     }
 }

src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs

@@ -116,7 +116,7 @@ public IGrainContext CreateContext(GrainAddress activationAddress)
                     _serviceProvider,
                     _sharedComponents);
 
-                RuntimeContext.SetExecutionContext(context, out var existingContext);
+                RuntimeContext.SetExecutionContext(context, out var originalContext);
 
                 try
                 {
@@ -126,7 +126,7 @@ public IGrainContext CreateContext(GrainAddress activationAddress)
                 }
                 finally
                 {
-                    RuntimeContext.SetExecutionContext(existingContext);
+                    RuntimeContext.ResetExecutionContext(originalContext);
                 }
 
                 return context;

src/Orleans.Runtime/Core/SystemTarget.cs

@@ -280,7 +280,7 @@ public void ReceiveMessage(object message)
                     {
                         this.MessagingTrace.OnEnqueueMessageOnActivation(msg, this);
                         var workItem = new RequestWorkItem(this, msg);
-                        this.WorkItemGroup.TaskScheduler.QueueWorkItem(workItem);
+                        this.WorkItemGroup.TaskScheduler.QueueRequestWorkItem(workItem);
                         break;
                     }
 

src/Orleans.Runtime/Scheduler/ActivationTaskScheduler.cs

@@ -38,7 +38,6 @@ internal ActivationTaskScheduler(WorkItemGroup workGroup, ILogger<ActivationTask
 
         public void RunTask(Task task)
         {
-            RuntimeContext.SetExecutionContext(workerGroup.GrainContext);
             bool done = TryExecuteTask(task);
             if (!done)
                 logger.LogWarning(

src/Orleans.Runtime/Scheduler/ClosureWorkItem.cs

@@ -27,9 +27,9 @@ public AsyncClosureWorkItem(Func<Task> closure, IGrainContext grainContext)
 
         public override async void Execute()
         {
+            RuntimeContext.SetExecutionContext(GrainContext, out var originalContext);
             try
             {
-                RuntimeContext.SetExecutionContext(this.GrainContext);
                 RequestContext.Clear();
                 await this.continuation();
                 this.completion.TrySetResult(true);
@@ -40,7 +40,7 @@ public override async void Execute()
             }
             finally
             {
-                RuntimeContext.ResetExecutionContext();
+                RuntimeContext.ResetExecutionContext(originalContext);
             }
         }
 
@@ -73,9 +73,9 @@ public AsyncClosureWorkItem(Func<Task<T>> closure, IGrainContext grainContext)
 
         public override async void Execute()
         {
+            RuntimeContext.SetExecutionContext(GrainContext, out var originalContext);
             try
             {
-                RuntimeContext.SetExecutionContext(this.GrainContext);
                 RequestContext.Clear();
                 var result = await this.continuation();
                 this.completion.TrySetResult(result);
@@ -86,7 +86,7 @@ public override async void Execute()
             }
             finally
             {
-                RuntimeContext.ResetExecutionContext();
+                RuntimeContext.ResetExecutionContext(originalContext);
             }
         }
 

src/Orleans.Runtime/Scheduler/RequestWorkItem.cs

@@ -5,28 +5,28 @@ namespace Orleans.Runtime.Scheduler
     internal sealed class RequestWorkItem : WorkItemBase
     {
         private readonly Message request;
-        private readonly SystemTarget target;
+        private readonly SystemTarget _target;
 
         public RequestWorkItem(SystemTarget t, Message m)
         {
-            target = t;
+            _target = t;
             request = m;
         }
 
         public override string Name => $"RequestWorkItem:Id={request.Id}";
 
-        public override IGrainContext GrainContext => this.target;
+        public override IGrainContext GrainContext => _target;
 
         public override void Execute()
         {
+            RuntimeContext.SetExecutionContext(_target, out var originalContext);
             try
             {
-                RuntimeContext.SetExecutionContext(this.target);
-                target.HandleNewRequest(request);
+                _target.HandleNewRequest(request);
             }
             finally
             {
-                RuntimeContext.ResetExecutionContext();
+                RuntimeContext.ResetExecutionContext(originalContext);
             }
         }
 

src/Orleans.Runtime/Scheduler/TaskSchedulerUtils.cs

@@ -7,52 +7,53 @@ namespace Orleans.Runtime.Scheduler
 {
     internal static class TaskSchedulerUtils
     {
-        private static readonly Action<object> TaskFunc = (state) => RunWorkItemTask((IWorkItem)state);
+        private static readonly Action<object> TaskFunc = RunWorkItemTask;
         private static readonly Action<object> ThreadPoolWorkItemTaskFunc = (state) => RunThreadPoolWorkItemTask((IThreadPoolWorkItem)state);
 
         private static void RunThreadPoolWorkItemTask(IThreadPoolWorkItem todo)
         {
             todo.Execute();
         }
 
-        private static void RunWorkItemTask(IWorkItem todo)
+        private static void RunWorkItemTask(object state)
         {
+            var workItem = (RequestWorkItem)state;
+            RuntimeContext.SetExecutionContext(workItem.GrainContext, out var originalContext);
             try
             {
-                RuntimeContext.SetExecutionContext(todo.GrainContext);
-                todo.Execute();
+                workItem.Execute();
             }
             finally
             {
-                RuntimeContext.ResetExecutionContext();
+                RuntimeContext.ResetExecutionContext(originalContext);
             }
         }
 
-        public static void QueueAction(this TaskScheduler taskScheduler, Action action)
+        public static void QueueAction(this ActivationTaskScheduler taskScheduler, Action action)
         {
             using var suppressExecutionContext = new ExecutionContextSuppressor(); 
 
             var task = new Task(action);
             task.Start(taskScheduler);
         }
 
-        public static void QueueAction(this TaskScheduler taskScheduler, Action<object> action, object state)
+        public static void QueueAction(this ActivationTaskScheduler taskScheduler, Action<object> action, object state)
         {
             using var suppressExecutionContext = new ExecutionContextSuppressor(); 
 
             var task = new Task(action, state);
             task.Start(taskScheduler);
         }
 
-        public static void QueueWorkItem(this TaskScheduler taskScheduler, IWorkItem todo)
+        public static void QueueRequestWorkItem(this ActivationTaskScheduler taskScheduler, RequestWorkItem requestWorkItem)
         {
             using var suppressExecutionContext = new ExecutionContextSuppressor(); 
 
-            var workItemTask = new Task(TaskFunc, todo);
+            var workItemTask = new Task(TaskFunc, requestWorkItem);
             workItemTask.Start(taskScheduler);
         }
 
-        public static void QueueThreadPoolWorkItem(this TaskScheduler taskScheduler, IThreadPoolWorkItem workItem)
+        public static void QueueThreadPoolWorkItem(this ActivationTaskScheduler taskScheduler, IThreadPoolWorkItem workItem)
         {
             using var suppressExecutionContext = new ExecutionContextSuppressor(); 
 

src/Orleans.Runtime/Scheduler/WorkItemGroup.cs

@@ -164,12 +164,12 @@ private static object DumpAsyncState(object o)
 
         // Execute one or more turns for this activation.
         // This method is always called in a single-threaded environment -- that is, no more than one
-        // thread will be in this method at once -- but other asynch threads may still be queueing tasks, etc.
+        // thread will be in this method at once -- but other async threads may still be queueing tasks, etc.
         public void Execute()
         {
+            RuntimeContext.SetExecutionContext(GrainContext, out var originalContext);
             try
             {
-                RuntimeContext.SetExecutionContext(this.GrainContext);
 
                 // Process multiple items -- drain the applicationMessageQueue (up to max items) for this physical activation
                 int count = 0;
@@ -264,7 +264,7 @@ public void Execute()
                     }
                 }
 
-                RuntimeContext.ResetExecutionContext();
+                RuntimeContext.ResetExecutionContext(originalContext);
             }
         }