SNOW-2159000 Allow thread pool to scale properly (#2231)

Co-authored-by: Mikołaj Kubik <mikolaj.kubik@snowflake.com>

Author: sfc-gh-rkowalski

src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java

@@ -12,7 +12,6 @@
 import java.util.concurrent.Future;
 import net.snowflake.client.core.HttpUtil;
 import net.snowflake.client.core.ObjectMapperFactory;
-import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SFSession;
 import net.snowflake.client.jdbc.SnowflakeConnectionV1;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
@@ -27,7 +26,7 @@
 
 /** Telemetry Service Interface */
 public class TelemetryClient implements Telemetry {
-  private static final SFLogger logger = SFLoggerFactory.getLogger(SFBaseSession.class);
+  private static final SFLogger logger = SFLoggerFactory.getLogger(TelemetryClient.class);
 
   private static final String SF_PATH_TELEMETRY = "/telemetry/send";
   private static final String SF_PATH_TELEMETRY_SESSIONLESS = "/telemetry/send/sessionless";

src/main/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryThreadPool.java

@@ -12,7 +12,8 @@
  * asynchronously to server
  */
 public class TelemetryThreadPool {
-  private ExecutorService uploader;
+  private static final int CORE_POOL_SIZE = 10;
+  private final ExecutorService uploader;
 
   private static TelemetryThreadPool instance;
 
@@ -27,15 +28,27 @@ public static TelemetryThreadPool getInstance() {
     return instance;
   }
 
+  /**
+   * Private constructor to initialize the singleton instance.
+   *
+   * <p>Configures a thread pool that scales dynamically based on workload. The pool starts with
+   * zero threads and will create new threads on demand up to a maximum of 10. If all 10 threads are
+   * active, new tasks are placed in an unbounded queue to await execution.
+   *
+   * <p>To conserve resources, threads that are idle for more than 30 seconds are terminated,
+   * allowing the pool to shrink back to zero during periods of inactivity.
+   */
   private TelemetryThreadPool() {
     uploader =
         new ThreadPoolExecutor(
-            0, // core size
-            10, // max size
-            1, // keep alive time
+            CORE_POOL_SIZE, // core size
+            CORE_POOL_SIZE, // max size
+            30L, // keep alive time
             TimeUnit.SECONDS,
             new LinkedBlockingQueue<>() // work queue
             );
+    // Allow core threads to time out and be terminated when idle.
+    ((ThreadPoolExecutor) uploader).allowCoreThreadTimeOut(true);
   }
 
   public void execute(Runnable task) {

src/test/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryThreadPoolTest.java

@@ -0,0 +1,89 @@
+package net.snowflake.client.jdbc.telemetryOOB;
+
+import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import net.snowflake.client.category.TestTags;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+@Tag(TestTags.OTHERS)
+class TelemetryThreadPoolTest {
+
+  /**
+   * This test confirms fixing the issue where the TelemetryThreadPool did not scale up to its
+   * maximum pool size due to its configuration with a core pool size of 0 and an unbounded work
+   * queue (LinkedBlockingQueue). https://github.com/snowflakedb/snowflake-jdbc/issues/2219
+   *
+   * <p>This test submits 100 tasks that each sleep for a short duration. If the pool were scaling,
+   * it would create up to 10 threads to handle the load concurrently. Previously only a single
+   * thread was ever created to process all tasks sequentially from the queue.
+   *
+   * @throws InterruptedException if the waiting thread is interrupted.
+   */
+  @Test
+  void testThreadPoolScaling() throws InterruptedException {
+    TelemetryThreadPool telemetryPool = TelemetryThreadPool.getInstance();
+    int taskCount = 100; // Submit more tasks than the max pool size (10)
+    ThreadPoolExecutor executor = getThreadPoolExecutor(telemetryPool);
+
+    final Set<Long> uniqueThreadIds = Collections.newSetFromMap(new ConcurrentHashMap<>());
+    final CountDownLatch completionLatch = new CountDownLatch(taskCount);
+
+    // Submit a burst of tasks. These tasks will be queued up and the pool should scale up to handle
+    // them
+    for (int i = 0; i < taskCount; i++) {
+      telemetryPool.execute(
+          () -> {
+            uniqueThreadIds.add(Thread.currentThread().getId());
+            try {
+              Thread.sleep(10);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            } finally {
+              completionLatch.countDown();
+            }
+          });
+    }
+
+    // Wait for all tasks to finish processing.
+    completionLatch.await(5, TimeUnit.SECONDS);
+
+    assertEquals(
+        10,
+        uniqueThreadIds.size(),
+        "The thread pool should have used 10 threads, but it used " + uniqueThreadIds.size());
+    assertEquals(10, executor.getPoolSize());
+
+    executor.setKeepAliveTime(
+        1L, TimeUnit.MILLISECONDS); // Reset keep-alive time to 0 to allow threads to terminate.
+
+    await()
+        .atMost(5, TimeUnit.SECONDS) // Max time to wait for the condition
+        .pollInterval(100, TimeUnit.MILLISECONDS) // Check every 100ms
+        .untilAsserted(
+            () -> {
+              assertEquals(
+                  0,
+                  executor.getPoolSize(),
+                  "The thread pool should have scaled down to 0 idle threads.");
+            });
+  }
+
+  private ThreadPoolExecutor getThreadPoolExecutor(TelemetryThreadPool telemetryThreadPool) {
+    try {
+      Field uploaderField = TelemetryThreadPool.class.getDeclaredField("uploader");
+      uploaderField.setAccessible(true);
+      return (ThreadPoolExecutor) uploaderField.get(telemetryThreadPool);
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      throw new RuntimeException("Failed to access the uploader field in TelemetryThreadPool", e);
+    }
+  }
+}