From 98b93472513fb051f5a3dd349cdaf561f99a754a Mon Sep 17 00:00:00 2001
From: Marty Tippin <120425148+tippmar-nr@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:25:44 -0500
Subject: [PATCH] fix: Kafka bootstrap cache wasn't implemented correctly

---
 .../Wrapper/Kafka/KafkaBuilderWrapper.cs      | 24 +++++++------------
 1 file changed, 8 insertions(+), 16 deletions(-)

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Kafka/KafkaBuilderWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Kafka/KafkaBuilderWrapper.cs
index 2dc440fcaf..b4948ebf8c 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Kafka/KafkaBuilderWrapper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Kafka/KafkaBuilderWrapper.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections;
+using System.Collections.Concurrent;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 using NewRelic.Reflection;
@@ -11,8 +12,7 @@ namespace NewRelic.Providers.Wrapper.Kafka
 {
     public class KafkaBuilderWrapper : IWrapper
     {
-        private Func<object, IEnumerable> _producerBuilderConfigGetter;
-        private Func<object, IEnumerable> _consumerBuilderConfigGetter;
+        private ConcurrentDictionary<Type, Func<object, IEnumerable>> _builderConfigGetterDictionary = new();
 
         private const string WrapperName = "KafkaBuilderWrapper";
         private const string BootstrapServersKey = "bootstrap.servers";
@@ -27,21 +27,13 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
         {
             var builder = instrumentedMethodCall.MethodCall.InvocationTarget;
 
-            dynamic configuration = null;
-
-            if (builder.GetType().Name == "ProducerBuilder`2")
-            {
-                var configGetter = _producerBuilderConfigGetter ??= VisibilityBypasser.Instance.GeneratePropertyAccessor<IEnumerable>(builder.GetType(), "Config");
-                configuration = configGetter(builder) as dynamic;
-            }
-            else if (builder.GetType().Name == "ConsumerBuilder`2")
+            if (!_builderConfigGetterDictionary.TryGetValue(builder.GetType(), out var configGetter))
             {
-                var configGetter = _consumerBuilderConfigGetter ??= VisibilityBypasser.Instance.GeneratePropertyAccessor<IEnumerable>(builder.GetType(), "Config");
-                configuration = configGetter(builder) as dynamic;
+                configGetter = VisibilityBypasser.Instance.GeneratePropertyAccessor<IEnumerable>(builder.GetType(), "Config");
+                _builderConfigGetterDictionary[builder.GetType()] = configGetter;
             }
 
-            if (configuration == null)
-                return Delegates.NoOp;
+            dynamic configuration = configGetter(builder);
 
             string bootstrapServers = null;
 
@@ -55,9 +47,9 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             }
 
             if (!string.IsNullOrEmpty(bootstrapServers))
-                return Delegates.GetDelegateFor<object>(onSuccess: (producerOrConsumerAsObject) =>
+                return Delegates.GetDelegateFor<object>(onSuccess: (builtObject) =>
                 {
-                    KafkaHelper.AddBootstrapServersToCache(producerOrConsumerAsObject, bootstrapServers);
+                    KafkaHelper.AddBootstrapServersToCache(builtObject, bootstrapServers);
                 });
 
             return Delegates.NoOp;