From 8165b65a4dfebe7836db8bca1e1238e5ca871fa6 Mon Sep 17 00:00:00 2001
From: Mitchell Cookson <31948306+Cooksauce@users.noreply.github.com>
Date: Mon, 29 Jan 2024 05:15:44 -0600
Subject: [PATCH] Add OpenTelemetry extension

---
 .../NATS.Net.OpenTelemetry.csproj             |  8 +++++++
 .../OpenTelemetryTracingExtensions.cs         | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 src/NATS.Net.OpenTelemetry/OpenTelemetryTracingExtensions.cs

diff --git a/src/NATS.Net.OpenTelemetry/NATS.Net.OpenTelemetry.csproj b/src/NATS.Net.OpenTelemetry/NATS.Net.OpenTelemetry.csproj
index f5ae8bf9b..89f1cf38b 100644
--- a/src/NATS.Net.OpenTelemetry/NATS.Net.OpenTelemetry.csproj
+++ b/src/NATS.Net.OpenTelemetry/NATS.Net.OpenTelemetry.csproj
@@ -16,4 +16,12 @@
 
     </PropertyGroup>
 
+    <ItemGroup>
+        <PackageReference Include="OpenTelemetry" Version="1.6.0" />
+    </ItemGroup>
+
+    <ItemGroup>
+      <ProjectReference Include="..\NATS.Client.Core\NATS.Client.Core.csproj" />
+    </ItemGroup>
+
 </Project>
diff --git a/src/NATS.Net.OpenTelemetry/OpenTelemetryTracingExtensions.cs b/src/NATS.Net.OpenTelemetry/OpenTelemetryTracingExtensions.cs
new file mode 100644
index 000000000..3e64c7c0c
--- /dev/null
+++ b/src/NATS.Net.OpenTelemetry/OpenTelemetryTracingExtensions.cs
@@ -0,0 +1,23 @@
+using NATS.Client.Core;
+
+// ReSharper disable once CheckNamespace - Place in OTEL namespace for discoverability
+namespace OpenTelemetry.Trace;
+
+public static class OpenTelemetryTracingExtensions
+{
+    /// <summary>
+    /// Enables trace collection on activity sources from the NATS.Client nuget package.
+    /// </summary>
+    /// <param name="builder"><see cref="TracerProviderBuilder"/> being configured.</param>
+    /// <param name="includeInternal">Include traces from internal messaging used for control flow and lower level usage during high level abstractions (e.g. JetStream, KvStore, etc)</param>
+    /// <returns>The same instance of <see cref="TracerProviderBuilder"/> to chain the calls.</returns>
+    public static TracerProviderBuilder AddNatsInstrumentation(this TracerProviderBuilder builder, bool includeInternal = false)
+    {
+        builder.AddSource("NATS.Client");
+
+        if (includeInternal)
+            builder.AddSource("NATS.Client.Internal");
+
+        return builder;
+    }
+}