From de3a017746163663eed6e85e901e94d03a9b5584 Mon Sep 17 00:00:00 2001
From: Marty Tippin <120425148+tippmar-nr@users.noreply.github.com>
Date: Mon, 18 Nov 2024 16:12:47 -0600
Subject: [PATCH] Code cleanup

---
 .../Wrapper/RabbitMq/BasicPublishWrapper.cs   |  8 +++--
 .../RabbitMq/HandleBasicDeliverWrapper.cs     | 30 ++++++++++++++-----
 .../Wrapper/RabbitMq/Instrumentation.xml      |  2 +-
 .../Wrapper/RabbitMq/QueuePurgeWrapper.cs     |  3 +-
 4 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/BasicPublishWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/BasicPublishWrapper.cs
index 92f836f11..0a493e289 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/BasicPublishWrapper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/BasicPublishWrapper.cs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 using System;
+using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 
@@ -40,9 +41,10 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
                 RabbitMqHelper.InsertDTHeaders(instrumentedMethodCall, transaction, BasicPropertiesIndex);
 
 
-            // TODO: probably need to do something special for v7 since the return type is ValueTask
-            //return Delegates.GetDelegateFor(segment);
-            return Delegates.NoOp;
+            // TODO: probably need to do something special for v7 since the return type is ValueTask<T>
+            return instrumentedMethodCall.IsAsync ?
+                Delegates.GetAsyncDelegateFor<Task>(agent, segment)
+              : Delegates.GetDelegateFor(segment);
         }
     }
 }
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/HandleBasicDeliverWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/HandleBasicDeliverWrapper.cs
index ea221d3c0..d551b349d 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/HandleBasicDeliverWrapper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/HandleBasicDeliverWrapper.cs
@@ -5,6 +5,7 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 using NewRelic.Agent.Extensions.SystemExtensions;
@@ -70,13 +71,28 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
                 serverPort: port,
                 routingKey: routingKey);
 
-            return Delegates.GetDelegateFor(
-                onFailure: transaction.NoticeError,
-                onComplete: () =>
-                {
-                    segment.End();
-                    transaction.End();
-                });
+            return instrumentedMethodCall.IsAsync
+                ? Delegates.GetAsyncDelegateFor<Task>(
+                    agent,
+                    segment,
+                    false,
+                    onComplete: (t) =>
+                    {
+                        if (t.IsFaulted)
+                        {
+                            transaction.NoticeError(t.Exception);
+                        }
+
+                        segment.End();
+                        transaction.End();
+                    })
+                : Delegates.GetDelegateFor(
+                    onFailure: transaction.NoticeError,
+                    onComplete: () =>
+                    {
+                        segment.End();
+                        transaction.End();
+                    });
 
             IEnumerable<string> GetHeaderValue(IDictionary<string, object> carrier, string key)
             {
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/Instrumentation.xml b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/Instrumentation.xml
index 53c2dcd32..0c7c5f630 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/Instrumentation.xml
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/Instrumentation.xml
@@ -15,7 +15,7 @@ SPDX-License-Identifier: Apache-2.0
          RabbitMQ v7+
          public async Task<BasicGetResult?> BasicGetAsync(string queue, bool autoAck, CancellationToken cancellationToken)
       -->
-      <match assemblyName="RabbitMQ.Client" className="RabbitMQ.Client.Channel" >
+      <match assemblyName="RabbitMQ.Client" className="RabbitMQ.Client.Impl.Channel" >
         <exactMethodMatcher methodName="BasicGetAsync" parameters="System.String,System.Boolean,System.Threading.CancellationToken"/>
       </match>
     </tracerFactory>
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/QueuePurgeWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/QueuePurgeWrapper.cs
index 1ce9437a4..8247d9b03 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/QueuePurgeWrapper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RabbitMq/QueuePurgeWrapper.cs
@@ -1,6 +1,7 @@
 // Copyright 2020 New Relic, Inc. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 using NewRelic.Agent.Extensions.SystemExtensions;
@@ -38,7 +39,7 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             // Routing key is not available for this method.
             // It only returns uint and invocationTarget does not have the value.
 
-            return Delegates.GetDelegateFor(segment);
+            return instrumentedMethodCall.IsAsync ? Delegates.GetAsyncDelegateFor<Task>(agent, segment) : Delegates.GetDelegateFor(segment);
         }
     }
 }