Use IHttpActivityFeature to get host activity (#1545)

Co-authored-by: Safia Abdalla <[email protected]>

Author: JamesNK

src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs

@@ -400,6 +400,16 @@ public void Initialize(ISystemClock? clock = null)
 
         private Activity? GetHostActivity()
         {
+#if NET6_0_OR_GREATER
+            // Feature always returns the host activity
+            var feature = HttpContext.Features.Get<IHttpActivityFeature>();
+            if (feature != null)
+            {
+                return feature.Activity;
+            }
+#endif
+
+            // If feature isn't available, or not supported, then fallback to Activity.Current.
             var activity = Activity.Current;
             while (activity != null)
             {

test/FunctionalTests/Server/DiagnosticsTests.cs

@@ -0,0 +1,59 @@
+#region Copyright notice and license
+
+// Copyright 2019 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+using System.Net;
+using Greet;
+using Grpc.AspNetCore.FunctionalTests.Infrastructure;
+using Grpc.Tests.Shared;
+using NUnit.Framework;
+
+namespace Grpc.AspNetCore.FunctionalTests.Server
+{
+    [TestFixture]
+    public class DiagnosticsTests : FunctionalTestBase
+    {
+        [Test]
+        public async Task HostActivityTags_ReturnedInTrailers_Success()
+        {
+            // Arrange
+            var requestMessage = new HelloRequest
+            {
+                Name = "World"
+            };
+
+            var ms = new MemoryStream();
+            MessageHelpers.WriteMessage(ms, requestMessage);
+
+            var httpRequest = GrpcHttpHelper.Create("Greet.Greeter/SayHello");
+            httpRequest.Content = new GrpcStreamContent(ms);
+            httpRequest.Headers.Add("return-tags-trailers", "true");
+
+            // Act
+            var response = await Fixture.Client.SendAsync(httpRequest).DefaultTimeout();
+
+            // Assert
+            var responseMessage = await response.GetSuccessfulGrpcMessageAsync<HelloReply>().DefaultTimeout();
+            Assert.AreEqual("Hello World", responseMessage.Message);
+            response.AssertTrailerStatus();
+
+            var trailers = response.TrailingHeaders;
+            Assert.AreEqual("0", trailers.GetValues("grpc.status_code").Single());
+            Assert.AreEqual("/Greet.Greeter/SayHello", trailers.GetValues("grpc.method").Single());
+        }
+    }
+}

testassets/FunctionalTestsWebsite/FunctionalTestsWebsite.csproj

@@ -29,6 +29,8 @@
     <Protobuf Include="..\Proto\echo.proto" Link="Protos\echo.proto" />
     <Protobuf Include="..\Proto\issue.proto" Link="Protos\issue.proto" />
     <Protobuf Include="..\Proto\test.proto" Link="Protos\test.proto" />
+
+	<Compile Include="..\..\test\Shared\ActivityReplacer.cs" Link="Infrastructure\ActivityReplacer.cs" />
   </ItemGroup>
 
 </Project>

testassets/FunctionalTestsWebsite/Startup.cs

@@ -16,12 +16,14 @@
 
 #endregion
 
+using System.Diagnostics;
 using System.IdentityModel.Tokens.Jwt;
 using System.Security.Claims;
 using FunctionalTestsWebsite.Infrastructure;
 using FunctionalTestsWebsite.Services;
 using Greet;
 using Grpc.AspNetCore.Server.Model;
+using Grpc.Tests.Shared;
 using Microsoft.AspNetCore.Authentication.JwtBearer;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.IdentityModel.Tokens;
@@ -132,6 +134,33 @@ public void Configure(IApplicationBuilder app)
                 await next();
             });
 
+            app.Use(async (context, next) =>
+            {
+                // Allow a call to specify activity tags are returned as trailers.
+                if (context.Request.Headers.TryGetValue("return-tags-trailers", out var value) &&
+                    bool.TryParse(value.ToString(), out var remove) && remove)
+                {
+                    logger.LogInformation("Replacing activity.");
+
+                    // Replace the activity to check that tags are added to the host activity.
+                    using (new ActivityReplacer("GrpcFunctionalTests"))
+                    {
+                        await next();
+                    }
+
+                    logger.LogInformation("Adding tags to trailers.");
+
+                    foreach (var tag in Activity.Current!.Tags)
+                    {
+                        context.Response.AppendTrailer(tag.Key, tag.Value);
+                    }
+                }
+                else
+                {
+                    await next();
+                }
+            });
+
             app.UseRouting();
 
             app.UseAuthorization();