Improvements to telemetry implementation (#1009)

## Description
- Improved the cache logic for feature flags 
- Minor latency improvement for telemetry collector (3-4 ms improvement
on hot paths). This is by skipping allocation when present and skipping
map fetch when new object is created.
- Moved system config info to be logged at connection start-up
asynchrounously (Note : we can't move it to static as we need it in our
log folder -i.e., we can add it only after logging is setup)
- For feature flag factory, we preserve the object instance per compute
- This will improve performance in case of multi-connection. Currently,
we clear out for every compute - i.e., make a call to the feature flag
endpoint on every connectionContext.
- Force flush synchronously so that telemetry logs are pushed during
connection close itself.

### For benchmarking improvements, there is ~30 ms improvement with
these changes.[ see the internal
table](https://docs.google.com/document/d/1oAu1pTCM9VFrPDCt0I_aWsSIoNGOapmGc_2Zob6QUiA/edit?tab=t.0#heading=h.ibcl7vrboyf2)
for more detail

## Testing
- Verified locally that system information is being logged correctly
`2025-09-17 12:08:27 INFO
com.databricks.client.jdbc.Driver#lambda$connect$0 -
DriverSystemConfiguration[driverName='oss-jdbc',
driverVersion='1.0.9-oss', osName='Mac OS X', osVersion='15.6.1',
osArch='aarch64', runtimeName='Java HotSpot(TM) 64-Bit Server VM',
runtimeVersion='21.0.7', runtimeVendor='Oracle Corporation',
localeName='en_US', processName='JUnitStarter',
defaultCharsetEncoding='UTF-8']
`
- Even if i close the JVM after conn.close(), the logs show up. This was
not the case earlier

## 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

@@ -5,6 +5,7 @@
 ### Added
 
 ### Updated
+- Telemetry data is now captured more efficiently and consistently due to enhancements in the log and connection close flush logic.
 
 ### Fixed
 - Fixed state leaking issue in thrift client.

src/main/java/com/databricks/client/jdbc/Driver.java

@@ -22,10 +22,11 @@
 import java.util.List;
 import java.util.Properties;
 import java.util.TimeZone;
+import java.util.concurrent.CompletableFuture;
 import java.util.logging.Logger;
 
 /** Databricks JDBC driver. */
-public class Driver implements IDatabricksDriver, java.sql.Driver {
+public class Driver implements IDatabricksDriver {
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(Driver.class);
   private static final Driver INSTANCE;
 
@@ -57,8 +58,8 @@ public Connection connect(String url, Properties info) throws DatabricksSQLExcep
     IDatabricksConnectionContext connectionContext =
         DatabricksConnectionContextFactory.create(url, info);
     DriverUtil.setUpLogging(connectionContext);
+    CompletableFuture.runAsync(() -> LOGGER.info(getDriverSystemConfiguration().toString()));
     UserAgentManager.setUserAgent(connectionContext);
-    LOGGER.info(getDriverSystemConfiguration().toString());
     DatabricksConnection connection = new DatabricksConnection(connectionContext);
     boolean isConnectionOpen = false;
     try {

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java

@@ -156,9 +156,9 @@ public void close() throws DatabricksSQLException {
     }
     this.session.close();
     TelemetryClientFactory.getInstance().closeTelemetryClient(connectionContext);
-    DatabricksHttpClientFactory.getInstance().removeClient(connectionContext);
     DatabricksClientConfiguratorManager.getInstance().removeInstance(connectionContext);
     DatabricksDriverFeatureFlagsContextFactory.removeInstance(connectionContext);
+    DatabricksHttpClientFactory.getInstance().removeClient(connectionContext);
     DatabricksThreadContextHolder.clearAllContext();
   }
 

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

@@ -796,7 +796,7 @@ public int getTelemetryBatchSize() {
 
   @Override
   public boolean isTelemetryEnabled() {
-    return getParameter(DatabricksJdbcUrlParams.ENABLE_TELEMETRY, "0").equals("1");
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_TELEMETRY).equals("1");
   }
 
   @Override

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

@@ -1,7 +1,5 @@
 package com.databricks.jdbc.common.safe;
 
-import static java.lang.Math.max;
-
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.DatabricksClientConfiguratorManager;
 import com.databricks.jdbc.common.util.DriverUtil;
@@ -12,13 +10,13 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.cache.*;
-import com.google.common.util.concurrent.Futures;
-import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import java.io.IOException;
 import java.util.Map;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpGet;
@@ -33,56 +31,58 @@ public class DatabricksDriverFeatureFlagsContext {
           "/api/2.0/connector-service/feature-flags/OSS_JDBC/%s",
           DriverUtil.getDriverVersionWithoutOSSSuffix());
   private static final int DEFAULT_TTL_SECONDS = 900; // 15 minutes
-  private static final int REFRESH_BEFORE_EXPIRY_SECONDS = 10; // refresh 10s before expiry
   private final String featureFlagEndpoint;
   private final IDatabricksConnectionContext connectionContext;
-  private LoadingCache<String, String> featureFlags;
-  private final ExecutorService asyncExecutor = Executors.newCachedThreadPool();
+  private final Cache<String, String> featureFlags;
+  private final ScheduledExecutorService scheduler =
+      Executors.newSingleThreadScheduledExecutor(
+          r -> {
+            Thread t = new Thread(r, "databricks-jdbc-feature-flags-refresh");
+            t.setDaemon(true);
+            return t;
+          });
+  private ScheduledFuture<?> scheduledRefreshTask;
+  private volatile int refreshIntervalSeconds = DEFAULT_TTL_SECONDS;
 
   public DatabricksDriverFeatureFlagsContext(IDatabricksConnectionContext connectionContext) {
     this.connectionContext = connectionContext;
-    this.featureFlags = createFeatureFlagsCache(DEFAULT_TTL_SECONDS);
+    this.featureFlags = CacheBuilder.newBuilder().build();
     this.featureFlagEndpoint =
         String.format(
             "https://%s%s", connectionContext.getHostForOAuth(), FEATURE_FLAGS_ENDPOINT_SUFFIX);
+    // Make an initial blocking call to fetch featureFlags
+    refreshAllFeatureFlags();
+    // Async fetch eventually
+    scheduleOrRescheduleRefresh(DEFAULT_TTL_SECONDS);
   }
 
   // Constructor for testing
   DatabricksDriverFeatureFlagsContext(
       IDatabricksConnectionContext connectionContext, Map<String, String> initialFlags) {
     this.connectionContext = connectionContext;
-    this.featureFlags = createFeatureFlagsCache(DEFAULT_TTL_SECONDS);
+    this.featureFlags = CacheBuilder.newBuilder().build();
     this.featureFlagEndpoint =
         String.format(
             "https://%s%s", connectionContext.getHostForOAuth(), FEATURE_FLAGS_ENDPOINT_SUFFIX);
     initialFlags.forEach(this.featureFlags::put);
+    scheduleOrRescheduleRefresh(DEFAULT_TTL_SECONDS);
   }
 
-  private LoadingCache<String, String> createFeatureFlagsCache(int ttlSeconds) {
-    return CacheBuilder.newBuilder()
-        .expireAfterWrite(ttlSeconds, TimeUnit.SECONDS)
-        .refreshAfterWrite(
-            max(
-                300,
-                ttlSeconds
-                    - REFRESH_BEFORE_EXPIRY_SECONDS), // refresh time should be minimum 5 minutes
-            TimeUnit.SECONDS)
-        .build(
-            new CacheLoader<>() {
-              @Override
-              public String load(String key) {
-                refreshAllFeatureFlags();
-                return featureFlags.getIfPresent(key) != null
-                    ? featureFlags.getIfPresent(key)
-                    : "false";
-              }
-
-              @Override
-              public ListenableFuture<String> reload(String key, String oldValue) {
-                asyncExecutor.submit(() -> refreshAllFeatureFlags());
-                return Futures.immediateFuture(oldValue); // keep old value until refresh is done
-              }
-            });
+  private void scheduleOrRescheduleRefresh(int ttlSeconds) {
+    this.refreshIntervalSeconds = ttlSeconds > 0 ? ttlSeconds : DEFAULT_TTL_SECONDS;
+    if (scheduler.isShutdown()) {
+      return;
+    }
+    if (scheduledRefreshTask != null && !scheduledRefreshTask.isCancelled()) {
+      scheduledRefreshTask.cancel(false);
+    }
+    // Schedule refresh at a fixed rate.
+    scheduledRefreshTask =
+        scheduler.scheduleAtFixedRate(
+            this::refreshAllFeatureFlags,
+            this.refreshIntervalSeconds,
+            this.refreshIntervalSeconds,
+            TimeUnit.SECONDS);
   }
 
   private void refreshAllFeatureFlags() {
@@ -121,7 +121,7 @@ void fetchAndSetFlagsFromServer(IDatabricksHttpClient httpClient, HttpGet reques
 
         Integer ttlSeconds = featureFlagsResponse.getTtlSeconds();
         if (ttlSeconds != null) {
-          featureFlags = createFeatureFlagsCache(ttlSeconds);
+          scheduleOrRescheduleRefresh(ttlSeconds);
         }
       } else {
         LOGGER.trace(
@@ -133,11 +133,23 @@ void fetchAndSetFlagsFromServer(IDatabricksHttpClient httpClient, HttpGet reques
   }
 
   public boolean isFeatureEnabled(String name) {
+    String value = featureFlags.getIfPresent(name);
+    return Boolean.parseBoolean(value);
+  }
+
+  public void shutdown() {
+    ScheduledFuture<?> task = scheduledRefreshTask;
+    if (task != null) {
+      task.cancel(false);
+    }
+    scheduler.shutdown();
     try {
-      return Boolean.parseBoolean(featureFlags.get(name));
-    } catch (Exception e) {
-      LOGGER.trace("Error fetching feature flag {}: {}", name, e.getMessage());
-      return false;
+      if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
+        scheduler.shutdownNow();
+      }
+    } catch (InterruptedException ie) {
+      Thread.currentThread().interrupt();
+      scheduler.shutdownNow();
     }
   }
 }

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

@@ -7,7 +7,7 @@
 
 /** Factory class to manage DatabricksDriverFeatureFlagsContext instances */
 public class DatabricksDriverFeatureFlagsContextFactory {
-  private static final Map<String, DatabricksDriverFeatureFlagsContext> contextMap =
+  private static final Map<String, FeatureFlagsContextHolder> contextMap =
       new ConcurrentHashMap<>();
 
   private DatabricksDriverFeatureFlagsContextFactory() {
@@ -22,9 +22,21 @@ private DatabricksDriverFeatureFlagsContextFactory() {
    */
   public static DatabricksDriverFeatureFlagsContext getInstance(
       IDatabricksConnectionContext context) {
-    return contextMap.computeIfAbsent(
-        context.getComputeResource().getUniqueIdentifier(),
-        k -> new DatabricksDriverFeatureFlagsContext(context));
+    String key = keyOf(context);
+    FeatureFlagsContextHolder holder =
+        contextMap.compute(
+            key,
+            (k, existing) -> {
+              if (existing == null) {
+                // First reference for this compute
+                return new FeatureFlagsContextHolder(
+                    new DatabricksDriverFeatureFlagsContext(context), 1);
+              }
+              // Additional reference for the same compute
+              existing.refCount.incrementAndGet();
+              return existing;
+            });
+    return holder.context;
   }
 
   /**
@@ -34,15 +46,33 @@ public static DatabricksDriverFeatureFlagsContext getInstance(
    */
   public static void removeInstance(IDatabricksConnectionContext connectionContext) {
     if (connectionContext != null) {
-      contextMap.remove(connectionContext.getComputeResource().getUniqueIdentifier());
+      String key = keyOf(connectionContext);
+      contextMap.computeIfPresent(
+          key,
+          (k, holder) -> {
+            // Last reference being removed: shutdown and remove entry
+            if (holder.refCount.get() <= 1) {
+              holder.context.shutdown();
+              return null;
+            }
+            // Still referenced elsewhere: just decrement
+            holder.refCount.decrementAndGet();
+            return holder;
+          });
     }
   }
 
   @VisibleForTesting
   static void setFeatureFlagsContext(
       IDatabricksConnectionContext connectionContext, Map<String, String> featureFlags) {
+    String key = keyOf(connectionContext);
     contextMap.put(
-        connectionContext.getComputeResource().getUniqueIdentifier(),
-        new DatabricksDriverFeatureFlagsContext(connectionContext, featureFlags));
+        key,
+        new FeatureFlagsContextHolder(
+            new DatabricksDriverFeatureFlagsContext(connectionContext, featureFlags), 1));
+  }
+
+  private static String keyOf(IDatabricksConnectionContext context) {
+    return context.getComputeResource().getUniqueIdentifier();
   }
 }

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

@@ -0,0 +1,13 @@
+package com.databricks.jdbc.common.safe;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+final class FeatureFlagsContextHolder {
+  final DatabricksDriverFeatureFlagsContext context;
+  AtomicInteger refCount;
+
+  FeatureFlagsContextHolder(DatabricksDriverFeatureFlagsContext context, int refCount) {
+    this.context = context;
+    this.refCount = new AtomicInteger(refCount);
+  }
+}

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

@@ -22,7 +22,6 @@
 import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.dbclient.impl.sqlexec.CommandBuilder;
-import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
@@ -48,7 +47,7 @@ public class DatabricksThriftServiceClient implements IDatabricksClient, IDatabr
   private final MetadataResultSetBuilder metadataResultSetBuilder;
 
   public DatabricksThriftServiceClient(IDatabricksConnectionContext connectionContext)
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException {
     this.connectionContext = connectionContext;
     this.thriftAccessor = new DatabricksThriftAccessor(connectionContext);
     this.metadataResultSetBuilder = new MetadataResultSetBuilder(connectionContext);