add CallContext.ResponseHeadersAsync() - more reliable way of getting headers in the middle of requests

Author: mgravell

src/protobuf-net.Grpc/CallContext.cs

@@ -4,6 +4,7 @@
 using System.ComponentModel;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace ProtoBuf.Grpc
 {
@@ -161,6 +162,13 @@ public CallContext(in CallOptions callOptions = default, CallContextFlags flags
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public Metadata ResponseHeaders() => MetadataContext?.Headers ?? ThrowNoContext<Metadata>();
 
+        /// <summary>
+        /// Get the response-headers from a client operation when they are available
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public Task<Metadata> ResponseHeadersAsync() => MetadataContext?.GetHeadersTask(true)
+            ?? Task.FromResult(ThrowNoContext<Metadata>()); // note this actually throws immediately; this is intentional
+
         /// <summary>
         /// Get the response-trailers from a client operation
         /// </summary>

src/protobuf-net.Grpc/Internal/MetadataContext.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 using Grpc.Core;
 
@@ -10,11 +11,42 @@ internal sealed class MetadataContext
         internal MetadataContext(object? state) => State = state;
 
         internal object? State { get; }
-        private Metadata? _headers, _trailers;
+        private Metadata? _trailers;
+        private object? _headersTaskOrSource;
+
         internal Metadata Headers
         {
-            get => _headers ?? Throw("Headers are not yet available");
+            get
+            {
+                var pending = GetHeadersTask(false);
+                return pending is object && pending.RanToCompletion()
+                    ? pending.Result
+                    : Throw("Headers are not yet available");
+            }
+        }
+
+        internal Task<Metadata>? GetHeadersTask(bool createIfMissing)
+        {
+            return _headersTaskOrSource switch
+            {
+                Task<Metadata> task => task,
+                TaskCompletionSource<Metadata> tcs => tcs.Task,
+                _ => createIfMissing ? InterlockedCreateSource() : null,
+            };
+
+            Task<Metadata> InterlockedCreateSource()
+            {
+                var newTcs = new TaskCompletionSource<Metadata>();
+                var existing = Interlocked.CompareExchange(ref _headersTaskOrSource, newTcs, null);
+                return existing switch
+                {
+                    Task<Metadata> task => task,
+                    TaskCompletionSource<Metadata> tcs => tcs.Task,
+                    _ => newTcs.Task,
+                };
+            }
         }
+
         internal Metadata Trailers
         {
             get => _trailers ?? Throw("Trailers are not yet available");
@@ -27,7 +59,8 @@ internal Metadata Trailers
         internal MetadataContext Reset()
         {
             Status = Status.DefaultSuccess;
-            _headers = _trailers = null;
+            _trailers = null;
+            _headersTaskOrSource = null;
             return this;
         }
 
@@ -57,24 +90,34 @@ internal void SetTrailers<T>(T call, Func<T, Status> getStatus, Func<T, Metadata
 
         internal ValueTask SetHeadersAsync(Task<Metadata> headers)
         {
+            var tcs = Interlocked.CompareExchange(ref _headersTaskOrSource, headers, null) as TaskCompletionSource<Metadata>;
             if (headers.RanToCompletion())
             {
-                _headers = headers.Result;
+                // headers are sync; update TCS if one
+                tcs?.TrySetResult(headers.Result);
                 return default;
             }
             else
             {
-                return Awaited(this, headers);
+                // headers are async (or faulted); pay the piper
+                return Awaited(this, tcs, headers);
             }
-            static async ValueTask Awaited(MetadataContext context, Task<Metadata> headers)
+
+            static async ValueTask Awaited(MetadataContext context, TaskCompletionSource<Metadata>? tcs, Task<Metadata> headers)
             {
                 try
                 {
-                    context._headers = await headers.ConfigureAwait(false);
+                    tcs?.TrySetResult(await headers.ConfigureAwait(false));
                 }
                 catch (RpcException fault)
                 {
                     context.SetTrailers(fault);
+                    tcs?.TrySetException(fault);
+                    throw;
+                }
+                catch (Exception ex)
+                {
+                    tcs?.TrySetException(ex);
                     throw;
                 }
             }

tests/protobuf-net.Grpc.Test.Integration/StreamTests.cs

@@ -517,7 +517,7 @@ public async Task DuplexEchoFault(Scenario scenario, int expectedCount, string m
             {
                 await foreach (var item in client.DuplexEcho(For(scenario, DEFAULT_SIZE), ctx))
                 {
-                    CheckHeaderState();
+                    await CheckHeaderStateAsync();
                     values.Add(item.Bar);
                 }
             });
@@ -526,7 +526,7 @@ public async Task DuplexEchoFault(Scenario scenario, int expectedCount, string m
             Assert.Equal(marker + " faultval", rpc.Trailers.GetString("faultkey"));
 
             _fixture?.Log("after await foreach");
-            CheckHeaderState();
+            await CheckHeaderStateAsync();
             Assert.Equal(string.Join(",", Enumerable.Range(0, expectedCount)), string.Join(",", values));
 
             if ((flags & CallContextFlags.CaptureMetadata) != 0)
@@ -538,7 +538,7 @@ public async Task DuplexEchoFault(Scenario scenario, int expectedCount, string m
                 Assert.Equal(marker + " detail", status.Detail);
             }
 
-            void CheckHeaderState()
+            async ValueTask CheckHeaderStateAsync()
             {
                 if (haveCheckedHeaders) return;
                 haveCheckedHeaders = true;
@@ -548,10 +548,10 @@ void CheckHeaderState()
                     switch (scenario)
                     {
                         case Scenario.FaultBeforeHeaders:
-                            Assert.Null(ctx.ResponseHeaders().GetString("prekey"));
+                            Assert.Null((await ctx.ResponseHeadersAsync()).GetString("prekey"));
                             break;
                         default:
-                            Assert.Equal("preval", ctx.ResponseHeaders().GetString("prekey"));
+                            Assert.Equal("preval", (await ctx.ResponseHeadersAsync()).GetString("prekey"));
                             break;
                     }
                 }