From af58ff9db27b8720f8cf950f6935dc1d2e24c1d8 Mon Sep 17 00:00:00 2001
From: Andrew Lock <andrew.lock@datadoghq.com>
Date: Fri, 20 Dec 2024 16:05:49 +0000
Subject: [PATCH] Fix BasicPublish for <.NET Core 3.1

---
 .../BasicPublishAsyncCachedStringsIntegration.cs     | 12 ++++++++++++
 .../RabbitMQ/BasicPublishAsyncIntegration.cs         | 12 ++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncCachedStringsIntegration.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncCachedStringsIntegration.cs
index 31aece78af5a..b39256d17fd8 100644
--- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncCachedStringsIntegration.cs
+++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncCachedStringsIntegration.cs
@@ -41,11 +41,23 @@ internal static CallTargetState OnMethodBegin<TTarget, TExchange, TRoutingKey, T
         return BasicPublishAsyncIntegration.OnMethodBegin(instance, exchangeString, routingKeyString, mandatory, basicProperties, body, default);
     }
 
+#if NETCOREAPP3_1_OR_GREATER
+    // We don't support ValueTask in < .NET Core 3.1, which means this doesn't work and is never called
     internal static TReturn? OnAsyncMethodEnd<TTarget, TReturn>(TTarget instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
     {
         state.Scope.DisposeWithException(exception);
         return returnValue;
     }
+#else
+    // If the method executes synchronously, then this works fine, but if it executes asynchronously we could miss the exception
+    // This is because the exception is thrown after the method returns, and we can't catch it here
+    // If/when we support handling ValueTask correctly in < .NET Core 3.1, we should remove this method
+    internal static CallTargetReturn<TReturn?> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
+    {
+        state.Scope.DisposeWithException(exception);
+        return new CallTargetReturn<TReturn?>(returnValue);
+    }
+#endif
 }
 
 /// <summary>
diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncIntegration.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncIntegration.cs
index 362972fada11..e2717885d6d9 100644
--- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncIntegration.cs
+++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/BasicPublishAsyncIntegration.cs
@@ -65,9 +65,21 @@ internal static CallTargetState OnMethodBegin<TTarget, TBasicProperties, TBody>(
         return new CallTargetState(scope);
     }
 
+#if NETCOREAPP3_1_OR_GREATER
+    // We don't support ValueTask in < .NET Core 3.1, which means this doesn't work and is never called
     internal static TReturn? OnAsyncMethodEnd<TTarget, TReturn>(TTarget instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
     {
         state.Scope.DisposeWithException(exception);
         return returnValue;
     }
+#else
+    // If the method executes synchronously, then this works fine, but if it executes asynchronously we could miss the exception
+    // This is because the exception is thrown after the method returns, and we can't catch it here
+    // If/when we support handling ValueTask correctly in < .NET Core 3.1, we should remove this method
+    internal static CallTargetReturn<TReturn?> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
+    {
+        state.Scope.DisposeWithException(exception);
+        return new CallTargetReturn<TReturn?>(returnValue);
+    }
+#endif
 }