Fix race condition between chunk release and download error handling (#1407)

## Description

Fixes #1367 — Race condition between `CHUNK_RELEASED` and
`DOWNLOAD_FAILED` causes invalid state transition warnings during Arrow
Cloud Fetch operations.

### Root Cause

When chunk download threads are running in parallel and their
connections are cut off by `IdleConnectionEvictor` (due to no
socket-level data movement caused by extreme resource pressure in
low-resource environments), the downloads fail. When the consumer
reaches a failed chunk, `rs.next()` throws a `DatabricksSQLException`.
This can trigger `rs.close()` — either explicitly by the user's error
handling code, implicitly by try-with-resources, or automatically by
frameworks that close resources on error. Additionally, `rs.close()` can
also be called independently by the user before consuming all data.

In either case, `doClose()` calls `shutdownNow()` on the download thread
pool and releases all chunks to `CHUNK_RELEASED`. But `shutdownNow()`
only sends an interrupt signal — it doesn't wait for threads to stop.
Download threads that are still processing their errors may try to set
`DOWNLOAD_FAILED` after their chunks have already been released. The
transition `CHUNK_RELEASED → DOWNLOAD_FAILED` is invalid in the chunk
lifecycle state machine, causing the error/warning.

### Race Condition Sequence

```
Thread A (consumer/main)              Thread B (download)
─────────────────────────             ─────────────────────
                                      downloading chunk [7]...
                                      connection killed by evictor
                                      IOException caught
                                      retries exhausted
                                      ↓
rs.next() → gets error                (processing error...)
rs.close()
  → doClose()
    → shutdownNow()                   (interrupt sent, but not yet processed)
    → chunk[7].releaseChunk()
      → CHUNK_RELEASED ✓
                                      → setStatus(DOWNLOAD_FAILED)
                                        CHUNK_RELEASED → DOWNLOAD_FAILED ✗
                                        Invalid transition!
```

### Fix

Added `awaitTermination(3, TimeUnit.SECONDS)` after `shutdownNow()` in
both `RemoteChunkProvider.doClose()` and
`StreamingChunkProvider.close()`. This ensures download threads complete
their error handling path (catch → finally → setStatus) before chunks
are released.

The 3-second timeout is a conservative upper bound. After
`shutdownNow()` interrupts threads, they exit their retry sleep
(`Thread.sleep(1500)`) immediately via `InterruptedException` and
process the error path in milliseconds. `awaitTermination()` returns as
soon as all threads finish — it does not wait the full 3 seconds.

### Files Changed

| File | Change |
|------|--------|
| `RemoteChunkProvider.java` | Added `awaitTermination(3s)` after
`shutdownNow()` in `doClose()`, added INFO log for close lifecycle |
| `StreamingChunkProvider.java` | Added `awaitTermination(3s)` after
`shutdownNow()` in `close()` |
| `NEXT_CHANGELOG.md` | Added changelog entry |

## Testing

### Reproduction (Docker + iptables)

Reproduced the error/warning logs from issue #1367 using:
1. Docker container with network throttled to 5mbit (slow chunk
downloads)
2. `iptables REJECT --tcp-reset` to kill connections mid-download
(simulating `IdleConnectionEvictor` behavior)
3. Calling `rs.close()` while download threads process errors
4. Temporary `MAX_RETRIES=0` + uninterruptible sleep before
`setStatus(DOWNLOAD_FAILED)` to widen the race window

**Before fix** — driver logs showed:
```
WARNING: Invalid state transition for chunk [1]: CHUNK_RELEASED -> DOWNLOAD_FAILED
WARNING: Invalid state transition for chunk [2]: CHUNK_RELEASED -> DOWNLOAD_FAILED
```

**After fix** — no invalid state transitions. `awaitTermination()`
waited for threads to finish before releasing chunks.

Verified for both `RemoteChunkProvider` and `StreamingChunkProvider`.

## Additional Notes to the Reviewer

- The `IdleConnectionEvictor` shares a connection pool between the
Thrift transport and chunk downloads (`DatabricksHttpClient` via
`DatabricksHttpClientFactory`). In resource-constrained environments,
CPU starvation causes the Java layer to stop reading from sockets,
making active download connections appear idle to the evictor.
- `shutdownNow()` is best-effort per Java docs — it cannot guarantee
threads will stop. The interrupt flag is set, but threads only respond
when they hit an interruptible operation (`Thread.sleep`, I/O). Code
after the interrupt (catch blocks, finally blocks, `setStatus()` calls)
is pure computation that doesn't check interrupts.
- The same race pattern existed in both `RemoteChunkProvider` (standard
Cloud Fetch) and `StreamingChunkProvider` (opt-in via
`EnableStreamingChunkProvider`). Both are fixed.
- `RemoteChunkProviderV2` (incubator) downloads chunks sequentially
without a thread pool, so it is not affected.

---------

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Co-authored-by: Samikshya Chand <148681192+samikshya-db@users.noreply.github.com>

Author: sreekanth-db

NEXT_CHANGELOG.md

@@ -10,6 +10,7 @@
 - Added support for using SQL SHOW commands for Thrift-mode metadata operations (`getTables`, `getColumns`, `getSchemas`, `getFunctions`, `getPrimaryKeys`, `getImportedKeys`, `getCrossReference`). Enable by setting `UseQueryForMetadata=1`. This aligns Thrift metadata behavior with Statement Execution API (SEA) mode.
 
 ### Fixed
+- Fixed race condition between chunk download error handling and result set close that could cause invalid state transition warnings (`CHUNK_RELEASED -> DOWNLOAD_FAILED`) during Arrow Cloud Fetch operations in resource-constrained environments.
 - Fixed `EnableBatchedInserts` silently falling back to individual execution when table or schema names contain special characters (e.g., hyphens) inside backtick-quoted identifiers. Added a warn log when the fallback occurs.
 - Fixed `IntervalConverter` crash (`IllegalArgumentException: Invalid interval metadata`) when INTERVAL columns are returned via CloudFetch. Arrow metadata from CloudFetch uses underscored format (`INTERVAL_YEAR_MONTH`, `INTERVAL_DAY_TIME`) which the driver's regex did not accept.
 - Fixed `Statement` being prematurely closed after queries that return inline results, which prevented re-execution, `getResultSet()`, and `getExecutionResult()` from working. Statements now remain open and reusable until explicitly closed by the caller.

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

@@ -241,6 +241,11 @@ protected ValueVector getColumnVector(int recordBatchIndex, int columnIndex) {
    * @param targetStatus new status to set
    */
   protected void setStatus(ChunkStatus targetStatus) {
+    if (getStatus() == ChunkStatus.CHUNK_RELEASED) {
+      LOGGER.debug(
+          "Skipping transition to {} — chunk [{}] already released", targetStatus, chunkIndex);
+      return;
+    }
     try {
       stateMachine.transition(targetStatus);
     } catch (DatabricksParsingException e) {

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

@@ -7,6 +7,8 @@
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp;
 import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowResultLink;
 import com.databricks.jdbc.model.core.ResultData;
@@ -16,9 +18,11 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class RemoteChunkProvider extends AbstractRemoteChunkProvider<ArrowResultChunk> {
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(RemoteChunkProvider.class);
   private static final String CHUNKS_DOWNLOADER_THREAD_POOL_PREFIX =
       "databricks-jdbc-chunks-downloader-";
   private ExecutorService chunkDownloaderExecutorService;
@@ -125,8 +129,26 @@ public void downloadNextChunks() {
   /** {@inheritDoc} */
   @Override
   protected void doClose() {
+    LOGGER.debug(
+        "doClose() called — shutting down executor and releasing all {} chunks (thread: {})",
+        chunkIndexToChunksMap.size(),
+        Thread.currentThread().getName());
     isClosed = true;
-    chunkDownloaderExecutorService.shutdownNow();
+    if (chunkDownloaderExecutorService != null) {
+      chunkDownloaderExecutorService.shutdownNow();
+      // Wait for download threads to finish error handling before releasing chunks.
+      // After shutdownNow() interrupts threads, they exit their retry sleep and process
+      // the error path (catch → finally → setStatus) in milliseconds. 3 seconds is a
+      // conservative upper bound to avoid racing with that error handling path.
+      try {
+        if (!chunkDownloaderExecutorService.awaitTermination(3, TimeUnit.SECONDS)) {
+          LOGGER.warn("Download threads did not terminate within timeout");
+        }
+      } catch (InterruptedException e) {
+        LOGGER.error(e, "Interrupted while waiting for download threads to terminate");
+        Thread.currentThread().interrupt();
+      }
+    }
     chunkIndexToChunksMap.values().forEach(ArrowResultChunk::releaseChunk);
     DatabricksThreadContextHolder.clearStatementInfo();
   }

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

@@ -18,6 +18,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
@@ -276,6 +277,18 @@ public void close() {
     // Shutdown download executor
     if (downloadExecutor != null) {
       downloadExecutor.shutdownNow();
+      // Wait for download threads to finish error handling before releasing chunks.
+      // After shutdownNow() interrupts threads, they exit their retry sleep and process
+      // the error path (catch → finally → setStatus) in milliseconds. 3 seconds is a
+      // conservative upper bound to avoid racing with that error handling path.
+      try {
+        if (!downloadExecutor.awaitTermination(3, TimeUnit.SECONDS)) {
+          LOGGER.warn("Download threads did not terminate within timeout");
+        }
+      } catch (InterruptedException e) {
+        LOGGER.error(e, "Interrupted while waiting for download threads to terminate");
+        Thread.currentThread().interrupt();
+      }
     }
 
     // Release all chunks