[PECOBLR-591] Add chunk latency telemetry (#877)

samikshya-db · web-flow · commit ba4c0dc63341 · 2025-07-04T16:13:20.000+05:30
* Add chunk latency

* Add chunk latency logs

* Add tests
diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java
@@ -23,6 +23,7 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.jdbc.telemetry.TelemetryClientFactory;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import com.google.common.annotations.VisibleForTesting;
 import java.sql.*;
 import java.util.*;
@@ -150,6 +151,7 @@ public void close() throws DatabricksSQLException {
     LOGGER.debug("public void close()");
     for (IDatabricksStatementInternal statement : statementSet) {
       statement.close(false);
+      ChunkLatencyHandler.getInstance().clearStatement(statement.getStatementId());
       statementSet.remove(statement);
     }
     this.session.close();
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java
@@ -1,16 +1,14 @@
 package com.databricks.jdbc.api.impl.arrow;
 
-import static com.databricks.jdbc.telemetry.TelemetryHelper.exportLatencyLog;
-
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.util.DatabricksThreadContextHolder;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
-import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.concurrent.ExecutionException;
@@ -49,7 +47,6 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
     boolean downloadSuccessful = false;
 
     // Sets context in the newly spawned thread
-    DatabricksThreadContextHolder.setChunkId(chunk.getChunkIndex());
     DatabricksThreadContextHolder.setConnectionContext(this.connectionContext);
     DatabricksThreadContextHolder.setStatementId(this.statementId);
 
@@ -67,7 +64,13 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
 
           chunk.downloadData(httpClient, chunkDownloader.getCompressionCodec());
           downloadSuccessful = true;
-        } catch (DatabricksParsingException | IOException e) {
+
+          // Record chunk download latency on successful download
+          long downloadLatency = System.currentTimeMillis() - startTime;
+          ChunkLatencyHandler.getInstance()
+              .recordChunkDownloadLatency(statementId, chunk.getChunkIndex(), downloadLatency);
+
+        } catch (IOException | DatabricksSQLException e) {
           retries++;
           if (retries >= MAX_RETRIES) {
             LOGGER.error(
@@ -80,7 +83,9 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
             throw new DatabricksSQLException(
                 "Failed to download chunk after multiple attempts",
                 e,
-                DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR);
+                statementId,
+                chunk.getChunkIndex(),
+                DatabricksDriverErrorCode.CHUNK_DOWNLOAD_ERROR.name());
           } else {
             LOGGER.warn(
                 String.format(
@@ -110,7 +115,6 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
         chunk.setStatus(ArrowResultChunk.ChunkStatus.DOWNLOAD_FAILED);
       }
 
-      exportLatencyLog(System.currentTimeMillis() - startTime);
       chunkDownloader.downloadProcessed(chunk.getChunkIndex());
       DatabricksThreadContextHolder.clearAllContext();
     }
diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java
@@ -10,6 +10,7 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import java.time.Instant;
 import java.util.Collection;
 import java.util.Map;
@@ -103,6 +104,7 @@ public ChunkLinkDownloadService(
     this.isDownloadInProgress = new AtomicBoolean(false);
     this.isDownloadChainStarted = new AtomicBoolean(false);
     this.isShutdown = false;
+    ChunkLatencyHandler.getInstance().initializeStatement(statementId, totalChunks);
 
     this.chunkIndexToLinkFuture = new ConcurrentHashMap<>();
     for (long i = 0; i < totalChunks; i++) {
diff --git a/src/main/java/com/databricks/jdbc/common/util/DatabricksThreadContextHolder.java b/src/main/java/com/databricks/jdbc/common/util/DatabricksThreadContextHolder.java
@@ -9,7 +9,6 @@ public class DatabricksThreadContextHolder {
   private static final ThreadLocal<IDatabricksConnectionContext> localConnectionContext =
       new ThreadLocal<>();
   private static final ThreadLocal<String> localStatementId = new ThreadLocal<>();
-  private static final ThreadLocal<Long> localChunkId = new ThreadLocal<>();
   private static final ThreadLocal<Integer> localRetryCount = new ThreadLocal<>();
   private static final ThreadLocal<StatementType> localStatementType = new ThreadLocal<>();
   private static final ThreadLocal<String> localSessionId = new ThreadLocal<>();
@@ -61,21 +60,12 @@ public static StatementType getStatementType() {
     return localStatementType.get();
   }
 
-  public static void setChunkId(Long chunkId) {
-    localChunkId.set(chunkId);
-  }
-
-  public static Long getChunkId() {
-    return localChunkId.get();
-  }
-
   public static void clearConnectionContext() {
     localConnectionContext.remove();
   }
 
   public static void clearStatementInfo() {
     localStatementId.remove();
-    localChunkId.remove();
     localStatementType.remove();
     localRetryCount.remove();
   }
diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksSQLException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksSQLException.java
@@ -25,6 +25,18 @@ public DatabricksSQLException(String reason, Throwable cause, String sqlState) {
         reason);
   }
 
+  // This constructor is used to export chunk download failure logs
+  public DatabricksSQLException(
+      String reason, Throwable cause, String statementId, Long chunkIndex, String sqlState) {
+    super(reason, sqlState, cause);
+    exportFailureLog(
+        DatabricksThreadContextHolder.getConnectionContext(),
+        DatabricksDriverErrorCode.CONNECTION_ERROR.name(),
+        reason,
+        chunkIndex,
+        statementId);
+  }
+
   public DatabricksSQLException(String reason, String sqlState) {
     super(reason, sqlState);
     exportFailureLog(DatabricksThreadContextHolder.getConnectionContext(), sqlState, reason);
diff --git a/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java b/src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java
@@ -121,9 +121,6 @@ public class DriverConnectionParameters {
   @JsonProperty("enable_sea_hybrid_results")
   boolean enableSeaHybridResults;
 
-  @JsonProperty("enable_complex_datatype_support")
-  boolean enableComplexSupport;
-
   @JsonProperty("allow_self_signed_support")
   boolean allowSelfSignedSupport;
 
@@ -331,11 +328,6 @@ public DriverConnectionParameters setEnableSeaHybridResults(boolean enableSeaHyb
     return this;
   }
 
-  public DriverConnectionParameters setEnableComplexSupport(boolean enableComplexSupport) {
-    this.enableComplexSupport = enableComplexSupport;
-    return this;
-  }
-
   public DriverConnectionParameters setAllowSelfSignedSupport(boolean allowSelfSignedSupport) {
     this.allowSelfSignedSupport = allowSelfSignedSupport;
     return this;
@@ -397,7 +389,6 @@ public String toString() {
         .add("nonProxyHosts", nonProxyHosts)
         .add("httpConnectionPoolSize", httpConnectionPoolSize)
         .add("enableSeaHybridResults", enableSeaHybridResults)
-        .add("enableComplexSupport", enableComplexSupport)
         .add("allowSelfSignedSupport", allowSelfSignedSupport)
         .add("useSystemTrustStore", useSystemTrustStore)
         .add("rowsFetchedPerBlock", rowsFetchedPerBlock)
diff --git a/src/main/java/com/databricks/jdbc/model/telemetry/latency/ChunkDetails.java b/src/main/java/com/databricks/jdbc/model/telemetry/latency/ChunkDetails.java
@@ -12,14 +12,20 @@ public class ChunkDetails {
   private Long slowestChunkLatencyMillis;
 
   @JsonProperty("total_chunks_present")
-  private Integer totalChunksPresent;
+  private Long totalChunksPresent;
 
   @JsonProperty("total_chunks_iterated")
-  private Integer totalChunksIterated;
+  private Long totalChunksIterated;
 
   @JsonProperty("sum_chunks_download_time_millis")
   private Long sumChunksDownloadTimeMillis;
 
+  public ChunkDetails(long totalChunks) {
+    this.totalChunksIterated = 0L;
+    this.sumChunksDownloadTimeMillis = 0L;
+    this.totalChunksPresent = totalChunks;
+  }
+
   public ChunkDetails setInitialChunkLatencyMillis(Long initialChunkLatencyMillis) {
     this.initialChunkLatencyMillis = initialChunkLatencyMillis;
     return this;
@@ -30,12 +36,12 @@ public ChunkDetails setSlowestChunkLatencyMillis(Long slowestChunkLatencyMillis)
     return this;
   }
 
-  public ChunkDetails setTotalChunksPresent(Integer totalChunksPresent) {
+  public ChunkDetails setTotalChunksPresent(Long totalChunksPresent) {
     this.totalChunksPresent = totalChunksPresent;
     return this;
   }
 
-  public ChunkDetails setTotalChunksIterated(Integer totalChunksIterated) {
+  public ChunkDetails setTotalChunksIterated(Long totalChunksIterated) {
     this.totalChunksIterated = totalChunksIterated;
     return this;
   }
@@ -45,6 +51,26 @@ public ChunkDetails setSumChunksDownloadTimeMillis(Long sumChunksDownloadTimeMil
     return this;
   }
 
+  public Long getInitialChunkLatencyMillis() {
+    return initialChunkLatencyMillis;
+  }
+
+  public Long getSlowestChunkLatencyMillis() {
+    return slowestChunkLatencyMillis;
+  }
+
+  public Long getTotalChunksPresent() {
+    return totalChunksPresent;
+  }
+
+  public Long getTotalChunksIterated() {
+    return totalChunksIterated;
+  }
+
+  public Long getSumChunksDownloadTimeMillis() {
+    return sumChunksDownloadTimeMillis;
+  }
+
   @Override
   public String toString() {
     return new ToStringer(ChunkDetails.class)
diff --git a/src/main/java/com/databricks/jdbc/telemetry/ITelemetryClient.java b/src/main/java/com/databricks/jdbc/telemetry/ITelemetryClient.java
@@ -6,4 +6,6 @@ public interface ITelemetryClient {
   void exportEvent(TelemetryFrontendLog event);
 
   void close();
+
+  void closeStatement(String statementId);
 }
diff --git a/src/main/java/com/databricks/jdbc/telemetry/NoopTelemetryClient.java b/src/main/java/com/databricks/jdbc/telemetry/NoopTelemetryClient.java
@@ -25,4 +25,9 @@ public void exportEvent(TelemetryFrontendLog event) {
   public void close() {
     // do nothing
   }
+
+  @Override
+  public void closeStatement(String statementId) {
+    // do nothing
+  }
 }
diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java
@@ -2,6 +2,8 @@
 
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.model.telemetry.TelemetryFrontendLog;
+import com.databricks.jdbc.model.telemetry.latency.ChunkDetails;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import com.databricks.sdk.core.DatabricksConfig;
 import java.util.LinkedList;
 import java.util.List;
@@ -51,6 +53,23 @@ public void exportEvent(TelemetryFrontendLog event) {
 
   @Override
   public void close() {
+    // Export any pending chunk latency telemetry before flushing
+    ChunkLatencyHandler.getInstance()
+        .getAllPendingChunkDetails()
+        .forEach(
+            (statementId, chunkDetails) -> {
+              TelemetryHelper.exportChunkLatencyTelemetry(chunkDetails, statementId);
+            });
+    flush();
+  }
+
+  @Override
+  public void closeStatement(String statementId) {
+    ChunkDetails chunkDetails =
+        ChunkLatencyHandler.getInstance().getChunkDetailsAndCleanup(statementId);
+    if (chunkDetails != null) {
+      TelemetryHelper.exportChunkLatencyTelemetry(chunkDetails, statementId);
+    }
     flush();
   }
 
diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java
@@ -12,6 +12,7 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.*;
+import com.databricks.jdbc.model.telemetry.latency.ChunkDetails;
 import com.databricks.sdk.core.DatabricksConfig;
 import com.databricks.sdk.core.ProxyConfig;
 import com.databricks.sdk.core.UserAgent;
@@ -85,6 +86,20 @@ public static void exportInitialTelemetryLog(IDatabricksConnectionContext connec
 
   public static void exportFailureLog(
       IDatabricksConnectionContext connectionContext, String errorName, String errorMessage) {
+    exportFailureLog(
+        connectionContext,
+        errorName,
+        errorMessage,
+        null,
+        DatabricksThreadContextHolder.getStatementId());
+  }
+
+  public static void exportFailureLog(
+      IDatabricksConnectionContext connectionContext,
+      String errorName,
+      String errorMessage,
+      Long chunkIndex,
+      String statementId) {
 
     // Connection context is not set in following scenarios:
     // a. Unit tests
@@ -101,10 +116,18 @@ public static void exportFailureLog(
                   new FrontendLogEntry()
                       .setSqlDriverLog(
                           new TelemetryEvent()
+                              .setSqlStatementId(statementId)
                               .setDriverConnectionParameters(
                                   getDriverConnectionParameter(connectionContext))
                               .setDriverErrorInfo(errorInfo)
                               .setDriverSystemConfiguration(getDriverSystemConfiguration())));
+      if (chunkIndex != null) {
+        // When chunkIndex is provided, we are exporting a chunk download failure log
+        telemetryFrontendLog
+            .getEntry()
+            .getSqlDriverLog()
+            .setSqlOperation(new SqlExecutionEvent().setChunkId(chunkIndex));
+      }
       ITelemetryClient client =
           TelemetryClientFactory.getInstance().getTelemetryClient(connectionContext);
       client.exportEvent(telemetryFrontendLog);
@@ -115,8 +138,7 @@ public static void exportLatencyLog(long executionTime) {
     SqlExecutionEvent executionEvent =
         new SqlExecutionEvent()
             .setDriverStatementType(DatabricksThreadContextHolder.getStatementType())
-            .setRetryCount(DatabricksThreadContextHolder.getRetryCount())
-            .setChunkId(DatabricksThreadContextHolder.getChunkId());
+            .setRetryCount(DatabricksThreadContextHolder.getRetryCount());
     exportLatencyLog(
         DatabricksThreadContextHolder.getConnectionContext(),
         executionTime,
@@ -125,6 +147,35 @@ public static void exportLatencyLog(long executionTime) {
         DatabricksThreadContextHolder.getSessionId());
   }
 
+  public static void exportChunkLatencyTelemetry(ChunkDetails chunkDetails, String statementId) {
+    if (chunkDetails == null) {
+      return;
+    }
+
+    IDatabricksConnectionContext connectionContext =
+        DatabricksThreadContextHolder.getConnectionContext();
+    if (connectionContext == null) {
+      return;
+    }
+
+    SqlExecutionEvent sqlExecutionEvent = new SqlExecutionEvent().setChunkDetails(chunkDetails);
+
+    TelemetryEvent telemetryEvent =
+        new TelemetryEvent()
+            .setSqlOperation(sqlExecutionEvent)
+            .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext));
+
+    TelemetryFrontendLog telemetryFrontendLog =
+        new TelemetryFrontendLog()
+            .setFrontendLogEventId(getEventUUID())
+            .setContext(getLogContext())
+            .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent));
+
+    TelemetryClientFactory.getInstance()
+        .getTelemetryClient(connectionContext)
+        .exportEvent(telemetryFrontendLog);
+  }
+
   @VisibleForTesting
   static void exportLatencyLog(
       IDatabricksConnectionContext connectionContext,
@@ -231,7 +282,6 @@ private static DriverConnectionParameters buildDriverConnectionParameters(
             .setNonProxyHosts(StringUtil.split(connectionContext.getNonProxyHosts()))
             .setHttpConnectionPoolSize(connectionContext.getHttpConnectionPoolSize())
             .setEnableSeaHybridResults(connectionContext.isSqlExecHybridResultsEnabled())
-            .setEnableComplexSupport(connectionContext.isComplexDatatypeSupportEnabled())
             .setAllowSelfSignedSupport(connectionContext.allowSelfSignedCerts())
             .setUseSystemTrustStore(connectionContext.useSystemTrustStore())
             .setRowsFetchedPerBlock(connectionContext.getRowsFetchedPerBlock())
diff --git a/src/main/java/com/databricks/jdbc/telemetry/latency/ChunkLatencyHandler.java b/src/main/java/com/databricks/jdbc/telemetry/latency/ChunkLatencyHandler.java
diff --git a/src/test/java/com/databricks/jdbc/telemetry/latency/ChunkLatencyHandlerTest.java b/src/test/java/com/databricks/jdbc/telemetry/latency/ChunkLatencyHandlerTest.java

Original file line number	Diff line number	Diff line change
`@@ -6,4 +6,6 @@ public interface ITelemetryClient {`
`6`	`6`	`void exportEvent(TelemetryFrontendLog event);`
`7`	`7`
`8`	`8`	`void close();`
	`9`	`+`
	`10`	`+ void closeStatement(String statementId);`
`9`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,4 +25,9 @@ public void exportEvent(TelemetryFrontendLog event) {`
`25`	`25`	`public void close() {`
`26`	`26`	`// do nothing`
`27`	`27`	`}`
	`28`	`+`
	`29`	`+ @Override`
	`30`	`+ public void closeStatement(String statementId) {`
	`31`	`+ // do nothing`
	`32`	`+ }`
`28`	`33`	`}`