Use JoinableTaskContext from MEF when available

Author: jcansdale

src/GitHub.App/Authentication/TwoFactorChallengeHandler.cs

@@ -6,9 +6,11 @@
 using ReactiveUI;
 using System.Threading.Tasks;
 using GitHub.Api;
-using GitHub.Helpers;
 using GitHub.Extensions;
 using GitHub.ViewModels.Dialog;
+using Microsoft.VisualStudio.Threading;
+using Microsoft.VisualStudio.Shell;
+using Task = System.Threading.Tasks.Task;
 
 namespace GitHub.Authentication
 {
@@ -17,6 +19,12 @@ namespace GitHub.Authentication
     [PartCreationPolicy(CreationPolicy.Shared)]
     public class TwoFactorChallengeHandler : ReactiveObject, IDelegatingTwoFactorChallengeHandler
     {
+        [ImportingConstructor]
+        public TwoFactorChallengeHandler([Import(AllowDefault = true)] JoinableTaskContext joinableTaskContext)
+        {            
+            JoinableTaskFactory = joinableTaskContext?.Factory ?? ThreadHelper.JoinableTaskFactory;
+        }
+
         ILogin2FaViewModel twoFactorDialog;
         public IViewModel CurrentViewModel
         {
@@ -33,7 +41,7 @@ public async Task<TwoFactorChallengeResult> HandleTwoFactorException(TwoFactorAu
         {
             Guard.ArgumentNotNull(exception, nameof(exception));
 
-            await ThreadingHelper.SwitchToMainThreadAsync();
+            await JoinableTaskFactory.SwitchToMainThreadAsync();
 
             var userError = new TwoFactorRequiredUserError(exception);
             var result = await twoFactorDialog.Show(userError);
@@ -50,8 +58,10 @@ public async Task<TwoFactorChallengeResult> HandleTwoFactorException(TwoFactorAu
 
         public async Task ChallengeFailed(Exception exception)
         {
-            await ThreadingHelper.SwitchToMainThreadAsync();
+            await JoinableTaskFactory.SwitchToMainThreadAsync();
             twoFactorDialog.Cancel();
         }
+
+        JoinableTaskFactory JoinableTaskFactory { get; }
     }
 }

src/GitHub.App/Services/RepositoryCloneService.cs

@@ -8,11 +8,11 @@
 using System.Threading.Tasks;
 using GitHub.Api;
 using GitHub.Extensions;
-using GitHub.Helpers;
 using GitHub.Logging;
 using GitHub.Models;
 using GitHub.Primitives;
 using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Threading;
 using Octokit.GraphQL;
 using Octokit.GraphQL.Model;
 using Rothko;
@@ -50,7 +50,8 @@ public RepositoryCloneService(
             IGraphQLClientFactory graphqlFactory,
             IGitHubContextService gitHubContextService,
             IUsageTracker usageTracker,
-            IGitHubServiceProvider sp)
+            IGitHubServiceProvider sp,
+            [Import(AllowDefault = true)] JoinableTaskContext joinableTaskContext)
         {
             this.operatingSystem = operatingSystem;
             this.vsGitServices = vsGitServices;
@@ -59,6 +60,7 @@ public RepositoryCloneService(
             this.gitHubContextService = gitHubContextService;
             this.usageTracker = usageTracker;
             dte = new Lazy<EnvDTE.DTE>(() => sp.GetService<EnvDTE.DTE>());
+            JoinableTaskFactory = joinableTaskContext?.Factory ?? ThreadHelper.JoinableTaskFactory;
 
             defaultClonePath = GetLocalClonePathFromGitProvider(operatingSystem.Environment.GetUserRepositoriesPath());
         }
@@ -208,9 +210,9 @@ public async Task CloneRepository(
             // vsGitServices.Clone() as this must be called on the main thread.
             if (!DestinationDirectoryExists(repositoryPath))
             {
-                await ThreadingHelper.SwitchToPoolThreadAsync();
+                await TaskScheduler.Default;
                 operatingSystem.Directory.CreateDirectory(repositoryPath);
-                await ThreadingHelper.SwitchToMainThreadAsync();
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
             }
 
             try
@@ -251,6 +253,8 @@ string GetLocalClonePathFromGitProvider(string fallbackPath)
 
         public string DefaultClonePath { get { return defaultClonePath; } }
 
+        JoinableTaskFactory JoinableTaskFactory { get; }
+
         class OrganizationAdapter
         {
             public IReadOnlyList<RepositoryListItemModel> Repositories { get; set; }

src/GitHub.App/ViewModels/GitHubPane/PullRequestDetailViewModel.cs

@@ -13,7 +13,6 @@
 using GitHub.Commands;
 using GitHub.Extensions;
 using GitHub.Factories;
-using GitHub.Helpers;
 using GitHub.Logging;
 using GitHub.Models;
 using GitHub.Services;
@@ -25,6 +24,9 @@
 using static System.FormattableString;
 using ReactiveCommand = ReactiveUI.ReactiveCommand;
 using GitHub.Primitives;
+using Microsoft.VisualStudio.Threading;
+using Microsoft.VisualStudio.Shell;
+using Task = System.Threading.Tasks.Task;
 
 namespace GitHub.ViewModels.GitHubPane
 {
@@ -85,7 +87,8 @@ public PullRequestDetailViewModel(
             ISyncSubmodulesCommand syncSubmodulesCommand,
             IViewViewModelFactory viewViewModelFactory,
             IGitService gitService,
-            IOpenIssueishDocumentCommand openDocumentCommand)
+            IOpenIssueishDocumentCommand openDocumentCommand,
+            [Import(AllowDefault = true)] JoinableTaskContext joinableTaskContext)
         {
             Guard.ArgumentNotNull(pullRequestsService, nameof(pullRequestsService));
             Guard.ArgumentNotNull(sessionManager, nameof(sessionManager));
@@ -106,6 +109,7 @@ public PullRequestDetailViewModel(
             this.viewViewModelFactory = viewViewModelFactory;
             this.gitService = gitService;
             this.openDocumentCommand = openDocumentCommand;
+            JoinableTaskFactory = joinableTaskContext?.Factory ?? ThreadHelper.JoinableTaskFactory;
 
             Files = files;
 
@@ -471,7 +475,7 @@ public override async Task Refresh()
         {
             try
             {
-                await ThreadingHelper.SwitchToMainThreadAsync();
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
 
                 Error = null;
                 OperationError = null;
@@ -734,5 +738,7 @@ public UpdateCommandState(
             public string SyncSubmodulesToolTip { get; }
             public int SubmodulesToSync { get; }
         }
+
+        JoinableTaskFactory JoinableTaskFactory { get; }
     }
 }

src/GitHub.VisualStudio/GitHubPackage.cs

@@ -381,7 +381,7 @@ async Task<object> CreateService(IAsyncServiceContainer container, CancellationT
                 Assumes.Present(sp);
 
                 var environment = new Rothko.Environment();
-                return new UsageService(sp, environment);
+                return new UsageService(sp, environment, ThreadHelper.JoinableTaskContext);
             }
             else if (serviceType == typeof(IUsageTracker))
             {
@@ -393,7 +393,7 @@ async Task<object> CreateService(IAsyncServiceContainer container, CancellationT
                 Assumes.Present(serviceProvider);
                 Assumes.Present(settings);
 
-                return new UsageTracker(serviceProvider, usageService, settings);
+                return new UsageTracker(serviceProvider, usageService, settings, ThreadHelper.JoinableTaskContext);
             }
             else if (serviceType == typeof(IVSGitExt))
             {

src/GitHub.VisualStudio/Services/LocalRepositories.cs

@@ -3,8 +3,10 @@
 using System.Linq;
 using System.Threading.Tasks;
 using GitHub.Extensions;
-using GitHub.Helpers;
 using GitHub.Models;
+using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Threading;
+using Task = System.Threading.Tasks.Task;
 
 namespace GitHub.Services
 {
@@ -21,17 +23,18 @@ public class LocalRepositories : ILocalRepositories
         readonly IVSGitServices vsGitServices;
 
         [ImportingConstructor]
-        public LocalRepositories(IVSGitServices vsGitServices)
+        public LocalRepositories(IVSGitServices vsGitServices, [Import(AllowDefault = true)] JoinableTaskContext joinableTaskContext)
         {
             this.vsGitServices = vsGitServices;
+            JoinableTaskFactory = joinableTaskContext?.Factory ?? ThreadHelper.JoinableTaskFactory;
         }
 
         /// <inheritdoc/>
         public async Task Refresh()
         {
-            await ThreadingHelper.SwitchToPoolThreadAsync();
+            await TaskScheduler.Default;
             var list = vsGitServices.GetKnownRepositories();
-            await ThreadingHelper.SwitchToMainThreadAsync();
+            await JoinableTaskFactory.SwitchToMainThreadAsync();
 
             repositories.Except(list).ToList().ForEach(x => repositories.Remove(x));
             list.Except(repositories).ToList().ForEach(x => repositories.Add(x));
@@ -42,5 +45,7 @@ readonly ObservableCollectionEx<LocalRepositoryModel> repositories
 
         /// <inheritdoc/>
         public IReadOnlyObservableCollection<LocalRepositoryModel> Repositories => repositories;
+
+        JoinableTaskFactory JoinableTaskFactory { get; }
     }
 }

src/GitHub.VisualStudio/Services/UsageService.cs

@@ -5,11 +5,12 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
-using GitHub.Helpers;
 using GitHub.Logging;
 using GitHub.Models;
 using Serilog;
 using Rothko;
+using Microsoft.VisualStudio.Threading;
+using Microsoft.VisualStudio.Shell;
 using Environment = System.Environment;
 using Task = System.Threading.Tasks.Task;
 
@@ -31,10 +32,12 @@ public sealed class UsageService : IUsageService, IDisposable
         Guid? userGuid;
 
         [ImportingConstructor]
-        public UsageService(IGitHubServiceProvider serviceProvider, IEnvironment environment)
+        public UsageService(IGitHubServiceProvider serviceProvider, IEnvironment environment,
+            [Import(AllowDefault = true)] JoinableTaskContext joinableTaskContext)
         {
             this.serviceProvider = serviceProvider;
             this.environment = environment;
+            JoinableTaskFactory = joinableTaskContext?.Factory ?? ThreadHelper.JoinableTaskFactory;
         }
 
         public void Dispose()
@@ -135,7 +138,7 @@ async Task Initialize()
         {
             if (storePath == null)
             {
-                await ThreadingHelper.SwitchToMainThreadAsync();
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
 
                 var program = serviceProvider.GetService<IProgram>();
 
@@ -186,5 +189,7 @@ class UserData
         {
             public Guid UserGuid { get; set; }
         }
+
+        JoinableTaskFactory JoinableTaskFactory { get; }
     }
 }

src/GitHub.VisualStudio/Services/UsageTracker.cs

@@ -1,14 +1,13 @@
 using System;
-using System.ComponentModel.Composition;
 using System.Globalization;
 using System.Linq;
 using System.Linq.Expressions;
 using System.Reflection;
 using System.Threading.Tasks;
-using GitHub.Helpers;
 using GitHub.Logging;
 using GitHub.Models;
 using GitHub.Settings;
+using Microsoft.VisualStudio.Threading;
 using Serilog;
 using Task = System.Threading.Tasks.Task;
 
@@ -31,11 +30,13 @@ public sealed class UsageTracker : IUsageTracker, IDisposable
         public UsageTracker(
             IGitHubServiceProvider gitHubServiceProvider,
             IUsageService service,
-            IPackageSettings settings)
+            IPackageSettings settings,
+            JoinableTaskContext joinableTaskContext)
         {
-            this.gitHubServiceProvider = gitHubServiceProvider;
+            makethis.gitHubServiceProvider = gitHubServiceProvider;
             this.service = service;
             this.userSettings = settings;
+            JoinableTaskFactory = joinableTaskContext.Factory;
             timer = StartTimer();
         }
 
@@ -69,7 +70,7 @@ async Task Initialize()
 
             // The services needed by the usage tracker are loaded when they are first needed to
             // improve the startup time of the extension.
-            await ThreadingHelper.SwitchToMainThreadAsync();
+            await JoinableTaskFactory.SwitchToMainThreadAsync();
 
             client = gitHubServiceProvider.TryGetService<IMetricsService>();
             connectionManager = gitHubServiceProvider.GetService<IConnectionManager>();
@@ -141,5 +142,7 @@ async Task<UsageModel> GetCurrentReport(UsageData data)
             current.Dimensions.IsEnterpriseUser = connectionManager.Connections.Any(x => !x.HostAddress.IsGitHubDotCom());
             return current;
         }
+
+        JoinableTaskFactory JoinableTaskFactory { get; }
     }
 }