Add ServiceRpcDescriptor.JoinableTaskFactory property

Directory.Packages.props

@@ -7,6 +7,7 @@
     <CodeAnalysisVersionForAnalyzers>3.11.0</CodeAnalysisVersionForAnalyzers>
     <CodeAnalysisVersion>4.4.0</CodeAnalysisVersion>
     <CodefixTestingVersion>1.1.2-beta1.23110.2</CodefixTestingVersion>
+    <VSThreadingVersion>17.6.29-alpha</VSThreadingVersion>
   </PropertyGroup>
   <ItemGroup>
     <PackageVersion Include="IsExternalInit" Version="1.0.3" />
@@ -24,10 +25,11 @@
     <PackageVersion Include="Microsoft.VisualStudioEng.MicroBuild.Core" Version="1.0.0" />
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="17.4.1" />
     <PackageVersion Include="Microsoft.VisualStudio.Internal.MicroBuild.NonShipping" Version="$(MicroBuildVersion)" />
-    <PackageVersion Include="Microsoft.VisualStudio.Threading.Analyzers" Version="17.5.21" />
+    <PackageVersion Include="Microsoft.VisualStudio.Threading" Version="$(VSThreadingVersion)" />
+    <PackageVersion Include="Microsoft.VisualStudio.Threading.Analyzers" Version="$(VSThreadingVersion)" />
     <PackageVersion Include="Microsoft.Windows.CsWin32" Version="0.2.188-beta" />
     <PackageVersion Include="NuGet.Protocol" Version="6.4.0" />
-    <PackageVersion Include="StreamJsonRpc" Version="2.15.3-alpha" />
+    <PackageVersion Include="StreamJsonRpc" Version="2.15.16-alpha" />
     <PackageVersion Include="System.Collections.Immutable" Version="7.0.0" />
     <PackageVersion Include="System.CommandLine" Version="2.0.0-beta4.22272.1" />
     <PackageVersion Include="TraceSource.ActivityTracing" Version="0.1.201-beta" />

src/Microsoft.ServiceHub.Framework/Container/GlobalBrokeredServiceContainer.ProfferedServiceFactory.cs

@@ -105,8 +105,7 @@ public void Dispose()
 				 : this.Descriptor;
 
 			IServiceBroker serviceBroker = this.Container.GetSecureServiceBroker(options);
