Silverlight/WPF/WP7 Async tasks - fESRI.ArcGIS.Client

using System;
using System.Reflection;
using System.Threading;
using System.Threading.Tasks;
using ESRI.ArcGIS.Client.Tasks;

namespace ESRI.Samples.Async.Tasks
{
    public static partial class Extensions
    {
    ///////////////////////////////////////
    // Base implementation
    ///////////////////////////////////////

    private static PropertyInfo QueryCountEventArgsCountProperty = typeof(QueryCountEventArgs).GetProperty("Count");
        private static PropertyInfo QueryEventArgsFeaturesetProperty = typeof(QueryEventArgs).GetProperty("FeatureSet");
        private static PropertyInfo RelationshipEventArgsResultProperty = typeof(RelationshipEventArgs).GetProperty("Result");
        private static EventInfo QueryExecuteCompletedEvent = typeof(QueryTask).GetEvent("ExecuteCompleted");
        private static EventInfo QueryExecuteCountCompletedEvent = typeof(QueryTask).GetEvent("ExecuteCountCompleted");
        private static EventInfo QueryExecuteRelationshipQueryCompletedEvent = typeof(QueryTask).GetEvent("ExecuteRelationshipQueryCompleted");

        public static Task<FeatureSet> ExecuteTaskAsync(this QueryTask task, Query query)
        {
            return ExecuteTaskAsync(task, query, CancellationToken.None);
        }

        public static Task<FeatureSet> ExecuteTaskAsync(this QueryTask task, Query query, CancellationToken cancellationToken)
        {
            return HandleTaskAndResult<FeatureSet, QueryEventArgs>(
                task, QueryExecuteCompletedEvent, QueryEventArgsFeaturesetProperty,
                () => { task.ExecuteAsync(query); }, cancellationToken);
        }

        public static Task<int> ExecuteCountTaskAsync(this QueryTask task, Query query)
        {
            return ExecuteCountTaskAsync(task, query, CancellationToken.None);
        }

        public static Task<int> ExecuteCountTaskAsync(this QueryTask task, Query query, CancellationToken cancellationToken)
        {
            return HandleTaskAndResult<int, QueryCountEventArgs>(
                task, QueryExecuteCountCompletedEvent, QueryCountEventArgsCountProperty,
                () => { task.ExecuteCountAsync(query); }, cancellationToken);
        }

        public static Task<RelationshipResult> ExecuteRelationshipQueryTaskAsync(this QueryTask task, RelationshipParameter parameter)
        {
            return ExecuteRelationshipQueryTaskAsync(task, parameter, CancellationToken.None);
        }

        public static Task<RelationshipResult> ExecuteRelationshipQueryTaskAsync(this QueryTask task, RelationshipParameter parameter, CancellationToken cancellationToken)
        {
            return HandleTaskAndResult<RelationshipResult, RelationshipEventArgs>(
                task, QueryExecuteRelationshipQueryCompletedEvent, RelationshipEventArgsResultProperty,
                () => { task.ExecuteRelationshipQueryAsync(parameter); }, cancellationToken);
        }


    ///////////////////////////////////////
    // Base implementation
    ///////////////////////////////////////
        private static System.Threading.Tasks.Task<T> HandleTaskAndResult<T,S>(
                this TaskBase task,
                EventInfo completedEvent,
                PropertyInfo resultProperty,
                Action ExecuteMethod,
                CancellationToken cancellationToken)
            where S : TaskEventArgs
        {
            var innertask = HandleTask<S>(task, completedEvent, ExecuteMethod, cancellationToken);
            TaskCompletionSource<T> tcs = new TaskCompletionSource<T>();
            innertask.ContinueWith((e) =>
            {
                if (e.IsCompleted)
                    tcs.SetResult((T)resultProperty.GetValue(e.Result, null));
                else if (e.IsFaulted)
                    tcs.SetException(e.Exception);
                else if (e.IsCanceled)
                    tcs.SetCanceled();
            });
            return tcs.Task;
        }

        private static System.Threading.Tasks.Task<T> HandleTask<T>(
            this TaskBase task,
            EventInfo completedEvent,
            Action ExecuteMethod,
            CancellationToken cancellationToken)
            where T : TaskEventArgs
        {
            TaskCompletionSource<T> tcs = new TaskCompletionSource<T>();
            EventHandler<T> completedHandler = null;
            EventHandler<TaskFailedEventArgs> failedHandler = null;
            Action unregister = () =>
            {
                completedEvent.RemoveEventHandler(task, completedHandler);
                task.Failed -= failedHandler;
            };
            if (cancellationToken.CanBeCanceled)
                cancellationToken.Register(() =>
                {
                    unregister();
                    task.CancelAsync();
                });

            completedHandler = delegate(object sender, T e)
            {
                unregister();
                if (!cancellationToken.IsCancellationRequested)
                    tcs.SetResult(e);
            };
            failedHandler = delegate(object sender, TaskFailedEventArgs e)
            {
                unregister();
                tcs.SetException(e.Error);
            };

            completedEvent.AddEventHandler(task, completedHandler);
            task.Failed += failedHandler;
            try
            {
                ExecuteMethod();
            }
            catch (System.Exception ex)
            {
                unregister();
                tcs.SetException(ex);
            }
            return tcs.Task;
        }
    }
}