Added IScheduler.Marshal()

This is to allow the scheduler to ensure you're within an expected
execution context, without needing to force a yield if you're already
in the expected or desired context

This includes major refactoring to collapse the implementation into Scheduler
for simplicity. Added unit tests to check behavior of Marshal() and Yield()
with respect to threads and the sync context

Author: AshleighAdams

src/CoroutineScheduler.Core/IScheduler.cs

@@ -23,6 +23,12 @@ public interface IScheduler
 	/// </summary>
 	/// <returns>An awaitable whose awaiter completes at the mercy of this scheduler.</returns>
 	YieldTask Yield();
+
+	/// <summary>
+	/// Marshal to this scheduler.<br/>
+	/// </summary>
+	/// <returns>An awaitable whose awaiter completes at the mercy of this scheduler for the purposes of marshalling in some way, such as between threads.</returns>
+	YieldTask Marshal();
 }
 
 /// <summary>
@@ -36,7 +42,6 @@ public record struct YieldTask
 	/// <summary>
 	/// A struct container to provide the compiler duck-typeing so an <see cref="IYieldAwaiter"/> can be <c>await</c>ed.
 	/// </summary>
-	/// <param name="awaiter"></param>
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	[DebuggerNonUserCode]
 	public YieldTask(IYieldAwaiter awaiter)

src/CoroutineScheduler.Core/PublicAPI.Unshipped.txt

@@ -1,5 +1,6 @@
 #nullable enable
 CoroutineScheduler.IScheduler
+CoroutineScheduler.IScheduler.Marshal() -> CoroutineScheduler.YieldTask
 CoroutineScheduler.IScheduler.SpawnTask(System.Func<System.Threading.Tasks.Task!>! func) -> System.Threading.Tasks.Task!
 CoroutineScheduler.IScheduler.Yield() -> CoroutineScheduler.YieldTask
 CoroutineScheduler.IYieldAwaiter

src/CoroutineScheduler/AsyncManualResetEvent.cs