-			descriptor = descriptor
-				.WithTraceSource(await this.Container.GetTraceSourceForConnectionAsync(serviceBroker, serviceMoniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+			descriptor = await this.Container.ApplyDescriptorSettingsAsync(descriptor, serviceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 
 			using (options.ApplyCultureToCurrentContext())
 			{
@@ -151,8 +150,7 @@ public void Dispose()
 			cancellationToken.ThrowIfCancellationRequested();
 
 			IServiceBroker serviceBroker = this.Container.GetSecureServiceBroker(options);
-			serviceDescriptor = serviceDescriptor
-				.WithTraceSource(await this.Container.GetTraceSourceForConnectionAsync(serviceBroker, serviceDescriptor.Moniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+			serviceDescriptor = await this.Container.ApplyDescriptorSettingsAsync(serviceDescriptor, serviceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 
 			object? liveObject = await this.InvokeFactoryAsync(serviceBroker, serviceDescriptor.Moniker, options, cancellationToken).ConfigureAwait(false);
 			try

src/Microsoft.ServiceHub.Framework/Container/GlobalBrokeredServiceContainer.ProfferedViewIntrinsicService.cs

@@ -76,8 +76,7 @@ internal ProfferedViewIntrinsicService(GlobalBrokeredServiceContainer container,
 				 : this.Descriptor;
 
 			IServiceBroker serviceBroker = this.Container.GetSecureServiceBroker(options);
-			descriptor = descriptor
-				.WithTraceSource(await this.Container.GetTraceSourceForConnectionAsync(serviceBroker, serviceMoniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+			descriptor = await this.Container.ApplyDescriptorSettingsAsync(descriptor, serviceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 
 			using (options.ApplyCultureToCurrentContext())
 			{
@@ -114,8 +113,7 @@ internal ProfferedViewIntrinsicService(GlobalBrokeredServiceContainer container,
 			where T : class
 		{
 			IServiceBroker serviceBroker = this.Container.GetSecureServiceBroker(options);
-			serviceDescriptor = serviceDescriptor
-				.WithTraceSource(await this.Container.GetTraceSourceForConnectionAsync(serviceBroker, serviceDescriptor.Moniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+			serviceDescriptor = await this.Container.ApplyDescriptorSettingsAsync(serviceDescriptor, serviceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 			var liveObject = (T?)await this.factory(view, serviceDescriptor.Moniker, options, serviceBroker, cancellationToken).ConfigureAwait(false);
 			return serviceDescriptor.ConstructLocalProxy(liveObject);
 		}

src/Microsoft.ServiceHub.Framework/Container/GlobalBrokeredServiceContainer.View.cs

@@ -223,8 +223,7 @@ public event EventHandler<BrokeredServicesChangedEventArgs>? AvailabilityChanged
 				(proffered, errorCode) = await this.TryGetProfferingSourceAsync(serviceDescriptor.Moniker, cancellationToken).ConfigureAwait(false);
 				if (proffered is object)
 				{
-					serviceDescriptor = serviceDescriptor
-						.WithTraceSource(await this.container.GetTraceSourceForConnectionAsync(this, serviceDescriptor.Moniker, options, clientRole: true, cancellationToken).ConfigureAwait(false));
+					serviceDescriptor = await this.container.ApplyDescriptorSettingsAsync(serviceDescriptor, this, options, clientRole: true, cancellationToken).ConfigureAwait(false);
 
 					GC.KeepAlive(typeof(ValueTask<T>)); // workaround CLR bug https://dev.azure.com/devdiv/DevDiv/_workitems/edit/1358442
 					T? proxy = proffered is ProfferedViewIntrinsicService viewIntrinsic

src/Microsoft.ServiceHub.Framework/Container/GlobalBrokeredServiceContainer.cs

@@ -269,6 +269,23 @@ public IServiceBroker GetSecureServiceBroker(ServiceActivationOptions options)
 	/// </devremarks>
 	public ValueTask<TraceSource?> GetTraceSourceForBrokeredServiceAsync(IServiceBroker serviceBroker, ServiceMoniker serviceMoniker, ServiceActivationOptions options, bool clientRole, CancellationToken cancellationToken) => this.GetTraceSourceForConnectionAsync(serviceBroker, serviceMoniker, options, clientRole, cancellationToken);
 
+	/// <summary>
+	/// Applies typical transformations on a descriptor fro brokered service clients and services.
+	/// </summary>
+	/// <param name="descriptor">The stock descriptor used for this service.</param>
+	/// <param name="serviceBroker">A service broker that may be used to acquire other, related services as necessary to mutate the <paramref name="descriptor"/>.</param>
+	/// <param name="serviceActivationOptions">The activation options for the service.</param>
+	/// <param name="clientRole">A value indicating whether the <paramref name="descriptor"/> is about to be used to activate a client proxy or client connection; use <see langword="false" /> when activating the service itself.</param>
+	/// <param name="cancellationToken">A cancellation token.</param>
+	/// <returns>The modified descriptor.</returns>
+	internal async ValueTask<ServiceRpcDescriptor> ApplyDescriptorSettingsAsync(ServiceRpcDescriptor descriptor, IServiceBroker serviceBroker, ServiceActivationOptions serviceActivationOptions, bool clientRole, CancellationToken cancellationToken)
+	{
+		TraceSource? traceSource = await this.GetTraceSourceForConnectionAsync(serviceBroker, descriptor.Moniker, serviceActivationOptions, clientRole, cancellationToken).ConfigureAwait(false);
+		return descriptor
+			.WithJoinableTaskFactory(this.joinableTaskFactory)
+			.WithTraceSource(traceSource);
+	}
+
 	/// <summary>
 	/// Registers a set of services with the global broker. This is separate from proffering a service. A service should be registered before it is proffered.
 	/// An <see cref="IServiceBroker.AvailabilityChanged"/> event is never

src/Microsoft.ServiceHub.Framework/Container/ServiceBrokerOfExportedServices.cs

@@ -96,8 +96,7 @@ public void RegisterAndProfferServices(GlobalBrokeredServiceContainer container)
 					return null;
 				}
 
-				ServiceRpcDescriptor descriptor = brokeredService.Descriptor
-					.WithTraceSource(await container.GetTraceSourceForBrokeredServiceAsync(contextualServiceBroker, serviceMoniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+				ServiceRpcDescriptor descriptor = await container.ApplyDescriptorSettingsAsync(brokeredService.Descriptor, contextualServiceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 
 				if (descriptor is ServiceJsonRpcDescriptor { MultiplexingStreamOptions: null } oldJsonRpcDescriptor)
 				{
@@ -163,8 +162,7 @@ public void RegisterAndProfferServices(GlobalBrokeredServiceContainer container)
 
 			try
 			{
-				serviceDescriptor = serviceDescriptor
-					.WithTraceSource(await container.GetTraceSourceForBrokeredServiceAsync(contextualServiceBroker, serviceDescriptor.Moniker, options, clientRole: false, cancellationToken).ConfigureAwait(false));
+				serviceDescriptor = await container.ApplyDescriptorSettingsAsync(serviceDescriptor, contextualServiceBroker, options, clientRole: false, cancellationToken).ConfigureAwait(false);
 
 				brokeredService = await export.Value.CreateBrokeredServiceAsync(cancellationToken).ConfigureAwait(false);
 				if (brokeredService is null)

src/Microsoft.ServiceHub.Framework/PublicAPI.Unshipped.txt

@@ -30,6 +30,8 @@ Microsoft.ServiceHub.Framework.ServiceHubHostRemoteServiceBroker.Dispose() -> vo
 Microsoft.ServiceHub.Framework.ServiceHubHostRemoteServiceBroker.GetPipeAsync(Microsoft.ServiceHub.Framework.ServiceMoniker! serviceMoniker, Microsoft.ServiceHub.Framework.ServiceActivationOptions options, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask<System.IO.Pipelines.IDuplexPipe?>
 Microsoft.ServiceHub.Framework.ServiceHubHostRemoteServiceBroker.GetProxyAsync<T>(Microsoft.ServiceHub.Framework.ServiceRpcDescriptor! serviceDescriptor, Microsoft.ServiceHub.Framework.ServiceActivationOptions options, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask<T?>
 Microsoft.ServiceHub.Framework.ServiceHubHostRemoteServiceBroker.ServiceHubHostRemoteServiceBroker(Microsoft.ServiceHub.Framework.IServiceBroker! inner) -> void
+Microsoft.ServiceHub.Framework.ServiceRpcDescriptor.JoinableTaskFactory.get -> Microsoft.VisualStudio.Threading.JoinableTaskFactory?
+Microsoft.ServiceHub.Framework.ServiceRpcDescriptor.WithJoinableTaskFactory(Microsoft.VisualStudio.Threading.JoinableTaskFactory? joinableTaskFactory) -> Microsoft.ServiceHub.Framework.ServiceRpcDescriptor!
 Microsoft.ServiceHub.Framework.Services.IBrokeredServiceManifest
 Microsoft.ServiceHub.Framework.Services.IBrokeredServiceManifest.GetAvailableServicesAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask<System.Collections.Generic.IReadOnlyCollection<Microsoft.ServiceHub.Framework.ServiceMoniker!>!>
 Microsoft.ServiceHub.Framework.Services.IBrokeredServiceManifest.GetAvailableVersionsAsync(string! serviceName, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask<System.Collections.Immutable.ImmutableSortedSet<System.Version?>!>

src/Microsoft.ServiceHub.Framework/ServiceJsonRpcDescriptor.cs

@@ -260,6 +260,7 @@ public override RpcConnection ConstructRpcConnection(IDuplexPipe pipe)
 		jsonRpc.ExceptionStrategy = this.ExceptionStrategy;
 		jsonRpc.ActivityTracingStrategy = new CorrelationManagerTracingStrategy { TraceSource = this.TraceSource };
 		jsonRpc.SynchronizationContext = new NonConcurrentSynchronizationContext(sticky: false);
+		jsonRpc.JoinableTaskFactory = this.JoinableTaskFactory;
 
 		// Only set the TraceSource if we've been given one, so we don't set it to null (which would throw).
 		if (this.TraceSource != null)

src/Microsoft.ServiceHub.Framework/ServiceRpcDescriptor.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.IO.Pipelines;
 using System.Reflection;
+using Microsoft.VisualStudio.Threading;
 using Nerdbank.Streams;
 
 namespace Microsoft.ServiceHub.Framework;
@@ -40,6 +41,7 @@ protected ServiceRpcDescriptor(ServiceRpcDescriptor copyFrom)
 		this.ClientInterface = copyFrom.ClientInterface;
 		this.TraceSource = copyFrom.TraceSource;
 		this.MultiplexingStream = copyFrom.MultiplexingStream;
+		this.JoinableTaskFactory = copyFrom.JoinableTaskFactory;
 	}
 
 	/// <summary>
@@ -64,6 +66,12 @@ protected ServiceRpcDescriptor(ServiceRpcDescriptor copyFrom)
 	/// <value><see langword="null"/> by default.</value>
 	public MultiplexingStream? MultiplexingStream { get; private set; }
 
+	/// <summary>
+	/// Gets the <see cref="Microsoft.VisualStudio.Threading.JoinableTaskFactory"/> that may be applied
+	/// to the constructed RPC connection.
+	/// </summary>
+	public JoinableTaskFactory? JoinableTaskFactory { get; private set; }
+
 	/// <summary>
 	/// Gets the interface type that the client's "callback" RPC target is expected to implement.
 	/// </summary>
@@ -176,6 +184,24 @@ public ServiceRpcDescriptor WithTraceSource(TraceSource? traceSource)
 		return result;
 	}
 
+	/// <summary>
+	/// Returns an instance of <see cref="ServiceRpcDescriptor"/> that resembles this one,
+	/// but with the <see cref="JoinableTaskFactory" /> property set to the specified value.
+	/// </summary>
+	/// <param name="joinableTaskFactory">The value for the modified property..</param>
+	/// <returns>A clone of this instance, with the property changed. Or this same instance if the property already matches.</returns>
+	public ServiceRpcDescriptor WithJoinableTaskFactory(JoinableTaskFactory? joinableTaskFactory)
+	{
+		if (this.JoinableTaskFactory == joinableTaskFactory)
+		{
+			return this;
+		}
+
+		ServiceRpcDescriptor result = this.Clone();
+		result.JoinableTaskFactory = joinableTaskFactory;
+		return result;
+	}
+
 	/// <summary>
 	/// Returns an instance of <see cref="ServiceRpcDescriptor"/> that resembles this one,
 	/// but with the <see cref="ServiceMoniker" /> property set to the specified value.

test/.editorconfig

@@ -62,3 +62,6 @@ dotnet_diagnostic.VSTHRD003.severity = suggestion
 
 # VSTHRD012: Provide JoinableTaskFactory where allowed
 dotnet_diagnostic.VSTHRD012.severity = suggestion
+
+# VSSDK005: Avoid instantiating JoinableTaskContext
+dotnet_diagnostic.VSSDK005.severity = none

test/Microsoft.ServiceHub.Analyzers.Tests/Helpers/ReferencesHelper.cs

@@ -9,5 +9,5 @@ internal static class ReferencesHelper
 	public static readonly ReferenceAssemblies DefaultReferences = ReferenceAssemblies.Net.Net60
 		.WithPackages(ImmutableArray.Create(
 			new PackageIdentity("System.Threading.Tasks.Extensions", "4.5.4"),
-			new PackageIdentity("Microsoft.VisualStudio.Threading", "17.1.46")));
+			new PackageIdentity("Microsoft.VisualStudio.Threading", "17.6.29-alpha")));
 }

test/Microsoft.ServiceHub.Framework.Tests/ServiceJsonRpcDescriptorTests.cs

@@ -16,6 +16,7 @@ public partial class ServiceJsonRpcDescriptorTests : TestBase
 	private static readonly ServiceMoniker SomeMoniker = new ServiceMoniker("Some name");
 	private static readonly ServiceMoniker SomeOtherMoniker = new ServiceMoniker("Some other name");
 	private static readonly MultiplexingStream.Options MultiplexingStreamOptions = new MultiplexingStream.Options();
+	private static readonly JoinableTaskFactory SomeJoinableTaskFactory = new JoinableTaskContext().Factory;
 
 	public ServiceJsonRpcDescriptorTests(ITestOutputHelper logger)
 		: base(logger)
@@ -402,6 +403,19 @@ public async Task WithTraceSource_ImpactOnConstructedServer()
 		Assert.NotEmpty(myListener.TracedMessages);
 	}
 
+	[Fact]
+	public void WithJoinableTaskFactory_ImpactsJsonRpcProperty()
+	{
+		ServiceJsonRpcDescriptor descriptor = new(SomeMoniker, ServiceJsonRpcDescriptor.Formatters.UTF8, ServiceJsonRpcDescriptor.MessageDelimiters.HttpLikeHeaders);
+		ServiceRpcDescriptor jtfDescriptor = descriptor.WithJoinableTaskFactory(SomeJoinableTaskFactory);
+
+		(System.IO.Pipelines.IDuplexPipe, System.IO.Pipelines.IDuplexPipe) pair = FullDuplexStream.CreatePipePair();
+		descriptor.ConstructRpc(new Calculator(), pair.Item1);
+		ICalculator calc = jtfDescriptor.ConstructRpc<ICalculator>(pair.Item2);
+
+		Assert.Same(SomeJoinableTaskFactory, ((IJsonRpcClientProxy)calc).JsonRpc.JoinableTaskFactory);
+	}
+
 	[Fact]
 	public async Task CamelCasingDoesNotImpactDictionaryKeys()
 	{

test/Microsoft.ServiceHub.Framework.Tests/ServiceRpcDescriptorTests.cs

@@ -4,12 +4,16 @@
 using System.Diagnostics;
 using System.IO.Pipelines;
 using Microsoft.ServiceHub.Framework;
+using Microsoft.VisualStudio.Threading;
+using Nerdbank.Streams;
 using Xunit;
 
 public class ServiceRpcDescriptorTests
 {
 	private static readonly ServiceMoniker SomeMoniker = new ServiceMoniker("Some name");
 	private static readonly ServiceMoniker SomeOtherMoniker = new ServiceMoniker("Some other name");
+	private static readonly TraceSource SomeTraceSource = new("test");
+	private static readonly JoinableTaskFactory SomeJoinableTaskFactory = new JoinableTaskContext().Factory;
 
 	[Fact]
 	public void Ctor_ValidatesInputs()
@@ -83,6 +87,33 @@ public void WithServiceMoniker_PreservesProperties()
 		Assert.True(copy2.SomeOtherSetting);
 	}
 
+	[Fact]
+	public void WithJoinableTaskFactory_ImpactsOnlyNewInstance()
+	{
+		MockServiceRpcDescriptor descriptor = new(SomeMoniker);
+		MockServiceRpcDescriptor copy = (MockServiceRpcDescriptor)descriptor.WithJoinableTaskFactory(SomeJoinableTaskFactory);
+		Assert.Same(SomeJoinableTaskFactory, copy.JoinableTaskFactory);
+		Assert.Null(descriptor.JoinableTaskFactory);
+	}
+
+	[Fact]
+	public void CopyConstructorCopiesProperties()
+	{
+		MultiplexingStream mxstream = MultiplexingStream.Create(Stream.Null);
+
+#pragma warning disable CS0618 // Type or member is obsolete
+		MockServiceRpcDescriptor descriptor = (MockServiceRpcDescriptor)new MockServiceRpcDescriptor(SomeMoniker)
+			.WithJoinableTaskFactory(SomeJoinableTaskFactory)
+			.WithMultiplexingStream(mxstream)
+			.WithTraceSource(SomeTraceSource);
+#pragma warning restore CS0618 // Type or member is obsolete
+
+		MockServiceRpcDescriptor copy = new(descriptor);
+		Assert.Same(mxstream, copy.MultiplexingStream);
+		Assert.Same(SomeTraceSource, copy.TraceSource);
+		Assert.Same(SomeJoinableTaskFactory, copy.JoinableTaskFactory);
+	}
+
 	private class MockServiceRpcDescriptor : ServiceRpcDescriptor
 	{
 		public MockServiceRpcDescriptor(ServiceMoniker serviceMoniker)