[PECOBLR-1190] Make both telemetryClient and factory on a host level (#1075)

## Description
- Minimise the number of OAuth requests from the driver. This PR reduces
it on the featureFlag and telemetry level
- Now that it is reduced, we also change logLevels (Note that even in
3.0.3 we were sending all telemetry logs, the 3.0.4 contains logLevels
in telemetry PR)

## Testing
- Unit tests

## Additional Notes to the Reviewer
<!-- Share any additional context or insights that may help the reviewer
understand the changes better. This could include challenges faced,
limitations, or compromises made during the development process.
Also, mention any areas of the code that you would like the reviewer to
focus on specifically. -->

Author: samikshya-db

NEXT_CHANGELOG.md

@@ -9,6 +9,7 @@
 - Added a new config attribute `DisableOauthRefreshToken` to control whether refresh tokens are requested in OAuth exchanges. By default, the driver does not include the `offline_access` scope. If `offline_access` is explicitly provided by the user, it is preserved and not removed.
 
 ### Updated
+- Minimized OAuth requests by reducing calls in feature flags and telemetry.
 - Updated sdk version from 0.65.0 to 0.69.0
 
 ### Fixed

src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactory.java

@@ -73,6 +73,6 @@ static void setFeatureFlagsContext(
   }
 
   private static String keyOf(IDatabricksConnectionContext context) {
-    return context.getComputeResource().getUniqueIdentifier();
+    return context.getHostForOAuth();
   }
 }

src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java

@@ -8,7 +8,8 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.sdk.core.DatabricksConfig;
 import com.google.common.annotations.VisibleForTesting;
