[ES-1765150] Coalesce concurrent expired-link refreshes into single batch RPC (#1277)

## Summary
- When multiple cloud fetch download threads detect expired presigned
URLs simultaneously, they each made independent FetchResults RPCs,
causing thread pool exhaustion under high concurrency
- Introduces a centralized coalescing mechanism in
`StreamingChunkProvider.getRefreshedLink()` that serializes refresh
requests under a lock, performs a single batch `fetchLinks()` call from
the lowest expired chunk offset, and distributes fresh links to all
expired chunks
- Adds `LinkRefresher` functional interface to decouple download tasks
from the coalescing logic, with fallback to single `refetchLink()` when
the batch doesn't include the target chunk

## Test plan
- [x] `StreamingChunkDownloadTaskTest` — 6/6 pass (uses `LinkRefresher`
mock)
- [x] `StreamingChunkProviderTest$LinkExpiryTests` — 9/9 pass (verify
batch coalescing behavior)
- [ ] Manual verification with high-concurrency cloud fetch workload to
confirm reduced FetchResults RPC count

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Gopal Lal <gopal.lal@databricks.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -20,6 +20,7 @@
 - Fixed socket leak in SDK HTTP client that prevented CRaC checkpointing. The SDK's connection pool was not shut down on `connection.close()`, leaving TCP sockets open.
 - Fixed Date fields within complex types (ARRAY, STRUCT, MAP) being returned as epoch day integers instead of proper date values.
 - Fixed `DatabaseMetaData.getColumns()` returning the column type name in `COLUMN_DEF` for columns with no default value. `COLUMN_DEF` now correctly returns `null` per the JDBC specification.
+- Coalesce concurrent expired cloud fetch link refreshes into a single batch FetchResults RPC to prevent thread pool exhaustion under high concurrency.
 
 ---
 *Note: When making changes, please add your change under the appropriate section

src/main/java/com/databricks/jdbc/api/impl/arrow/LinkRefresher.java

@@ -0,0 +1,22 @@
+package com.databricks.jdbc.api.impl.arrow;
+
+import com.databricks.jdbc.model.core.ExternalLink;
+import java.sql.SQLException;
+
+/**
+ * Callback interface for refreshing expired chunk links. Used by {@link StreamingChunkDownloadTask}
+ * to delegate link refresh to {@link StreamingChunkProvider}, which coalesces concurrent refresh
+ * requests into a single batch RPC.
+ */
+@FunctionalInterface
+interface LinkRefresher {
+  /**
+   * Refreshes an expired link for the given chunk.
+   *
+   * @param chunkIndex The chunk index whose link has expired
+   * @param rowOffset The row offset of the chunk (used by Thrift for FETCH_ABSOLUTE)
+   * @return The refreshed ExternalLink with a new expiration time
+   * @throws SQLException if the refresh operation fails
+   */
+  ExternalLink refreshLink(long chunkIndex, long rowOffset) throws SQLException;
+}

src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTask.java

@@ -5,15 +5,15 @@
 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 java.io.IOException;
 import java.sql.SQLException;
 import java.util.concurrent.Callable;
 
 /**
- * A download task for streaming chunk provider. Simpler than ChunkDownloadTask - uses
- * ChunkLinkFetcher directly for link refresh instead of ChunkLinkDownloadService.
+ * A download task for streaming chunk provider. Uses a {@link LinkRefresher} callback for link
+ * refresh, which enables the {@link StreamingChunkProvider} to coalesce concurrent expired-link
+ * refreshes into a single batch RPC.
  */
 public class StreamingChunkDownloadTask implements Callable<Void> {
 
@@ -26,19 +26,19 @@ public class StreamingChunkDownloadTask implements Callable<Void> {
   private final ArrowResultChunk chunk;
   private final IDatabricksHttpClient httpClient;
   private final CompressionCodec compressionCodec;
-  private final ChunkLinkFetcher linkFetcher;
+  private final LinkRefresher linkRefresher;
   private final double cloudFetchSpeedThreshold;
 
   public StreamingChunkDownloadTask(
       ArrowResultChunk chunk,
       IDatabricksHttpClient httpClient,
       CompressionCodec compressionCodec,
-      ChunkLinkFetcher linkFetcher,
+      LinkRefresher linkRefresher,
       double cloudFetchSpeedThreshold) {
     this.chunk = chunk;
     this.httpClient = httpClient;
     this.compressionCodec = compressionCodec;
-    this.linkFetcher = linkFetcher;
+    this.linkRefresher = linkRefresher;
     this.cloudFetchSpeedThreshold = cloudFetchSpeedThreshold;
   }
 
@@ -50,12 +50,12 @@ public Void call() throws DatabricksSQLException {
     try {
       while (!downloadSuccessful) {
         try {
-          // Check if link is expired and refresh if needed
+          // Check if link is expired and refresh if needed.
+          // The LinkRefresher (StreamingChunkProvider.getRefreshedLink) updates the chunk's
+          // link directly under the refetchLock, so we don't need to set it here.
           if (chunk.isChunkLinkInvalid()) {
             LOGGER.debug("Link invalid for chunk {}, refetching", chunk.getChunkIndex());
-            ExternalLink freshLink =
-                linkFetcher.refetchLink(chunk.getChunkIndex(), chunk.getStartRowOffset());
-            chunk.setChunkLink(freshLink);
+            linkRefresher.refreshLink(chunk.getChunkIndex(), chunk.getStartRowOffset());
           }
 
           // Perform the download

src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java

@@ -99,6 +99,11 @@ public class StreamingChunkProvider implements ChunkProvider {
   // that must be consistent within the loop.
   private final ReentrantLock downloadLock = new ReentrantLock();
 
+  // Synchronization for coalescing concurrent expired-link refreshes.
+  // When multiple download threads detect expired links simultaneously, only one thread
+  // performs the actual batch FetchResults RPC while others wait and reuse the result.
+  private final Object refetchLock = new Object();
+
   // Executors
   private final ExecutorService downloadExecutor;
   private final Thread linkPrefetchThread;
@@ -511,12 +516,112 @@ private void triggerDownloads() {
     }
   }
 
+  /**
+   * Coalesces concurrent expired-link refresh requests into a single batch RPC.
+   *
+   * <p>When multiple download threads detect expired links simultaneously, the first thread to
+   * acquire the lock performs a single {@link ChunkLinkFetcher#fetchLinks} call from the lowest
+   * expired chunk offset. The response refreshes all expired in-memory chunks. Subsequent threads
+   * find their chunks already refreshed via the double-check pattern.
+   *
+   * <p>If the batch response does not include the requested chunk, falls back to a single {@link
+   * ChunkLinkFetcher#refetchLink} call.
+   *
+   * @param chunkIndex The chunk index whose link has expired
+   * @param rowOffset The row offset of the chunk (used by Thrift)
+   * @return The refreshed ExternalLink
+   * @throws SQLException if the refresh fails
+   */
+  ExternalLink getRefreshedLink(long chunkIndex, long rowOffset) throws SQLException {
+    synchronized (refetchLock) {
+      // Double-check: another thread may have already refreshed while we waited on the lock
+      ArrowResultChunk targetChunk = chunks.get(chunkIndex);
+      if (targetChunk != null && !targetChunk.isChunkLinkInvalid()) {
+        return targetChunk.getChunkLink();
+      }
+
+      // Find the minimum expired chunk index among pre-download chunks.
+      // Only called from StreamingChunkDownloadTask before download starts, so the
+      // target chunk is always pre-download. We also skip chunks that have already been
+      // downloaded (DOWNLOAD_IN_PROGRESS or later) since their links are no longer needed.
+      long minExpiredIndex = Long.MAX_VALUE;
+      long minExpiredRowOffset = 0;
+      int expiredCount = 0;
+      for (ArrowResultChunk c : chunks.values()) {
+        ChunkStatus status = c.getStatus();
+        if (status != ChunkStatus.PENDING
+            && status != ChunkStatus.URL_FETCHED
+            && status != ChunkStatus.DOWNLOAD_FAILED
+            && status != ChunkStatus.DOWNLOAD_RETRY) {
+          continue; // Already downloading or downloaded — link refresh not needed
+        }
+        if (c.isChunkLinkInvalid()) {
+          expiredCount++;
+          if (c.getChunkIndex() < minExpiredIndex) {
+            minExpiredIndex = c.getChunkIndex();
+            minExpiredRowOffset = c.getStartRowOffset();
+          }
+        }
+      }
+
+      if (minExpiredIndex == Long.MAX_VALUE) {
+        // No expired chunks found — race condition resolved by another thread
+        if (targetChunk != null) {
+          return targetChunk.getChunkLink();
+        }
+        throw new DatabricksSQLException(
+            "Chunk " + chunkIndex + " not found during link refresh",
+            DatabricksDriverErrorCode.CHUNK_READY_ERROR);
+      }
+
+      LOGGER.info(
+          "Coalesced link refresh: fetching from chunk {} (row offset {}) for {} expired chunks",
+          minExpiredIndex,
+          minExpiredRowOffset,
+          expiredCount);
+
+      // Single batch FetchResults RPC from the lowest expired offset
+      ChunkLinkFetchResult result = linkFetcher.fetchLinks(minExpiredIndex, minExpiredRowOffset);
+
+      // Update ALL pre-download chunks that received fresh links.
+      // Always overwrite even if the current link hasn't expired yet, since the
+      // server-provided link has a later expiry and prevents near-expiry races.
+      for (ExternalLink link : result.getChunkLinks()) {
+        ArrowResultChunk c = chunks.get(link.getChunkIndex());
+        if (c != null) {
+          ChunkStatus status = c.getStatus();
+          if (status == ChunkStatus.PENDING
+              || status == ChunkStatus.URL_FETCHED
+              || status == ChunkStatus.DOWNLOAD_FAILED
+              || status == ChunkStatus.DOWNLOAD_RETRY) {
+            c.setChunkLink(link);
+          }
+        }
+      }
+
+      // Check if our target chunk was refreshed by the batch
+      targetChunk = chunks.get(chunkIndex);
+      if (targetChunk != null && !targetChunk.isChunkLinkInvalid()) {
+        return targetChunk.getChunkLink();
+      }
+
+      // Fallback: batch response did not include the requested chunk
+      LOGGER.warn(
+          "Batch refresh did not include chunk {}, falling back to single refetch", chunkIndex);
+      ExternalLink fallbackLink = linkFetcher.refetchLink(chunkIndex, rowOffset);
+      if (targetChunk != null) {
+        targetChunk.setChunkLink(fallbackLink);
+      }
+      return fallbackLink;
+    }
+  }
+
   private void submitDownloadTask(ArrowResultChunk chunk) {
     LOGGER.debug("Submitting download task for chunk {}", chunk.getChunkIndex());
 
     StreamingChunkDownloadTask task =
         new StreamingChunkDownloadTask(
-            chunk, httpClient, compressionCodec, linkFetcher, cloudFetchSpeedThreshold);
+            chunk, httpClient, compressionCodec, this::getRefreshedLink, cloudFetchSpeedThreshold);
 
     downloadExecutor.submit(task);
   }

src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkDownloadTaskTest.java

@@ -38,7 +38,7 @@ public class StreamingChunkDownloadTaskTest {
 
   @Mock private ArrowResultChunk chunk;
   @Mock private IDatabricksHttpClient httpClient;
-  @Mock private ChunkLinkFetcher linkFetcher;
+  @Mock private LinkRefresher linkRefresher;
 
   private StreamingChunkDownloadTask downloadTask;
   private CompletableFuture<Void> downloadFuture;
@@ -48,7 +48,7 @@ void setUp() {
     downloadFuture = new CompletableFuture<>();
     downloadTask =
         new StreamingChunkDownloadTask(
-            chunk, httpClient, CompressionCodec.NONE, linkFetcher, CLOUD_FETCH_SPEED_THRESHOLD);
+            chunk, httpClient, CompressionCodec.NONE, linkRefresher, CLOUD_FETCH_SPEED_THRESHOLD);
   }
 
   @Test
@@ -137,17 +137,16 @@ void testLinkRefreshWhenLinkIsInvalid() throws Exception {
     when(chunk.isChunkLinkInvalid()).thenReturn(true, false);
 
     ExternalLink freshLink = mock(ExternalLink.class);
-    when(linkFetcher.refetchLink(5L, 100L)).thenReturn(freshLink);
+    when(linkRefresher.refreshLink(5L, 100L)).thenReturn(freshLink);
 
     doNothing()
         .when(chunk)
         .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
 
     downloadTask.call();
 
-    // Verify link was refreshed
-    verify(linkFetcher, times(1)).refetchLink(5L, 100L);
-    verify(chunk, times(1)).setChunkLink(freshLink);
+    // Verify link was refreshed (setChunkLink is done inside getRefreshedLink, not here)
+    verify(linkRefresher, times(1)).refreshLink(5L, 100L);
     verify(chunk, times(1))
         .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
     assertTrue(downloadFuture.isDone());
@@ -162,7 +161,7 @@ void testLinkRefreshFailureAndRetry() throws Exception {
     when(chunk.isChunkLinkInvalid()).thenReturn(true, true, false);
 
     ExternalLink freshLink = mock(ExternalLink.class);
-    when(linkFetcher.refetchLink(5L, 100L)).thenReturn(freshLink);
+    when(linkRefresher.refreshLink(5L, 100L)).thenReturn(freshLink);
 
     // First download fails with IOException, retry succeeds
     doThrow(new SocketException("Connection reset"))
@@ -172,9 +171,8 @@ void testLinkRefreshFailureAndRetry() throws Exception {
 
     downloadTask.call();
 
-    // Should refresh link twice (once per attempt)
-    verify(linkFetcher, times(2)).refetchLink(5L, 100L);
-    verify(chunk, times(2)).setChunkLink(freshLink);
+    // Should refresh link twice (once per attempt; setChunkLink done inside getRefreshedLink)
+    verify(linkRefresher, times(2)).refreshLink(5L, 100L);
     verify(chunk, times(2))
         .downloadData(httpClient, CompressionCodec.NONE, CLOUD_FETCH_SPEED_THRESHOLD);
     verify(chunk, times(1)).setStatus(ChunkStatus.DOWNLOAD_RETRY);
@@ -253,7 +251,7 @@ void testStatusTransitionsDuringRetries() throws Exception {
             spiedChunk,
             httpClient,
             CompressionCodec.NONE,
-            linkFetcher,
+            linkRefresher,
             CLOUD_FETCH_SPEED_THRESHOLD);
 
     // Execute the task