From 5e3be6e212865643fe8e78a8e3b9e4d9ce1bf619 Mon Sep 17 00:00:00 2001
From: moozzyk <pawelka@microsoft.com>
Date: Wed, 2 Nov 2016 17:03:30 -0700
Subject: [PATCH] Awaiting tasks returned from hubs and getting the actual
 result

---
 .../RpcEndpoint.cs                            | 23 ++++++++++++++-----
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.AspNetCore.SignalR/RpcEndpoint.cs b/src/Microsoft.AspNetCore.SignalR/RpcEndpoint.cs
index c7fdef4dd7..c2d16922d4 100644
--- a/src/Microsoft.AspNetCore.SignalR/RpcEndpoint.cs
+++ b/src/Microsoft.AspNetCore.SignalR/RpcEndpoint.cs
@@ -13,8 +13,8 @@ namespace Microsoft.AspNetCore.SignalR
     // REVIEW: Should there be an RPC package?
     public class RpcEndpoint<T> : EndPoint where T : class
     {
-        private readonly Dictionary<string, Func<Connection, InvocationDescriptor, InvocationResultDescriptor>> _callbacks
-            = new Dictionary<string, Func<Connection, InvocationDescriptor, InvocationResultDescriptor>>(StringComparer.OrdinalIgnoreCase);
+        private readonly Dictionary<string, Func<Connection, InvocationDescriptor, Task<InvocationResultDescriptor>>> _callbacks
+            = new Dictionary<string, Func<Connection, InvocationDescriptor, Task<InvocationResultDescriptor>>>(StringComparer.OrdinalIgnoreCase);
         private readonly Dictionary<string, Type[]> _paramTypes = new Dictionary<string, Type[]>();
 
         private readonly ILogger _logger;
@@ -61,10 +61,10 @@ namespace Microsoft.AspNetCore.SignalR
                 }
 
                 InvocationResultDescriptor result;
-                Func<Connection, InvocationDescriptor, InvocationResultDescriptor> callback;
+                Func<Connection, InvocationDescriptor, Task<InvocationResultDescriptor>> callback;
                 if (_callbacks.TryGetValue(invocationDescriptor.Method, out callback))
                 {
-                    result = callback(connection, invocationDescriptor);
+                    result = await callback(connection, invocationDescriptor);
                 }
                 else
                 {
@@ -110,7 +110,7 @@ namespace Microsoft.AspNetCore.SignalR
                     _logger.LogDebug("RPC method '{methodName}' is bound", methodName);
                 }
 
-                _callbacks[methodName] = (connection, invocationDescriptor) =>
+                _callbacks[methodName] = async (connection, invocationDescriptor) =>
                 {
                     var invocationResult = new InvocationResultDescriptor()
                     {
@@ -139,7 +139,18 @@ namespace Microsoft.AspNetCore.SignalR
                                 .Zip(parameters, (a, p) => Convert.ChangeType(a, p.ParameterType))
                                 .ToArray();
 
-                            invocationResult.Result = methodInfo.Invoke(value, args);
+                            var result = methodInfo.Invoke(value, args);
+                            var resultTask = result as Task;
+                            if (resultTask != null)
+                            {
+                                await resultTask;
+                                var property = resultTask.GetType().GetProperty("Result");
+                                invocationResult.Result = property?.GetValue(resultTask);
+                            }
+                            else
+                            {
+                                invocationResult.Result = result;
+                            }
                         }
                         catch (TargetInvocationException ex)
                         {