Skip to content

Commit

Permalink
Remove another reference field from async state machines
Browse files Browse the repository at this point in the history
The async state machine Task-derived type currently adds three fields:
- The StateMachine
- An Action field for caching any delegate created to MoveNext
- The ExecutionContext to flow to the next MoveNext invocation

The other pending PR gets rid of the Action field by using the unused Action field from the base Task for that purpose.

This PR gets rid of the ExecutionContext field by using the unused state object field from the base Task for that purpose.  The field is exposed via the public AsyncState property, so this also uses a bit from the state flags field to prevent this state object from being returned from that property.

The combination of removing those two fields shaves 16 bytes off of every `async Task` state machine box on 64-bit.  The only remaining field added by the state machine type is for the state machine itself, which is required.
  • Loading branch information
stephentoub committed Mar 21, 2023
1 parent 1448de8 commit 7474e69
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

Expand Down Expand Up @@ -290,12 +289,35 @@ private static void ExecutionContextCallback(object? s)
private Action? _moveNextAction;
/// <summary>The state machine itself.</summary>
public TStateMachine? StateMachine; // mutable struct; do not make this readonly. SOS DumpAsync command depends on this name.
/// <summary>Captured ExecutionContext with which to invoke <see cref="MoveNextAction"/>; may be null.</summary>
public ExecutionContext? Context;

public AsyncStateMachineBox()
{
// The async state machine uses the base Task's state object field to store the captured execution context.
// Ensure that state object isn't published out for others to see.
Debug.Assert((m_stateFlags & (int)InternalTaskOptions.PromiseTask) != 0, "Expected state flags to already be configured.");
Debug.Assert(m_stateObject is null, "Expected to be able to use the state object field for ExecutionContext.");
m_stateFlags |= (int)InternalTaskOptions.HiddenState;
}

/// <summary>A delegate to the <see cref="MoveNext()"/> method.</summary>
public Action MoveNextAction => _moveNextAction ??= new Action(MoveNext);

/// <summary>Captured ExecutionContext with which to invoke <see cref="MoveNextAction"/>; may be null.</summary>
/// <remarks>
/// This uses the base Task.m_stateObject field to store the context, as that field is otherwise unused for state machine boxes.
/// This *must* not be set to anything other than null or an ExecutionContext, or it will result in a type safety hole.
/// We also don't want this ExecutionContext exposed out to consumers of the Task via Task.AsyncState, so
/// the ctor sets the HiddenState option to prevent this from leaking out.
/// </remarks>
public ref ExecutionContext? Context
{
get
{
Debug.Assert(m_stateObject is null || m_stateObject is ExecutionContext, $"{nameof(m_stateObject)} must only be for ExecutionContext but contained {m_stateObject}.");
return ref Unsafe.As<object?, ExecutionContext?>(ref m_stateObject);
}
}

internal sealed override void ExecuteFromThreadPool(Thread threadPoolThread) => MoveNext(threadPoolThread);

/// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
//
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-

using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
Expand Down Expand Up @@ -120,7 +119,7 @@ public class Task : IAsyncResult, IDisposable
[ThreadStatic]
internal static Task? t_currentTask; // The currently executing task.

internal static int s_taskIdCounter; // static counter used to generate unique task IDs
private static int s_taskIdCounter; // static counter used to generate unique task IDs

private int m_taskId; // this task's unique ID. initialized only if it is ever requested

Expand All @@ -132,7 +131,7 @@ public class Task : IAsyncResult, IDisposable
// the completion event which will be set when the Future class calls Finish().
// But the event would now be signalled if Cancel() is called

internal object? m_stateObject; // A state object that can be optionally supplied, passed to action.
private protected object? m_stateObject; // A state object that can be optionally supplied, passed to action.
internal TaskScheduler? m_taskScheduler; // The task scheduler this task runs under.

internal volatile int m_stateFlags; // SOS DumpAsync command depends on this name
Expand Down Expand Up @@ -566,6 +565,7 @@ internal void TaskConstructorCore(Delegate? action, object? state, CancellationT
int illegalInternalOptions =
(int)(internalOptions &
~(InternalTaskOptions.PromiseTask |
InternalTaskOptions.HiddenState |
InternalTaskOptions.ContinuationTask |
InternalTaskOptions.LazyCancellation |
InternalTaskOptions.QueuedByRuntime));
Expand Down Expand Up @@ -1446,7 +1446,7 @@ WaitHandle IAsyncResult.AsyncWaitHandle
/// Gets the state object supplied when the <see cref="Task">Task</see> was created,
/// or null if none was supplied.
/// </summary>
public object? AsyncState => m_stateObject;
public object? AsyncState => (m_stateFlags & (int)InternalTaskOptions.HiddenState) == 0 ? m_stateObject : null;

/// <summary>
/// Gets an indication of whether the asynchronous operation completed synchronously.
Expand Down Expand Up @@ -6717,6 +6717,11 @@ internal enum InternalTaskOptions
ContinuationTask = 0x0200,
PromiseTask = 0x0400,

/// <summary>
/// The state object should not be returned from the AsyncState property.
/// </summary>
HiddenState = 0x0800,

/// <summary>
/// Store the presence of TaskContinuationOptions.LazyCancellation, since it does not directly
/// translate into any TaskCreationOptions.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,19 @@ public void AsyncTaskMethodBuilder_Completed_RemovedFromTracking()
}).Dispose();
}

[Fact]
public void AsyncTaskMethodBuilder_NullStateEvenAfterSuspend()
{
Task t = AwaitSomething();
Assert.Null(t.AsyncState);

static async Task AwaitSomething()
{
Assert.NotNull(ExecutionContext.Capture());
await new TaskCompletionSource().Task;
}
}

#region Helper Methods / Classes

[MethodImpl(MethodImplOptions.NoInlining)]
Expand Down

0 comments on commit 7474e69

Please sign in to comment.