diff --git a/src/Orleans.Core.Abstractions/Runtime/RuntimeContext.cs b/src/Orleans.Core.Abstractions/Runtime/RuntimeContext.cs index 1aea63ee1a..f6f08f95fb 100644 --- a/src/Orleans.Core.Abstractions/Runtime/RuntimeContext.cs +++ b/src/Orleans.Core.Abstractions/Runtime/RuntimeContext.cs @@ -23,31 +23,22 @@ internal static class RuntimeContext /// Sets the current grain context. /// /// The new context. + /// The current context at the time of the call. [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static void SetExecutionContext(IGrainContext newContext) + internal static void SetExecutionContext(IGrainContext newContext, out IGrainContext currentContext) { + currentContext = _threadLocalContext; _threadLocalContext = newContext; } /// - /// Sets the current grain context. - /// - /// The new context. - /// The existing context. - [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static void SetExecutionContext(IGrainContext newContext, out IGrainContext existingContext) - { - existingContext = _threadLocalContext; - _threadLocalContext = newContext; - } - - /// - /// Resets the current grain context. + /// Resets the current grain context to the provided original context. /// + /// The original context. [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static void ResetExecutionContext() + internal static void ResetExecutionContext(IGrainContext originalContext) { - _threadLocalContext = null; + _threadLocalContext = originalContext; } } } diff --git a/src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs b/src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs index 7a70c4d0cc..7ed002d75e 100644 --- a/src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs +++ b/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; diff --git a/src/Orleans.Runtime/Core/SystemTarget.cs b/src/Orleans.Runtime/Core/SystemTarget.cs index 4d8a3a61ec..d388e0b80f 100644 --- a/src/Orleans.Runtime/Core/SystemTarget.cs +++ b/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; } diff --git a/src/Orleans.Runtime/Scheduler/ActivationTaskScheduler.cs b/src/Orleans.Runtime/Scheduler/ActivationTaskScheduler.cs index 53825c751a..6c01eea98c 100644 --- a/src/Orleans.Runtime/Scheduler/ActivationTaskScheduler.cs +++ b/src/Orleans.Runtime/Scheduler/ActivationTaskScheduler.cs @@ -38,7 +38,6 @@ internal ActivationTaskScheduler(WorkItemGroup workGroup, ILogger 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> 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); } } diff --git a/src/Orleans.Runtime/Scheduler/RequestWorkItem.cs b/src/Orleans.Runtime/Scheduler/RequestWorkItem.cs index 79401a35b6..3d624cc9d4 100644 --- a/src/Orleans.Runtime/Scheduler/RequestWorkItem.cs +++ b/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); } } diff --git a/src/Orleans.Runtime/Scheduler/TaskSchedulerUtils.cs b/src/Orleans.Runtime/Scheduler/TaskSchedulerUtils.cs index a4ebbf4cc5..57deb2cb01 100644 --- a/src/Orleans.Runtime/Scheduler/TaskSchedulerUtils.cs +++ b/src/Orleans.Runtime/Scheduler/TaskSchedulerUtils.cs @@ -7,7 +7,7 @@ namespace Orleans.Runtime.Scheduler { internal static class TaskSchedulerUtils { - private static readonly Action TaskFunc = (state) => RunWorkItemTask((IWorkItem)state); + private static readonly Action TaskFunc = RunWorkItemTask; private static readonly Action ThreadPoolWorkItemTaskFunc = (state) => RunThreadPoolWorkItemTask((IThreadPoolWorkItem)state); private static void RunThreadPoolWorkItemTask(IThreadPoolWorkItem todo) @@ -15,20 +15,21 @@ 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(); @@ -36,7 +37,7 @@ public static void QueueAction(this TaskScheduler taskScheduler, Action action) task.Start(taskScheduler); } - public static void QueueAction(this TaskScheduler taskScheduler, Action action, object state) + public static void QueueAction(this ActivationTaskScheduler taskScheduler, Action action, object state) { using var suppressExecutionContext = new ExecutionContextSuppressor(); @@ -44,15 +45,15 @@ public static void QueueAction(this TaskScheduler taskScheduler, Action 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(); diff --git a/src/Orleans.Runtime/Scheduler/WorkItemGroup.cs b/src/Orleans.Runtime/Scheduler/WorkItemGroup.cs index f4a5d453c2..b44b706089 100644 --- a/src/Orleans.Runtime/Scheduler/WorkItemGroup.cs +++ b/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); } }