From 709b4ab0e6fa786b1379c6efc58ff305f495e460 Mon Sep 17 00:00:00 2001
From: Lorna <115649563+sfc-gh-ext-simba-lb@users.noreply.github.com>
Date: Thu, 27 Jul 2023 10:47:05 -0700
Subject: [PATCH] SNOW-872499: Disable max retries in chunk downloader if
 maxHttpRetries is 0 (#1488)

disable max retries in chunk downloader if maxHttpRetries is 0
---
 .../client/jdbc/SnowflakeChunkDownloader.java       | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeChunkDownloader.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeChunkDownloader.java
index 00c94b2ab..a2bd4135e 100644
--- a/src/main/java/net/snowflake/client/jdbc/SnowflakeChunkDownloader.java
+++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeChunkDownloader.java
@@ -636,7 +636,7 @@ public SnowflakeResultChunk getNextChunkToConsume()
   private void waitForChunkReady(SnowflakeResultChunk currentChunk) throws InterruptedException {
     int retry = 0;
     long startTime = System.currentTimeMillis();
-    while (currentChunk.getDownloadState() != DownloadState.SUCCESS && retry < maxHttpRetries) {
+    while (true) {
       logger.debug(
           "Thread {} is waiting for #chunk{} to be ready, current" + "chunk state is: {}, retry={}",
           Thread.currentThread().getId(),
@@ -644,7 +644,8 @@ private void waitForChunkReady(SnowflakeResultChunk currentChunk) throws Interru
           currentChunk.getDownloadState(),
           retry);
 
-      if (currentChunk.getDownloadState() != DownloadState.FAILURE) {
+      if (currentChunk.getDownloadState() != DownloadState.FAILURE
+          && currentChunk.getDownloadState() != DownloadState.SUCCESS) {
         // if the state is not failure, we should keep waiting; otherwise, we skip
         // waiting
         if (!currentChunk
@@ -669,6 +670,7 @@ private void waitForChunkReady(SnowflakeResultChunk currentChunk) throws Interru
         }
       }
 
+      // retry if chunk is not successfully downloaded
       if (currentChunk.getDownloadState() != DownloadState.SUCCESS) {
         retry++;
         // timeout or failed
@@ -715,6 +717,13 @@ private void waitForChunkReady(SnowflakeResultChunk currentChunk) throws Interru
           nextChunkToDownload = nextChunkToConsume + 1;
         }
       }
+
+      // exit if chunk has downloaded or we have hit max retry
+      // maxHttpRetries = 0 will retry indefinitely
+      if (currentChunk.getDownloadState() == DownloadState.SUCCESS
+          || (maxHttpRetries > 0 && retry >= maxHttpRetries)) {
+        break;
+      }
     }
     if (currentChunk.getDownloadState() == DownloadState.SUCCESS) {
       logger.debug("ready to consume #chunk{}, succeed retry={}", nextChunkToConsume, retry);