From 5e29b8361247fb6dc97fe885f5b6ebce21c6aab7 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 15 Jun 2026 11:13:03 -0400 Subject: [PATCH] add queryExceutor connection property --- .../bigquery/jdbc/BigQueryConnection.java | 16 ++++--- .../cloud/bigquery/jdbc/BigQueryJdbcMdc.java | 9 +++- .../bigquery/jdbc/BigQueryJdbcUrlUtility.java | 8 ++++ .../cloud/bigquery/jdbc/DataSource.java | 27 ++++++++++++ .../cloud/bigquery/jdbc/DataSourceTest.java | 43 +++++++++++++++++++ 5 files changed, 96 insertions(+), 7 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index cbbf50acc392..d6a3b4968467 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -174,6 +174,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { boolean useQueryCache; String queryDialect; int metadataFetchThreadCount; + int queryProcessThreadCount; boolean allowLargeResults; String destinationTable; String destinationDataset; @@ -340,6 +341,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.filterTablesOnDefaultDataset = ds.getFilterTablesOnDefaultDataset(); this.requestGoogleDriveScope = ds.getRequestGoogleDriveScope(); this.metadataFetchThreadCount = ds.getMetadataFetchThreadCount(); + this.queryProcessThreadCount = ds.getQueryProcessThreadCount(); this.requestReason = ds.getRequestReason(); this.connectionPoolSize = ds.getConnectionPoolSize(); this.listenerPoolSize = ds.getListenerPoolSize(); @@ -347,12 +349,12 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.headerProvider = createHeaderProvider(); this.bigQuery = getBigQueryConnection(); + // Fixed thread pool limits concurrent calls to prevent BigQuery metadata API throttling + // (independent tasks). this.metadataExecutor = BigQueryJdbcMdc.newFixedThreadPool(metadataFetchThreadCount); - // Use a bounded cached thread pool to prevent unbounded thread creation (and OOMs) - // under heavy load, while ensuring a limit (e.g., 100) high enough to prevent deadlocks - // between interdependent producer/consumer tasks (like nextPageWorker and - // populateBufferWorker). - this.queryExecutor = BigQueryJdbcMdc.newBoundedCachedThreadPool(100); + // Cached thread pool prevents deadlocks between interdependent producer/consumer query + // execution tasks. + this.queryExecutor = BigQueryJdbcMdc.newBoundedCachedThreadPool(queryProcessThreadCount); } } @@ -708,6 +710,10 @@ int getMetadataFetchThreadCount() { return this.metadataFetchThreadCount; } + int getQueryProcessThreadCount() { + return this.queryProcessThreadCount; + } + boolean isEnableWriteAPI() { return enableWriteAPI; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java index 61326e809d4a..0d0ee979a2af 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java @@ -62,6 +62,7 @@ static void clear() { static ExecutorService newFixedThreadPool(int nThreads, ThreadFactory threadFactory) { MdcThreadPoolExecutor executor = new MdcThreadPoolExecutor( + "Metadata Fetch Pool", nThreads, nThreads, 60L, @@ -86,6 +87,7 @@ static ExecutorService newFixedThreadPool(int nThreads) { */ static ExecutorService newBoundedCachedThreadPool(int maxThreads, ThreadFactory threadFactory) { return new MdcThreadPoolExecutor( + "Query Executor Pool", 0, maxThreads, 60L, @@ -125,8 +127,10 @@ public Thread newThread(Runnable r) { } private static class MdcThreadPoolExecutor extends ThreadPoolExecutor { + private final String poolName; public MdcThreadPoolExecutor( + String poolName, int corePoolSize, int maximumPoolSize, long keepAliveTime, @@ -134,6 +138,7 @@ public MdcThreadPoolExecutor( BlockingQueue workQueue, ThreadFactory threadFactory) { super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory); + this.poolName = poolName; } private final AtomicBoolean warningLogged = new AtomicBoolean(false); @@ -149,8 +154,8 @@ private void monitorQueueSaturation(int queueSize) { if (queueSize >= warnThreshold) { if (warningLogged.compareAndSet(false, true)) { LOG.warning( - "Thread pool is saturating. Max pool size: %d, Active threads: %d, Queued tasks: %d. Consider increasing the thread count property.", - maxPoolSize, getActiveCount(), queueSize); + "[%s] Thread pool is saturating. Max pool size: %d, Active threads: %d, Queued tasks: %d. Consider increasing the metadataFetchThreadCount property.", + poolName, maxPoolSize, getActiveCount(), queueSize); } } else if (queueSize <= recoveryThreshold) { if (warningLogged.get()) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java index 0a19bed7a2c8..b9211577bcdf 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java @@ -142,6 +142,8 @@ protected boolean removeEldestEntry(Map.Entry> eldes Pattern.CASE_INSENSITIVE); static final String METADATA_FETCH_THREAD_COUNT_PROPERTY_NAME = "MetaDataFetchThreadCount"; static final int DEFAULT_METADATA_FETCH_THREAD_COUNT_VALUE = 32; + static final String QUERY_PROCESS_THREAD_COUNT_PROPERTY_NAME = "QueryProcessThreadCount"; + static final int DEFAULT_QUERY_PROCESS_THREAD_COUNT_VALUE = 100; static final String RETRY_TIMEOUT_IN_SECS_PROPERTY_NAME = "Timeout"; static final long DEFAULT_RETRY_TIMEOUT_IN_SECS_VALUE = 0L; static final String JOB_TIMEOUT_PROPERTY_NAME = "JobTimeout"; @@ -540,6 +542,12 @@ protected boolean removeEldestEntry(Map.Entry> eldes "The number of threads used to call a DatabaseMetaData method.") .setDefaultValue(String.valueOf(DEFAULT_METADATA_FETCH_THREAD_COUNT_VALUE)) .build(), + BigQueryConnectionProperty.newBuilder() + .setName(QUERY_PROCESS_THREAD_COUNT_PROPERTY_NAME) + .setDescription( + "The maximum number of threads used to process query results concurrently.") + .setDefaultValue(String.valueOf(DEFAULT_QUERY_PROCESS_THREAD_COUNT_VALUE)) + .build(), BigQueryConnectionProperty.newBuilder() .setName(ENABLE_WRITE_API_PROPERTY_NAME) .setDescription( diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java index d7de950a935f..fed3fb53cdb5 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java @@ -86,6 +86,7 @@ public class DataSource implements javax.sql.DataSource { private Boolean filterTablesOnDefaultDataset; private Integer requestGoogleDriveScope; private Integer metadataFetchThreadCount; + private Integer queryProcessThreadCount; private String sslTrustStorePath; private String sslTrustStorePassword; private Map labels; @@ -248,6 +249,9 @@ public class DataSource implements javax.sql.DataSource { .put( BigQueryJdbcUrlUtility.METADATA_FETCH_THREAD_COUNT_PROPERTY_NAME, (ds, val) -> ds.setMetadataFetchThreadCount(Integer.parseInt(val))) + .put( + BigQueryJdbcUrlUtility.QUERY_PROCESS_THREAD_COUNT_PROPERTY_NAME, + (ds, val) -> ds.setQueryProcessThreadCount(Integer.parseInt(val))) .put( BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PROPERTY_NAME, DataSource::setSSLTrustStorePath) @@ -565,6 +569,11 @@ Properties createProperties() { BigQueryJdbcUrlUtility.METADATA_FETCH_THREAD_COUNT_PROPERTY_NAME, String.valueOf(this.metadataFetchThreadCount)); } + if (this.queryProcessThreadCount != null) { + connectionProperties.setProperty( + BigQueryJdbcUrlUtility.QUERY_PROCESS_THREAD_COUNT_PROPERTY_NAME, + String.valueOf(this.queryProcessThreadCount)); + } if (this.sslTrustStorePath != null) { connectionProperties.setProperty( BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PROPERTY_NAME, @@ -1085,6 +1094,24 @@ public void setMetadataFetchThreadCount(Integer metadataFetchThreadCount) { this.metadataFetchThreadCount = metadataFetchThreadCount; } + public Integer getQueryProcessThreadCount() { + return queryProcessThreadCount != null + ? queryProcessThreadCount + : BigQueryJdbcUrlUtility.DEFAULT_QUERY_PROCESS_THREAD_COUNT_VALUE; + } + + public void setQueryProcessThreadCount(Integer queryProcessThreadCount) { + if (queryProcessThreadCount != null) { + // Must be at least 4 to avoid thread starvation deadlocks between interdependent + // producer-consumer tasks (like nextPageWorker and populateBufferWorker). + validateMin( + queryProcessThreadCount, + 4, + BigQueryJdbcUrlUtility.QUERY_PROCESS_THREAD_COUNT_PROPERTY_NAME); + } + this.queryProcessThreadCount = queryProcessThreadCount; + } + public String getSSLTrustStorePath() { return sslTrustStorePath; } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/DataSourceTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/DataSourceTest.java index f6b82211c142..d7ae5466adfe 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/DataSourceTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/DataSourceTest.java @@ -83,4 +83,47 @@ public void testMetadataFetchThreadCountValidation() { .contains( "Invalid value for MetaDataFetchThreadCount. It must be greater than or equal to 1")); } + + @Test + public void testQueryProcessThreadCountValidation() { + DataSource dataSource = new DataSource(); + + // Setting values >= 4 should succeed + dataSource.setQueryProcessThreadCount(4); + assertEquals(4, dataSource.getQueryProcessThreadCount()); + + dataSource.setQueryProcessThreadCount(10); + assertEquals(10, dataSource.getQueryProcessThreadCount()); + + dataSource.setQueryProcessThreadCount(null); // Should fallback to default + assertEquals( + BigQueryJdbcUrlUtility.DEFAULT_QUERY_PROCESS_THREAD_COUNT_VALUE, + dataSource.getQueryProcessThreadCount()); + + // Setting value < 4 (e.g., 3, 0, -5) should throw BigQueryJdbcRuntimeException + BigQueryJdbcRuntimeException ex3 = + assertThrows( + BigQueryJdbcRuntimeException.class, () -> dataSource.setQueryProcessThreadCount(3)); + assertTrue( + ex3.getMessage() + .contains( + "Invalid value for QueryProcessThreadCount. It must be greater than or equal to 4")); + + BigQueryJdbcRuntimeException ex0 = + assertThrows( + BigQueryJdbcRuntimeException.class, () -> dataSource.setQueryProcessThreadCount(0)); + assertTrue( + ex0.getMessage() + .contains( + "Invalid value for QueryProcessThreadCount. It must be greater than or equal to 4")); + + BigQueryJdbcRuntimeException exNeg = + assertThrows( + BigQueryJdbcRuntimeException.class, () -> dataSource.setQueryProcessThreadCount(-5)); + assertTrue( + exNeg + .getMessage() + .contains( + "Invalid value for QueryProcessThreadCount. It must be greater than or equal to 4")); + } }