Code cleanup

newrelic · Nov 18, 2024 · de3a017 · de3a017
1 parent badab6a
commit de3a017
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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);
         }
     }
 }