@@ -14,6 +14,7 @@ CoroutineScheduler.InternalReleaseException.InternalReleaseException(string! mes
 CoroutineScheduler.InternalReleaseException.InternalReleaseException(string! message, System.Exception! innerException) -> void
 CoroutineScheduler.InternalReleaseException.InternalReleaseException(System.Runtime.Serialization.SerializationInfo! info, System.Runtime.Serialization.StreamingContext context) -> void
 CoroutineScheduler.Scheduler
+CoroutineScheduler.Scheduler.Marshal() -> CoroutineScheduler.YieldTask
 CoroutineScheduler.Scheduler.Resume() -> void
 CoroutineScheduler.Scheduler.Scheduler() -> void
 CoroutineScheduler.Scheduler.SpawnTask(System.Func<System.Threading.Tasks.Task!>! func) -> System.Threading.Tasks.Task!

src/CoroutineScheduler/PublicAPI.Unshipped.txt

@@ -1,11 +1,11 @@
 #if DEBUG
 #define DEBUG_ASYNC
 #endif
-
 #if !DEBUG_ASYNC
 using System.Diagnostics;
 #endif
 
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
@@ -46,17 +46,25 @@ internal UnhandledExceptionEventArgs(Exception ex)
 /// </summary>
 public sealed class Scheduler : IScheduler
 {
-	private AsyncManualResetEvent Awaiter { get; } = new AsyncManualResetEvent();
 	private SchedulerSynchronizationContext SyncContext { get; }
+	private YieldAwaiter NormalAwaiter { get; }
+	private YieldAwaiter MarshalAwaiter { get; }
 
 	/// <summary>
 	/// Create a new scheduler
 	/// </summary>
 	public Scheduler()
 	{
-		SyncContext = new(SyncronizationContextPost);
+		SyncContext = new(PostContinuation);
+		NormalAwaiter = new(this, true);
+		MarshalAwaiter = new(this, false);
 	}
 
+	private ContextCallback? Runner { get; set; }
+
+	private int? ResumingOnThread { get; set; }
+	private bool RequiresMarshalling => !ResumingOnThread.HasValue || ResumingOnThread.Value != Thread.CurrentThread.ManagedThreadId;
+
 	/// <summary>
 	/// Resume any tasks that are at the point of invocation are currently awaiting `Yield()`, and <br />
 	/// accept any posted continuations from the syncronization context for the duration of the call.
@@ -69,19 +77,59 @@ public void Resume()
 	{
 		var original = SynchronizationContext.Current;
 		SynchronizationContext.SetSynchronizationContext(SyncContext);
+		ResumingOnThread = Thread.CurrentThread.ManagedThreadId;
 		{
-			Awaiter.Release();
+			// cache this allocation
+			[DebuggerNonUserCode]
+			[MethodImpl(MethodImplOptions.AggressiveInlining)]
+			static void executionContextRunner(object obj) => (obj as Action)!();
+			Runner ??= executionContextRunner!;
+
+			ReleaseImplicit();
+
+			int alreadyQueuedCount = explicitQueue.Count;
+			while (alreadyQueuedCount > 0)
+			{
+				if (!explicitQueue.TryDequeue(out var workItem))
+					throw new InternalReleaseException("Failed to dequeue the next continuation.");
+
+				alreadyQueuedCount--;
+
+				if (workItem.Context is null)
+					workItem.Continuation();
+				else
+					ExecutionContext.Run(workItem.Context, Runner, workItem.Continuation);
+
+				ReleaseImplicit();
+			}
 		}
+		ResumingOnThread = null;
 		SynchronizationContext.SetSynchronizationContext(original);
 	}
 
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	private void ReleaseImplicit()
+	{
+		while (implicitQueue.TryDequeue(out var workItem))
+			workItem.Continuation(workItem.Context);
+	}
+
 	/// <summary>
 	/// Wait until <see cref="Resume()"/> is next invoked.
 	/// </summary>
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public YieldTask Yield()
 	{
-		return new(Awaiter);
+		return new(NormalAwaiter);
+	}
+
+	/// <summary>
+	/// Wait until <see cref="Resume()"/> is next invoked.
+	/// </summary>
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public YieldTask Marshal()
+	{
+		return new(MarshalAwaiter);
 	}
 
 	/// <inheritdoc/>
@@ -129,8 +177,75 @@ async void SpawnTaskInternal()
 #if !DEBUG_ASYNC
 	[DebuggerNonUserCode]
 #endif
-	private void SyncronizationContextPost(SendOrPostCallback cb, object? state)
+	private void PostContinuation(SendOrPostCallback cb, object? state)
+	{
+		if (RequiresMarshalling)
+			implicitQueue.Enqueue(new() { Continuation = cb, Context = state });
+		else
+			cb(state);
+	}
+
+#if !DEBUG_ASYNC
+	[DebuggerNonUserCode]
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+#endif
+	private void Yield(Action action, ExecutionContext? ctx)
 	{
-		Awaiter.Post(cb, state);
+		explicitQueue.Enqueue(new() { Continuation = action, Context = ctx });
+	}
+
+	private struct PostedWorkItem
+	{
+		public SendOrPostCallback Continuation { get; set; }
+		public object? Context { get; set; }
+	}
+
+	private struct YieldedWorkItem
+	{
+		public Action Continuation { get; set; }
+		public ExecutionContext? Context { get; set; }
+	}
+
+	private readonly ConcurrentQueue<YieldedWorkItem> explicitQueue = new();
+	private readonly ConcurrentQueue<PostedWorkItem> implicitQueue = new();
+
+	private class YieldAwaiter : IYieldAwaiter
+	{
+		public Scheduler Self { get; }
+		public bool ForceYield { get; }
+
+		public YieldAwaiter(Scheduler self, bool forceYield)
+		{
+			Self = self;
+			ForceYield = forceYield;
+		}
+
+		bool IYieldAwaiter.IsCompleted => ForceYield switch
+		{
+			true => false,
+			false => !Self.RequiresMarshalling,
+		};
+#if !DEBUG_ASYNC
+		[DebuggerNonUserCode]
+#endif
+		void INotifyCompletion.OnCompleted(Action continuation)
+		{
+			Self.Yield(continuation, ExecutionContext.Capture());
+		}
+#if !DEBUG_ASYNC
+		[DebuggerNonUserCode]
+#endif
+		void ICriticalNotifyCompletion.UnsafeOnCompleted(Action continuation)
+		{
+			Self.Yield(continuation, null);
+		}
+
+		[MethodImpl(MethodImplOptions.AggressiveInlining)]
+#if !DEBUG_ASYNC
+		[DebuggerNonUserCode]
+#endif
+		void IYieldAwaiter.GetResult()
+		{
+		}
 	}
 }