From 8cc9a85134ef93ee17accc74bf84835873e8c71f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Bylund?= <bjorn.bylund@bookbeat.com>
Date: Tue, 2 Jan 2024 08:40:29 +0100
Subject: [PATCH] Azure service bus new relic distributed tracing headers

---
 .../KnightBus.Azure.ServiceBus.csproj                       | 1 +
 .../src/KnightBus.Azure.ServiceBus/ServiceBus.cs            | 6 ++++++
 .../src/KnightBus.NewRelic/NewRelicMessageMiddleware.cs     | 4 ++++
 3 files changed, 11 insertions(+)

diff --git a/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/KnightBus.Azure.ServiceBus.csproj b/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/KnightBus.Azure.ServiceBus.csproj
index e35180c2..b96dc12d 100644
--- a/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/KnightBus.Azure.ServiceBus.csproj
+++ b/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/KnightBus.Azure.ServiceBus.csproj
@@ -17,6 +17,7 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Messaging.ServiceBus" Version="7.13.1" />
+    <PackageReference Include="NewRelic.Agent.Api" Version="10.2.0" />
   </ItemGroup>
   <ItemGroup>
     <None Include="knighbus-64.png" Pack="true" Visible="false" PackagePath="" />
diff --git a/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/ServiceBus.cs b/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/ServiceBus.cs
index 07c07168..e6e63ad2 100644
--- a/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/ServiceBus.cs
+++ b/knightbus-azureservicebus/src/KnightBus.Azure.ServiceBus/ServiceBus.cs
@@ -178,6 +178,12 @@ private async Task<ServiceBusMessage> CreateMessageAsync<T>(T body) where T : IM
                 }
             }
 
+            NewRelic.Api.Agent.NewRelic.GetAgent().CurrentTransaction.InsertDistributedTraceHeaders(message.ApplicationProperties,
+                ((carrier, key, value) =>
+                {
+                    carrier.Add(key, value);
+                }));
+
             return message;
         }
 
diff --git a/knightbus-newrelic/src/KnightBus.NewRelic/NewRelicMessageMiddleware.cs b/knightbus-newrelic/src/KnightBus.NewRelic/NewRelicMessageMiddleware.cs
index c41dcbfe..cdb78add 100644
--- a/knightbus-newrelic/src/KnightBus.NewRelic/NewRelicMessageMiddleware.cs
+++ b/knightbus-newrelic/src/KnightBus.NewRelic/NewRelicMessageMiddleware.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using KnightBus.Core;
@@ -22,6 +23,9 @@ public async Task ProcessAsync<T>(IMessageStateHandler<T> messageStateHandler, I
             var messageName = typeof(T).FullName;
             NewRelic.Api.Agent.NewRelic.SetTransactionName("Message", messageName);
             var transaction = _agent.CurrentTransaction;
+            transaction.AcceptDistributedTraceHeaders(messageStateHandler.MessageProperties,
+                ((carrier, key) => carrier.TryGetValue(key, out var value) ? new[] { value } : null),
+                TransportType.AMQP);
             try
             {
                 await next.ProcessAsync(messageStateHandler, cancellationToken).ConfigureAwait(false);