Awaitable TaskScheduler

gistfile1.cs

        /// <summary>
        /// Gets an awaitable object that immediately schedules any continuations on the given scheduler.
        /// </summary>
        /// <param name="scheduler">The scheduler to use for continuations.</param>
        /// <returns>An awaitable object.</returns>
        internal static TaskSchedulerAwaiter GetAwaiter(this TaskScheduler scheduler)
        {
            Contract.Requires<ArgumentNullException>(scheduler != null);
            return new TaskSchedulerAwaiter(scheduler);
        }

    /// <summary>
    /// An awaitable that switches the caller to a specific <see cref="TaskScheduler"/>.
    /// </summary>
    public struct TaskSchedulerAwaitable
    {
        /// <summary>
        /// The scheduler to switch to, if any.
        /// </summary>
        private readonly TaskScheduler scheduler;

        /// <summary>
        /// Initializes a new instance of the <see cref="TaskSchedulerAwaitable"/> struct.
        /// </summary>
        /// <param name="scheduler">The scheduler to switch the caller to.  Use <c>null</c> to avoid yielding.</param>
        public TaskSchedulerAwaitable(TaskScheduler scheduler)
        {
            this.scheduler = scheduler;
        }

        /// <summary>
        /// Gets the awaiter.
        /// </summary>
        public TaskSchedulerAwaiter GetAwaiter()
        {
            return new TaskSchedulerAwaiter(this.scheduler, TaskCreationOptions.None);
        }
    }

    /// <summary>
    /// An awaitable object that schedules tasks on a specific task scheduler.
    /// </summary>
    public struct TaskSchedulerAwaiter : INotifyCompletion
    {
        /// <summary>
        /// The scheduler to use for continuation tasks.
        /// </summary>
        private readonly TaskScheduler scheduler;

        /// <summary>
        /// Task creation options.
        /// </summary>
        private readonly TaskCreationOptions options;

        /// <summary>
        /// Initializes a new instance of the <see cref="TaskSchedulerAwaiter"/> struct.
        /// </summary>
        /// <param name="scheduler">The scheduler to use for continuation tasks.  May be <c>null</c> to indicate IsCompleted should be <c>true</c></param>
        /// <param name="options">Options to apply to OnCompleted tasks.</param>
        public TaskSchedulerAwaiter(TaskScheduler scheduler, TaskCreationOptions options = TaskCreationOptions.None)
        {
            this.scheduler = scheduler;
            this.options = options;
        }

        /// <summary>
        /// Gets a value indicating whether the awaitable task has completed.
        /// </summary>
        public bool IsCompleted
        {
            get { return this.scheduler == null; }
        }

        /// <summary>
        /// Schedules some action to execute after the awaitable task has completed.
        /// </summary>
        /// <param name="action">The action to execute.</param>
        public void OnCompleted(Action action)
        {
            Requires.NotNull(action, "action");

            // Normally OnCompleted would use ContinueWith to cause the given action to execute
            // after the awaited operation had completed.  However for TaskScheduler awaits,
            // all we're waiting for is time on the given scheduler, so schedule it immediately.
            Assumes.NotNull(this.scheduler); // IsCompleted returns true in this case, so this method should not be called.
            Task.Factory.StartNew(action, CancellationToken.None, this.options, this.scheduler);
        }

        /// <summary>
        /// Gets the result of the awaited task.
        /// </summary>
        public void GetResult()
        {
            // Scheduling a task scheduler yields no result.
        }
    }