Merge pull request #540 from microsoft/fix537

Fix `IProgress<T>` when T is a MessagePack `[Union]` type
microsoft · Sep 3, 2020 · fde90ef · fde90ef
2 parents 2d412ae + d9171bf
commit fde90ef
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 2 deletions.
diff --git a/src/StreamJsonRpc.Tests/JsonRpcMessagePackLengthTests.cs b/src/StreamJsonRpc.Tests/JsonRpcMessagePackLengthTests.cs
@@ -5,12 +5,12 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using MessagePack;
 using MessagePack.Formatters;
 using MessagePack.Resolvers;
-using Microsoft;
 using Microsoft.VisualStudio.Threading;
 using StreamJsonRpc;
 using StreamJsonRpc.Protocol;
@@ -31,6 +31,10 @@ internal interface IMessagePackServer
         Task<string?> AcceptUnionTypeAndReturnStringAsync(UnionBaseClass value, CancellationToken cancellationToken);
 
         Task AcceptUnionTypeAsync(UnionBaseClass value, CancellationToken cancellationToken);
+
+        Task ProgressUnionType(IProgress<UnionBaseClass> progress, CancellationToken cancellationToken);
+
+        IAsyncEnumerable<UnionBaseClass> GetAsyncEnumerableOfUnionType(CancellationToken cancellationToken);
     }
 
     [Fact]
@@ -334,6 +338,37 @@ public async Task UnionType_NamedParameter_NoReturnValue_Proxy()
         Assert.IsType<UnionDerivedClass>(server.ReceivedValue);
     }
 
+    [Fact]
+    public async Task UnionType_AsIProgressTypeArgument()
+    {
+        var server = new MessagePackServer();
+        this.serverRpc.AllowModificationWhileListening = true;
+        this.serverRpc.AddLocalRpcTarget(server);
+        var clientProxy = this.clientRpc.Attach<IMessagePackServer>();
+
+        var reportSource = new TaskCompletionSource<UnionBaseClass>();
+        var progress = new Progress<UnionBaseClass>(v => reportSource.SetResult(v));
+        await clientProxy.ProgressUnionType(progress, this.TimeoutToken);
+        Assert.IsType<UnionDerivedClass>(await reportSource.Task.WithCancellation(this.TimeoutToken));
+    }
+
+    [Fact]
+    public async Task UnionType_AsAsyncEnumerableTypeArgument()
+    {
+        var server = new MessagePackServer();
+        this.serverRpc.AllowModificationWhileListening = true;
+        this.serverRpc.AddLocalRpcTarget(server);
+        var clientProxy = this.clientRpc.Attach<IMessagePackServer>();
+
+        UnionBaseClass? actualItem = null;
+        await foreach (UnionBaseClass item in clientProxy.GetAsyncEnumerableOfUnionType(this.TimeoutToken))
+        {
+            actualItem = item;
+        }
+
+        Assert.IsType<UnionDerivedClass>(actualItem);
+    }
+
     protected override void InitializeFormattersAndHandlers(bool controlledFlushingClient)
     {
         this.serverMessageFormatter = new MessagePackFormatter();
@@ -407,6 +442,18 @@ public Task AcceptUnionTypeAsync(UnionBaseClass value, CancellationToken cancell
         }
 
         public UnionBaseClass ReturnUnionType() => new UnionDerivedClass();
+
+        public Task ProgressUnionType(IProgress<UnionBaseClass> progress, CancellationToken cancellationToken)
+        {
+            progress.Report(new UnionDerivedClass());
+            return Task.CompletedTask;
+        }
+
+        public async IAsyncEnumerable<UnionBaseClass> GetAsyncEnumerableOfUnionType([EnumeratorCancellation] CancellationToken cancellationToken)
+        {
+            await Task.Yield();
+            yield return new UnionDerivedClass();
+        }
     }
 
     private class DelayedFlushingHandler : LengthHeaderMessageHandler, IControlledFlushHandler

diff --git a/src/StreamJsonRpc/Reflection/MessageFormatterProgressTracker.cs b/src/StreamJsonRpc/Reflection/MessageFormatterProgressTracker.cs
@@ -269,7 +269,9 @@ public JsonProgress(JsonRpc rpc, object token)
             /// <param name="value">The typed value that will be send in the notification to be reported by the original <see cref="IProgress{T}"/> instance.</param>
             public void Report(T value)
             {
-                this.rpc.NotifyAsync(ProgressRequestSpecialMethod, this.token, value).ContinueWith(
+                var arguments = new object?[] { this.token, value };
+                var argumentDeclaredTypes = new Type[] { this.token.GetType(), typeof(T) };
+                this.rpc.NotifyAsync(ProgressRequestSpecialMethod, arguments, argumentDeclaredTypes).ContinueWith(
                     (t, s) => ((JsonRpc)s).TraceSource.TraceEvent(System.Diagnostics.TraceEventType.Error, (int)JsonRpc.TraceEvents.ProgressNotificationError, "Failed to send progress update. {0}", t.Exception.InnerException ?? t.Exception),
                     this.rpc,
                     CancellationToken.None,