-import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
@@ -23,10 +24,10 @@ public class TelemetryClientFactory {
   private static final TelemetryClientFactory INSTANCE = new TelemetryClientFactory();
 
   @VisibleForTesting
-  final LinkedHashMap<String, TelemetryClient> telemetryClients = new LinkedHashMap<>();
+  final Map<String, TelemetryClientHolder> telemetryClientHolders = new ConcurrentHashMap<>();
 
   @VisibleForTesting
-  final LinkedHashMap<String, TelemetryClient> noauthTelemetryClients = new LinkedHashMap<>();
+  final Map<String, TelemetryClientHolder> noauthTelemetryClientHolders = new ConcurrentHashMap<>();
 
   private final ExecutorService telemetryExecutorService;
 
@@ -59,24 +60,60 @@ public ITelemetryClient getTelemetryClient(IDatabricksConnectionContext connecti
     DatabricksConfig databricksConfig =
         TelemetryHelper.getDatabricksConfigSafely(connectionContext);
     if (databricksConfig != null) {
-      return telemetryClients.computeIfAbsent(
-          connectionContext.getConnectionUuid(),
-          k ->
-              new TelemetryClient(
-                  connectionContext, getTelemetryExecutorService(), databricksConfig));
+      String key = keyOf(connectionContext);
+      TelemetryClientHolder holder =
+          telemetryClientHolders.compute(
+              key,
+              (k, existing) -> {
+                if (existing == null) {
+                  return new TelemetryClientHolder(
+                      new TelemetryClient(
+                          connectionContext, getTelemetryExecutorService(), databricksConfig),
+                      1);
+                }
+                existing.refCount.incrementAndGet();
+                return existing;
+              });
+      return holder.client;
     }
     // Use no-auth telemetry client if connection creation failed.
-    return noauthTelemetryClients.computeIfAbsent(
-        connectionContext.getConnectionUuid(),
-        k -> new TelemetryClient(connectionContext, getTelemetryExecutorService()));
+    String key = keyOf(connectionContext);
+    TelemetryClientHolder holder =
+        noauthTelemetryClientHolders.compute(
+            key,
+            (k, existing) -> {
+              if (existing == null) {
+                return new TelemetryClientHolder(
+                    new TelemetryClient(connectionContext, getTelemetryExecutorService()), 1);
+              }
+              existing.refCount.incrementAndGet();
+              return existing;
+            });
+    return holder.client;
   }
 
   public void closeTelemetryClient(IDatabricksConnectionContext connectionContext) {
-    closeTelemetryClient(
-        telemetryClients.remove(connectionContext.getConnectionUuid()), "telemetry client");
-    closeTelemetryClient(
-        noauthTelemetryClients.remove(connectionContext.getConnectionUuid()),
-        "unauthenticated telemetry client");
+    String key = keyOf(connectionContext);
+    telemetryClientHolders.computeIfPresent(
+        key,
+        (k, holder) -> {
+          if (holder.refCount.get() <= 1) {
+            closeTelemetryClient(holder.client, "telemetry client");
+            return null;
+          }
+          holder.refCount.decrementAndGet();
+          return holder;
+        });
+    noauthTelemetryClientHolders.computeIfPresent(
+        key,
+        (k, holder) -> {
+          if (holder.refCount.get() <= 1) {
+            closeTelemetryClient(holder.client, "unauthenticated telemetry client");
+            return null;
+          }
+          holder.refCount.decrementAndGet();
+          return holder;
+        });
   }
 
   public ExecutorService getTelemetryExecutorService() {
@@ -109,12 +146,12 @@ static ITelemetryPushClient getTelemetryPushClient(
   @VisibleForTesting
   public void reset() {
     // Close all existing clients
-    telemetryClients.values().forEach(TelemetryClient::close);
-    noauthTelemetryClients.values().forEach(TelemetryClient::close);
+    telemetryClientHolders.values().forEach(holder -> holder.client.close());
+    noauthTelemetryClientHolders.values().forEach(holder -> holder.client.close());
 
     // Clear the maps
-    telemetryClients.clear();
-    noauthTelemetryClients.clear();
+    telemetryClientHolders.clear();
+    noauthTelemetryClientHolders.clear();
   }
 
   private void closeTelemetryClient(ITelemetryClient client, String clientType) {
@@ -126,4 +163,18 @@ private void closeTelemetryClient(ITelemetryClient client, String clientType) {
       }
     }
   }
+
+  private static final class TelemetryClientHolder {
+    final TelemetryClient client;
+    final AtomicInteger refCount;
+
+    TelemetryClientHolder(TelemetryClient client, int initialCount) {
+      this.client = client;
+      this.refCount = new AtomicInteger(initialCount);
+    }
+  }
+
+  private static String keyOf(IDatabricksConnectionContext context) {
+    return context.getHostForOAuth();
+  }
 }

src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java

@@ -88,7 +88,7 @@ private static void exportTelemetryEvent(
       TelemetryLogLevel logLevel) {
     if (connectionContext == null
         || telemetryDetails == null
-        || logLevel.toInt() <= connectionContext.getTelemetryLogLevel().toInt()) {
+        || logLevel.toInt() < connectionContext.getTelemetryLogLevel().toInt()) {
       // We don't export telemetry logs in the following three scenarios:
       // 1. When the context is not set.
       // 2. When telemetry details are not set.

src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java

@@ -44,8 +44,8 @@ public void testGetNoOpTelemetryClient() throws Exception {
     ITelemetryClient telemetryClient =
         TelemetryClientFactory.getInstance().getTelemetryClient(context);
     assertInstanceOf(NoopTelemetryClient.class, telemetryClient);
-    assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
-    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
+    assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
+    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
   }
 
   @Test
@@ -58,11 +58,11 @@ public void testGetAuthenticatedTelemetryClient() throws Exception {
     ITelemetryClient telemetryClient =
         TelemetryClientFactory.getInstance().getTelemetryClient(context);
     assertInstanceOf(TelemetryClient.class, telemetryClient);
-    assertEquals(1, TelemetryClientFactory.getInstance().telemetryClients.size());
-    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
+    assertEquals(1, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
+    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
     TelemetryClientFactory.getInstance().closeTelemetryClient(context);
-    assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
-    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
+    assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
+    assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
     TelemetryClientFactory.getInstance().closeTelemetryClient(context);
   }
 
@@ -80,8 +80,8 @@ public void testGetNoOpTelemetryClientWhenDatabricksConfigIsNull() throws Except
           TelemetryClientFactory.getInstance().getTelemetryClient(context);
 
       assertInstanceOf(NoopTelemetryClient.class, telemetryClient);
-      assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
-      assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
+      assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
+      assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
       TelemetryClientFactory.getInstance().closeTelemetryClient(context);
     }
   }

src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java

@@ -205,6 +205,7 @@ public void testGetDatabricksConfigSafely_HandlesNullContext() {
   public void testTelemetryNotAllowedUsecase() {
     // Clear thread context to ensure telemetry is not allowed
     when(connectionContext.forceEnableTelemetry()).thenReturn(false);
+    when(connectionContext.getHostForOAuth()).thenReturn("test-host");
     when(connectionContext.isTelemetryEnabled()).thenReturn(false);
     assertFalse(isTelemetryAllowedForConnection(connectionContext));
     when(connectionContext.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE);
@@ -216,6 +217,7 @@ public void testTelemetryNotAllowedUsecase() {
   public void testTelemetryAllowedWithForceTelemetryFlag() {
     when(connectionContext.getTelemetryLogLevel()).thenReturn(TelemetryLogLevel.DEBUG);
     when(connectionContext.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE);
+    when(connectionContext.getHostForOAuth()).thenReturn("test-host");
     enableFeatureFlagForTesting(connectionContext, Collections.emptyMap());
     assertTrue(() -> isTelemetryAllowedForConnection(connectionContext));
   }
@@ -260,6 +262,29 @@ void testExportTelemetryLog_EmitsWhenEventLevelHigherThanConfigured() {
     }
   }
 
+  @Test
+  void testExportTelemetryLog_EmitsWhenEventLevelEqualToConfigured() {
+    // Configured level: INFO (4); Event level: INFO (4) -> should export (4 >= 4)
+    when(connectionContext.getTelemetryLogLevel()).thenReturn(TelemetryLogLevel.INFO);
+    StatementTelemetryDetails details =
+        new StatementTelemetryDetails("stmt-eq").setOperationLatencyMillis(15L);
+
+    ITelemetryClient clientMock = mock(ITelemetryClient.class);
+    TelemetryClientFactory factoryMock = mock(TelemetryClientFactory.class);
+
+    try (MockedStatic<TelemetryClientFactory> mocked =
+        Mockito.mockStatic(TelemetryClientFactory.class)) {
+      mocked.when(TelemetryClientFactory::getInstance).thenReturn(factoryMock);
+      when(factoryMock.getTelemetryClient(connectionContext)).thenReturn(clientMock);
+
+      TelemetryHelper.exportTelemetryLog(details, TelemetryLogLevel.INFO);
+
+      mocked.verify(TelemetryClientFactory::getInstance, times(1));
+      verify(factoryMock, times(1)).getTelemetryClient(connectionContext);
+      verify(clientMock, times(1)).exportEvent(any());
+    }
+  }
+
   static Stream<Object[]> failureLogParameters() {
     return Stream.of(
         new Object[] {"test-statement-id", null},