From d5107c7d6f8a8f3e2a0f4aa4958344787e491827 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Wed, 25 Mar 2026 13:15:54 -0700 Subject: [PATCH 1/6] Make gateway connection acquire timeout tunable via system property Add system property COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS and environment variable COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS to override the gateway path TCP connection timeout (default 45s). Follows the same pattern as COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS with system property -> env variable -> default fallback, validation for non-positive values, and NumberFormatException handling. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cosmos/implementation/ConfigsTests.java | 53 +++++++++++++++++++ .../azure/cosmos/implementation/Configs.java | 48 ++++++++++++++++- 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 9eb43cea33bb..d8ba45c41609 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -9,6 +9,7 @@ import org.testng.annotations.Test; import java.net.URI; +import java.time.Duration; import java.util.EnumSet; import static com.azure.cosmos.implementation.Configs.isThinClientEnabled; @@ -259,4 +260,56 @@ public void thinClientEndpointTest() { System.clearProperty("COSMOS.THINCLIENT_ENDPOINT"); } } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutDefaultTest() { + // Default connection acquire timeout should be 45 seconds + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutOverrideTest() { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "30"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(30)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutRejectsZeroAndNegative() { + // Zero should fall back to default (45s) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "0"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + + // Negative should fall back to default (45s) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "-1"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutRejectsNonNumeric() { + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "abc"); + try { + // Non-numeric should fall back to default (45s) + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 01e3b70ef646..a37cfaa3dd3d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -139,7 +139,9 @@ public class Configs { // Reactor Netty Constants private static final Duration MAX_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60); - private static final Duration CONNECTION_ACQUIRE_TIMEOUT = Duration.ofSeconds(45); + private static final int DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS = 45; + private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS = "COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"; + private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE = "COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"; private static final String REACTOR_NETTY_CONNECTION_POOL_NAME = "reactor-netty-connection-pool"; private static final int DEFAULT_HTTP_RESPONSE_TIMEOUT_IN_SECONDS = 60; private static final int DEFAULT_QUERY_PLAN_RESPONSE_TIMEOUT_IN_SECONDS = 5; @@ -594,7 +596,49 @@ public static Duration getMaxIdleConnectionTimeout() { } public static Duration getConnectionAcquireTimeout() { - return CONNECTION_ACQUIRE_TIMEOUT; + return Duration.ofSeconds(getConnectionAcquireTimeoutInSeconds()); + } + + public static int getConnectionAcquireTimeoutInSeconds() { + int value = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + + String valueFromSystemProperty = System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + try { + value = Integer.parseInt(valueFromSystemProperty); + } catch (NumberFormatException e) { + logger.warn( + "Invalid non-numeric value '{}' for system property {}. Falling back to environment variable or default.", + valueFromSystemProperty, + CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + valueFromSystemProperty = null; + } + } + + if (valueFromSystemProperty == null || valueFromSystemProperty.isEmpty()) { + String valueFromEnvVariable = System.getenv(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + try { + value = Integer.parseInt(valueFromEnvVariable); + } catch (NumberFormatException e) { + logger.warn( + "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}s.", + valueFromEnvVariable, + CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + } + } + } + + if (value <= 0) { + logger.warn( + "Invalid connection acquire timeout: {}s. Must be > 0. Falling back to default: {}s.", + value, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + return DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + } + + return value; } /** From a7700ce1ca261dc01a7e0b2cb9b6b20557acef48 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Wed, 25 Mar 2026 13:21:53 -0700 Subject: [PATCH 2/6] Simplify getConnectionAcquireTimeoutInSeconds using firstNonNull/emptyToNull pattern Address review feedback: use the same concise style as getBulkTransactionalBatchFlushIntervalInMs instead of verbose if/try/catch blocks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hpkScenarioAadManagedIdentity.scala | 157 ++++++++++++++++++ .../cosmos/implementation/ConfigsTests.java | 30 ---- .../azure/cosmos/implementation/Configs.java | 43 +---- 3 files changed, 161 insertions(+), 69 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala diff --git a/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala b/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala new file mode 100644 index 000000000000..30356c0d2f08 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala @@ -0,0 +1,157 @@ +// Databricks notebook source +println("SCENARIO: hpkScenarioAadManagedIdentity") + +val authType = "ManagedIdentity" +val cosmosEndpoint = "https://benchmark-cosmos-lx0.documents.azure.com:443/" +val subscriptionId = "b31b6408-0fb5-4688-9a3c-33ffb3983297" +val tenantId = "b5a53dd4-9da5-494f-adab-d04a5754fc6f" +val resourceGroupName = "lx-cosmos-benchmark" +val cosmosDatabaseName = "hpk-mi-testdb" +val cosmosContainerName = "hpk-mi-testcontainer" +val cosmosHpkContainerName = cosmosContainerName + "-hpk" + +val cfg = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.auth.type" -> authType, + "spark.cosmos.account.subscriptionId" -> subscriptionId, + "spark.cosmos.account.tenantId" -> tenantId, + "spark.cosmos.account.resourceGroupName" -> resourceGroupName, + "spark.cosmos.database" -> cosmosDatabaseName, + "spark.cosmos.container" -> cosmosContainerName, + "spark.cosmos.enforceNativeTransport" -> "true", + "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", +) + +val cfgWithAutoSchemaInference = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.auth.type" -> authType, + "spark.cosmos.account.subscriptionId" -> subscriptionId, + "spark.cosmos.account.tenantId" -> tenantId, + "spark.cosmos.account.resourceGroupName" -> resourceGroupName, + "spark.cosmos.database" -> cosmosDatabaseName, + "spark.cosmos.container" -> cosmosContainerName, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.enforceNativeTransport" -> "true", + "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", +) + +val cfgHpk = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.auth.type" -> authType, + "spark.cosmos.account.subscriptionId" -> subscriptionId, + "spark.cosmos.account.tenantId" -> tenantId, + "spark.cosmos.account.resourceGroupName" -> resourceGroupName, + "spark.cosmos.database" -> cosmosDatabaseName, + "spark.cosmos.container" -> cosmosHpkContainerName, + "spark.cosmos.enforceNativeTransport" -> "true", + "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", +) + +val cfgHpkWithAutoSchemaInference = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.auth.type" -> authType, + "spark.cosmos.account.subscriptionId" -> subscriptionId, + "spark.cosmos.account.tenantId" -> tenantId, + "spark.cosmos.account.resourceGroupName" -> resourceGroupName, + "spark.cosmos.database" -> cosmosDatabaseName, + "spark.cosmos.container" -> cosmosHpkContainerName, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.enforceNativeTransport" -> "true", + "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", +) + +// COMMAND ---------- + +// Section 1: Setup Catalog and create database + standard container +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI", "com.azure.cosmos.spark.CosmosCatalog") +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.accountEndpoint", cosmosEndpoint) +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.auth.type", authType) +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.subscriptionId", subscriptionId) +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.tenantId", tenantId) +spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.resourceGroupName", resourceGroupName) + +// create database +spark.sql(s"CREATE DATABASE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName};") + +// create standard container (single partition key on /id) +spark.sql(s"CREATE TABLE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName}.${cosmosContainerName} using cosmos.oltp " + + s"TBLPROPERTIES(partitionKeyPath = '/id', manualThroughput = '400')") + +// COMMAND ---------- + +// Section 2: Ingest data into standard container +spark.createDataFrame(Seq(("cat-alive", "Schrodinger cat", 2, true), ("cat-dead", "Schrodinger cat", 2, false))) + .toDF("id","name","age","isAlive") + .write + .format("cosmos.oltp") + .options(cfg) + .mode("APPEND") + .save() + +// COMMAND ---------- + +// Section 3: Read and query standard container +val df = spark.read.format("cosmos.oltp").options(cfgWithAutoSchemaInference).load() +df.printSchema() +df.show() + +// COMMAND ---------- + +import org.apache.spark.sql.functions.col + +// Query to find the live cat and increment age +df.filter(col("isAlive") === true) + .withColumn("age", col("age") + 1) + .show() + +// COMMAND ---------- + +// Section 4: Create container with Hierarchical Partition Key (HPK) +// HPK uses multiple partition key paths: /tenantId, /userId, /sessionId +spark.sql(s"CREATE TABLE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName} using cosmos.oltp " + + s"TBLPROPERTIES(partitionKeyPath = '/tenantId,/userId,/sessionId', manualThroughput = '400')") + +// COMMAND ---------- + +// Section 5: Ingest data into HPK container +spark.createDataFrame(Seq( + ("1", "tenant-A", "user-1", "session-100", "login", "2024-01-01"), + ("2", "tenant-A", "user-1", "session-101", "query", "2024-01-02"), + ("3", "tenant-A", "user-2", "session-200", "login", "2024-01-01"), + ("4", "tenant-B", "user-3", "session-300", "login", "2024-01-03"), + ("5", "tenant-B", "user-3", "session-300", "logout", "2024-01-03") +)) + .toDF("id", "tenantId", "userId", "sessionId", "action", "timestamp") + .write + .format("cosmos.oltp") + .options(cfgHpk) + .mode("APPEND") + .save() + +// COMMAND ---------- + +// Section 6: Read and query HPK container +val dfHpk = spark.read.format("cosmos.oltp").options(cfgHpkWithAutoSchemaInference).load() +dfHpk.printSchema() +dfHpk.show() + +// COMMAND ---------- + +import org.apache.spark.sql.functions.col + +// Query: filter by first level of hierarchical partition key (tenantId) +dfHpk.filter(col("tenantId") === "tenant-A").show() + +// Query: filter by first and second level (tenantId + userId) +dfHpk.filter(col("tenantId") === "tenant-A" && col("userId") === "user-1").show() + +// Query: filter by all three levels (tenantId + userId + sessionId) +dfHpk.filter(col("tenantId") === "tenant-B" && col("userId") === "user-3" && col("sessionId") === "session-300").show() + +// COMMAND ---------- + +// Section 7: Update throughput on HPK container +spark.sql(s"ALTER TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName} " + + s"SET TBLPROPERTIES('manualThroughput' = '1100')") + +// COMMAND ---------- + +// Section 8: Cleanup +spark.sql(s"DROP TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosContainerName};") +spark.sql(s"DROP TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName};") diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index d8ba45c41609..888becfe697d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -282,34 +282,4 @@ public void connectionAcquireTimeoutOverrideTest() { System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); } } - - @Test(groups = { "unit" }) - public void connectionAcquireTimeoutRejectsZeroAndNegative() { - // Zero should fall back to default (45s) - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "0"); - try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); - } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); - } - - // Negative should fall back to default (45s) - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "-1"); - try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); - } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); - } - } - - @Test(groups = { "unit" }) - public void connectionAcquireTimeoutRejectsNonNumeric() { - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "abc"); - try { - // Non-numeric should fall back to default (45s) - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); - } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); - } - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index a37cfaa3dd3d..2e828cf9556b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -600,45 +600,10 @@ public static Duration getConnectionAcquireTimeout() { } public static int getConnectionAcquireTimeoutInSeconds() { - int value = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; - - String valueFromSystemProperty = System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { - try { - value = Integer.parseInt(valueFromSystemProperty); - } catch (NumberFormatException e) { - logger.warn( - "Invalid non-numeric value '{}' for system property {}. Falling back to environment variable or default.", - valueFromSystemProperty, - CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - valueFromSystemProperty = null; - } - } - - if (valueFromSystemProperty == null || valueFromSystemProperty.isEmpty()) { - String valueFromEnvVariable = System.getenv(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE); - if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { - try { - value = Integer.parseInt(valueFromEnvVariable); - } catch (NumberFormatException e) { - logger.warn( - "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}s.", - valueFromEnvVariable, - CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE, - DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - } - } - } - - if (value <= 0) { - logger.warn( - "Invalid connection acquire timeout: {}s. Must be > 0. Falling back to default: {}s.", - value, - DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - return DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; - } - - return value; + return Integer.parseInt(System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS, + firstNonNull( + emptyToNull(System.getenv().get(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE)), + String.valueOf(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS)))); } /** From 0bc9906f89136a1a80da089bdf04260fc337f417 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Wed, 25 Mar 2026 13:25:43 -0700 Subject: [PATCH 3/6] Address review feedback: combine methods, add validation, fix tests - Combine getConnectionAcquireTimeoutInSeconds into getConnectionAcquireTimeout - Add NumberFormatException handling and >0 validation with fallback - Add zero/negative rejection tests matching thin client pattern - Remove accidentally committed scala file Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hpkScenarioAadManagedIdentity.scala | 157 ------------------ .../cosmos/implementation/ConfigsTests.java | 19 +++ .../azure/cosmos/implementation/Configs.java | 28 +++- 3 files changed, 40 insertions(+), 164 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala diff --git a/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala b/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala deleted file mode 100644 index 30356c0d2f08..000000000000 --- a/sdk/cosmos/azure-cosmos-spark_3/test-databricks/notebooks/hpkScenarioAadManagedIdentity.scala +++ /dev/null @@ -1,157 +0,0 @@ -// Databricks notebook source -println("SCENARIO: hpkScenarioAadManagedIdentity") - -val authType = "ManagedIdentity" -val cosmosEndpoint = "https://benchmark-cosmos-lx0.documents.azure.com:443/" -val subscriptionId = "b31b6408-0fb5-4688-9a3c-33ffb3983297" -val tenantId = "b5a53dd4-9da5-494f-adab-d04a5754fc6f" -val resourceGroupName = "lx-cosmos-benchmark" -val cosmosDatabaseName = "hpk-mi-testdb" -val cosmosContainerName = "hpk-mi-testcontainer" -val cosmosHpkContainerName = cosmosContainerName + "-hpk" - -val cfg = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.auth.type" -> authType, - "spark.cosmos.account.subscriptionId" -> subscriptionId, - "spark.cosmos.account.tenantId" -> tenantId, - "spark.cosmos.account.resourceGroupName" -> resourceGroupName, - "spark.cosmos.database" -> cosmosDatabaseName, - "spark.cosmos.container" -> cosmosContainerName, - "spark.cosmos.enforceNativeTransport" -> "true", - "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", -) - -val cfgWithAutoSchemaInference = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.auth.type" -> authType, - "spark.cosmos.account.subscriptionId" -> subscriptionId, - "spark.cosmos.account.tenantId" -> tenantId, - "spark.cosmos.account.resourceGroupName" -> resourceGroupName, - "spark.cosmos.database" -> cosmosDatabaseName, - "spark.cosmos.container" -> cosmosContainerName, - "spark.cosmos.read.inferSchema.enabled" -> "true", - "spark.cosmos.enforceNativeTransport" -> "true", - "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", -) - -val cfgHpk = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.auth.type" -> authType, - "spark.cosmos.account.subscriptionId" -> subscriptionId, - "spark.cosmos.account.tenantId" -> tenantId, - "spark.cosmos.account.resourceGroupName" -> resourceGroupName, - "spark.cosmos.database" -> cosmosDatabaseName, - "spark.cosmos.container" -> cosmosHpkContainerName, - "spark.cosmos.enforceNativeTransport" -> "true", - "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", -) - -val cfgHpkWithAutoSchemaInference = Map("spark.cosmos.accountEndpoint" -> cosmosEndpoint, - "spark.cosmos.auth.type" -> authType, - "spark.cosmos.account.subscriptionId" -> subscriptionId, - "spark.cosmos.account.tenantId" -> tenantId, - "spark.cosmos.account.resourceGroupName" -> resourceGroupName, - "spark.cosmos.database" -> cosmosDatabaseName, - "spark.cosmos.container" -> cosmosHpkContainerName, - "spark.cosmos.read.inferSchema.enabled" -> "true", - "spark.cosmos.enforceNativeTransport" -> "true", - "spark.cosmos.read.consistencyStrategy" -> "LatestCommitted", -) - -// COMMAND ---------- - -// Section 1: Setup Catalog and create database + standard container -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI", "com.azure.cosmos.spark.CosmosCatalog") -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.accountEndpoint", cosmosEndpoint) -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.auth.type", authType) -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.subscriptionId", subscriptionId) -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.tenantId", tenantId) -spark.conf.set(s"spark.sql.catalog.cosmosCatalogMI.spark.cosmos.account.resourceGroupName", resourceGroupName) - -// create database -spark.sql(s"CREATE DATABASE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName};") - -// create standard container (single partition key on /id) -spark.sql(s"CREATE TABLE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName}.${cosmosContainerName} using cosmos.oltp " + - s"TBLPROPERTIES(partitionKeyPath = '/id', manualThroughput = '400')") - -// COMMAND ---------- - -// Section 2: Ingest data into standard container -spark.createDataFrame(Seq(("cat-alive", "Schrodinger cat", 2, true), ("cat-dead", "Schrodinger cat", 2, false))) - .toDF("id","name","age","isAlive") - .write - .format("cosmos.oltp") - .options(cfg) - .mode("APPEND") - .save() - -// COMMAND ---------- - -// Section 3: Read and query standard container -val df = spark.read.format("cosmos.oltp").options(cfgWithAutoSchemaInference).load() -df.printSchema() -df.show() - -// COMMAND ---------- - -import org.apache.spark.sql.functions.col - -// Query to find the live cat and increment age -df.filter(col("isAlive") === true) - .withColumn("age", col("age") + 1) - .show() - -// COMMAND ---------- - -// Section 4: Create container with Hierarchical Partition Key (HPK) -// HPK uses multiple partition key paths: /tenantId, /userId, /sessionId -spark.sql(s"CREATE TABLE IF NOT EXISTS cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName} using cosmos.oltp " + - s"TBLPROPERTIES(partitionKeyPath = '/tenantId,/userId,/sessionId', manualThroughput = '400')") - -// COMMAND ---------- - -// Section 5: Ingest data into HPK container -spark.createDataFrame(Seq( - ("1", "tenant-A", "user-1", "session-100", "login", "2024-01-01"), - ("2", "tenant-A", "user-1", "session-101", "query", "2024-01-02"), - ("3", "tenant-A", "user-2", "session-200", "login", "2024-01-01"), - ("4", "tenant-B", "user-3", "session-300", "login", "2024-01-03"), - ("5", "tenant-B", "user-3", "session-300", "logout", "2024-01-03") -)) - .toDF("id", "tenantId", "userId", "sessionId", "action", "timestamp") - .write - .format("cosmos.oltp") - .options(cfgHpk) - .mode("APPEND") - .save() - -// COMMAND ---------- - -// Section 6: Read and query HPK container -val dfHpk = spark.read.format("cosmos.oltp").options(cfgHpkWithAutoSchemaInference).load() -dfHpk.printSchema() -dfHpk.show() - -// COMMAND ---------- - -import org.apache.spark.sql.functions.col - -// Query: filter by first level of hierarchical partition key (tenantId) -dfHpk.filter(col("tenantId") === "tenant-A").show() - -// Query: filter by first and second level (tenantId + userId) -dfHpk.filter(col("tenantId") === "tenant-A" && col("userId") === "user-1").show() - -// Query: filter by all three levels (tenantId + userId + sessionId) -dfHpk.filter(col("tenantId") === "tenant-B" && col("userId") === "user-3" && col("sessionId") === "session-300").show() - -// COMMAND ---------- - -// Section 7: Update throughput on HPK container -spark.sql(s"ALTER TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName} " + - s"SET TBLPROPERTIES('manualThroughput' = '1100')") - -// COMMAND ---------- - -// Section 8: Cleanup -spark.sql(s"DROP TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosContainerName};") -spark.sql(s"DROP TABLE cosmosCatalogMI.${cosmosDatabaseName}.${cosmosHpkContainerName};") diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 888becfe697d..09a4eedafecc 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -282,4 +282,23 @@ public void connectionAcquireTimeoutOverrideTest() { System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); } } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutRejectsZeroAndNegative() { + // Zero should fall back to default (45s) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "0"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + + // Negative should fall back to default (45s) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "-1"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 2e828cf9556b..ff32b5c01ad3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -596,14 +596,28 @@ public static Duration getMaxIdleConnectionTimeout() { } public static Duration getConnectionAcquireTimeout() { - return Duration.ofSeconds(getConnectionAcquireTimeoutInSeconds()); - } + int timeoutInSeconds; + try { + timeoutInSeconds = Integer.parseInt(System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS, + firstNonNull( + emptyToNull(System.getenv().get(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE)), + String.valueOf(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS)))); + } catch (NumberFormatException e) { + logger.warn( + "Invalid non-numeric value for connection acquire timeout. Falling back to default: {}s.", + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + } - public static int getConnectionAcquireTimeoutInSeconds() { - return Integer.parseInt(System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS, - firstNonNull( - emptyToNull(System.getenv().get(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE)), - String.valueOf(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS)))); + if (timeoutInSeconds <= 0) { + logger.warn( + "Invalid connection acquire timeout: {}s. Must be > 0. Falling back to default: {}s.", + timeoutInSeconds, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + } + + return Duration.ofSeconds(timeoutInSeconds); } /** From 279293bd409d6d7a7904859a47ab8b50d4a3971e Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Wed, 25 Mar 2026 13:58:46 -0700 Subject: [PATCH 4/6] Address review comments and update CHANGELOG - Refactor getConnectionAcquireTimeout() to align with getThinClientConnectionTimeoutInSeconds() pattern: read system property first, then env var, then default - with source-specific warning messages - Add unit test for non-numeric value fallback behavior - Add CHANGELOG entry for connection acquire timeout override Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cosmos/implementation/ConfigsTests.java | 11 ++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure/cosmos/implementation/Configs.java | 39 +++++++++++++------ 3 files changed, 40 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 09a4eedafecc..30ef6e904835 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -301,4 +301,15 @@ public void connectionAcquireTimeoutRejectsZeroAndNegative() { System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); } } + + @Test(groups = { "unit" }) + public void connectionAcquireTimeoutRejectsNonNumericValue() { + // Non-numeric value should fall back to default (45s) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "abc"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + } + } } diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 886a0fa8c86d..a077e312c349 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -21,6 +21,7 @@ * Added `appendUserAgentSuffix` method to `AsyncDocumentClient` to allow downstream libraries to append to the user agent after client construction. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505) * Added aggressive HTTP timeout policies for document operations routed to Gateway V2. - [PR 47879](https://github.com/Azure/azure-sdk-for-java/pull/47879) * Added a default connect timeout of 5s for Gateway V2 (thin client) data-plane endpoints. - See [PR 48174](https://github.com/Azure/azure-sdk-for-java/pull/48174) +* Added system property `COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS` and environment variable `COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS` to allow overriding the gateway connection acquire timeout (default 45s). - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) ### 4.78.0 (2026-02-10) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index ff32b5c01ad3..c7198dd2d1d1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -596,17 +596,34 @@ public static Duration getMaxIdleConnectionTimeout() { } public static Duration getConnectionAcquireTimeout() { - int timeoutInSeconds; - try { - timeoutInSeconds = Integer.parseInt(System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS, - firstNonNull( - emptyToNull(System.getenv().get(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE)), - String.valueOf(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS)))); - } catch (NumberFormatException e) { - logger.warn( - "Invalid non-numeric value for connection acquire timeout. Falling back to default: {}s.", - DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + int timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + + String valueFromSystemProperty = System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + try { + timeoutInSeconds = Integer.parseInt(valueFromSystemProperty); + } catch (NumberFormatException e) { + logger.warn( + "Invalid non-numeric value '{}' for system property {}. Falling back to environment variable or default.", + valueFromSystemProperty, + CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + valueFromSystemProperty = null; + } + } + + if (valueFromSystemProperty == null || valueFromSystemProperty.isEmpty()) { + String valueFromEnvVariable = System.getenv(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + try { + timeoutInSeconds = Integer.parseInt(valueFromEnvVariable); + } catch (NumberFormatException e) { + logger.warn( + "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}s.", + valueFromEnvVariable, + CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + } + } } if (timeoutInSeconds <= 0) { From 93b73f3f68fcd05c5802b0b97097b4dd35da5608 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 26 Mar 2026 08:29:35 -0700 Subject: [PATCH 5/6] Refactor connection timeout constants from seconds to milliseconds - Rename DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS (45) to DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS (45000) - Rename DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS (5) to DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS (5000) - Update system properties and env vars to use _IN_MS suffix - Rename getThinClientConnectionTimeoutInSeconds() to getThinClientConnectionTimeoutInMs() - Change minimum value validation from > 0 to >= 500 - Update callers, tests, and documentation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../CONNECT_TIMEOUT_TESTING_README.md | 6 +- .../Http2ConnectTimeoutBifurcationTests.java | 4 +- .../cosmos/implementation/ConfigsTests.java | 100 ++++++++++++------ .../azure/cosmos/implementation/Configs.java | 76 ++++++------- .../implementation/http/HttpClientConfig.java | 2 +- .../http/ReactorNettyClient.java | 2 +- 6 files changed, 111 insertions(+), 79 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/CONNECT_TIMEOUT_TESTING_README.md b/sdk/cosmos/azure-cosmos-tests/CONNECT_TIMEOUT_TESTING_README.md index 7895b856b331..03b7dfc31a75 100644 --- a/sdk/cosmos/azure-cosmos-tests/CONNECT_TIMEOUT_TESTING_README.md +++ b/sdk/cosmos/azure-cosmos-tests/CONNECT_TIMEOUT_TESTING_README.md @@ -20,7 +20,7 @@ Only `iptables DROP SYN` prevents the TCP handshake, triggering the real netty ` - Docker Desktop with Linux containers - Docker memory: **8 GB+** - A Cosmos DB account with thin client enabled -- System properties: `COSMOS.THINCLIENT_ENABLED=true`, `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS=5` (default; override to 1 for fast-fail local tests) +- System properties: `COSMOS.THINCLIENT_ENABLED=true`, `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS=5000` (default; override to 1000 for fast-fail local tests) - Credentials in `sdk/cosmos/cosmos-v4.properties` ## Build & Run @@ -36,7 +36,7 @@ docker run --rm --cap-add=NET_ADMIN --memory 8g \ cosmos-netem-test bash -c ' cd /workspace && \ java -DCOSMOS.THINCLIENT_ENABLED=true \ - -DCOSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS=1 \ + -DCOSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS=1000 \ -DCOSMOS.HTTP2_ENABLED=true \ org.testng.TestNG /workspace/azure-cosmos-tests/src/test/resources/manual-thinclient-network-delay-testng.xml \ -verbose 2 @@ -111,4 +111,4 @@ is required to route non-SYN traffic to band 3 (no delay). - Tests run **sequentially** — tc/iptables are interface-global - `--cap-add=NET_ADMIN` required for both `tc` and `iptables` - `@AfterClass` removes all iptables rules (`alwaysRun=true`) -- System property `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS=1` sets the 1s bifurcated timeout +- System property `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS=1000` sets the 1s bifurcated timeout diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/Http2ConnectTimeoutBifurcationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/Http2ConnectTimeoutBifurcationTests.java index e93b4cf56353..88665e2b5d6b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/Http2ConnectTimeoutBifurcationTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/Http2ConnectTimeoutBifurcationTests.java @@ -73,7 +73,7 @@ public Http2ConnectTimeoutBifurcationTests(CosmosClientBuilder clientBuilder) { @BeforeClass(groups = {TEST_GROUP}, timeOut = TIMEOUT) public void beforeClass() { System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); - // Use the default THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS (5s) — no override. + // Use the default THINCLIENT_CONNECTION_TIMEOUT_IN_MS (5000ms) — no override. // Tests are designed around the 5s default to match production behavior. this.client = getClientBuilder().buildAsyncClient(); @@ -394,7 +394,7 @@ public void connectTimeout_GwV1_Metadata_UnaffectedByGwV2Drop() throws Exception /** * Measures the precise connect timeout boundary. * - * With THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS=5 (default), a single TCP connect attempt + * With THINCLIENT_CONNECTION_TIMEOUT_IN_MS=5000 (default), a single TCP connect attempt * to a blackholed port 10250 should fail in ~5s. We measure multiple individual * attempts by using an e2e timeout of 12s (enough for 2 connect attempts at 5s each). * diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 30ef6e904835..a04a474faccf 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -183,42 +183,58 @@ public void thinClientEnabledTest() { @Test(groups = { "unit" }) public void thinClientConnectionTimeoutDefaultTest() { - // Default thin client connection timeout should be 5 seconds - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); + // Default thin client connection timeout should be 5000 milliseconds + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); try { - assertThat(Configs.getThinClientConnectionTimeoutInSeconds()).isEqualTo(5); + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(5_000); } finally { - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); } } @Test(groups = { "unit" }) public void thinClientConnectionTimeoutOverrideTest() { - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); - System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS", "3"); + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); + System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS", "3000"); try { - assertThat(Configs.getThinClientConnectionTimeoutInSeconds()).isEqualTo(3); + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(3_000); } finally { - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); } } @Test(groups = { "unit" }) - public void thinClientConnectionTimeoutRejectsZeroAndNegative() { - // Zero should fall back to default (5s) - System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS", "0"); + public void thinClientConnectionTimeoutRejectsInvalidValues() { + // Zero should fall back to default (5000ms) + System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS", "0"); try { - assertThat(Configs.getThinClientConnectionTimeoutInSeconds()).isEqualTo(5); + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(5_000); } finally { - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); } - // Negative should fall back to default (5s) - System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS", "-1"); + // Negative should fall back to default (5000ms) + System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS", "-1"); try { - assertThat(Configs.getThinClientConnectionTimeoutInSeconds()).isEqualTo(5); + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(5_000); } finally { - System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); + } + + // Below 500ms should fall back to default (5000ms) + System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS", "499"); + try { + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(5_000); + } finally { + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); + } + + // Exactly 500ms should be accepted + System.setProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS", "500"); + try { + assertThat(Configs.getThinClientConnectionTimeoutInMs()).isEqualTo(500); + } finally { + System.clearProperty("COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"); } } @@ -263,42 +279,58 @@ public void thinClientEndpointTest() { @Test(groups = { "unit" }) public void connectionAcquireTimeoutDefaultTest() { - // Default connection acquire timeout should be 45 seconds - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + // Default connection acquire timeout should be 45000 milliseconds + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(45_000)); } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); } } @Test(groups = { "unit" }) public void connectionAcquireTimeoutOverrideTest() { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "30"); + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS", "30000"); try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(30)); + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(30_000)); } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); } } @Test(groups = { "unit" }) - public void connectionAcquireTimeoutRejectsZeroAndNegative() { - // Zero should fall back to default (45s) - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "0"); + public void connectionAcquireTimeoutRejectsInvalidValues() { + // Zero should fall back to default (45000ms) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS", "0"); try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(45_000)); } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); } - // Negative should fall back to default (45s) - System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS", "-1"); + // Negative should fall back to default (45000ms) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS", "-1"); try { - assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofSeconds(45)); + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(45_000)); } finally { - System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"); + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); + } + + // Below 500ms should fall back to default (45000ms) + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS", "499"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(45_000)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); + } + + // Exactly 500ms should be accepted + System.setProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS", "500"); + try { + assertThat(Configs.getConnectionAcquireTimeout()).isEqualTo(Duration.ofMillis(500)); + } finally { + System.clearProperty("COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index c7198dd2d1d1..337055c6947f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -62,9 +62,9 @@ public class Configs { // Data plane requests are routed to the thin client regional endpoint (from RegionalRoutingContext) // which uses a non-443 port. These get a shorter 5s connect/acquire timeout. // Metadata requests target Gateway V1 endpoint (port 443) and retain the full 45s/60s timeout (unchanged). - private static final int DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS = 5; - private static final String THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS = "COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"; - private static final String THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS_VARIABLE = "COSMOS_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS"; + private static final int DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS = 5_000; + private static final String THINCLIENT_CONNECTION_TIMEOUT_IN_MS = "COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS"; + private static final String THINCLIENT_CONNECTION_TIMEOUT_IN_MS_VARIABLE = "COSMOS_THINCLIENT_CONNECTION_TIMEOUT_IN_MS"; private static final String MAX_HTTP_BODY_LENGTH_IN_BYTES = "COSMOS.MAX_HTTP_BODY_LENGTH_IN_BYTES"; private static final String MAX_HTTP_INITIAL_LINE_LENGTH_IN_BYTES = "COSMOS.MAX_HTTP_INITIAL_LINE_LENGTH_IN_BYTES"; @@ -139,9 +139,9 @@ public class Configs { // Reactor Netty Constants private static final Duration MAX_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60); - private static final int DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS = 45; - private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS = "COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"; - private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE = "COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS"; + private static final int DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS = 45_000; + private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_MS = "COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS"; + private static final String CONNECTION_ACQUIRE_TIMEOUT_IN_MS_VARIABLE = "COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_MS"; private static final String REACTOR_NETTY_CONNECTION_POOL_NAME = "reactor-netty-connection-pool"; private static final int DEFAULT_HTTP_RESPONSE_TIMEOUT_IN_SECONDS = 60; private static final int DEFAULT_QUERY_PLAN_RESPONSE_TIMEOUT_IN_SECONDS = 5; @@ -596,61 +596,61 @@ public static Duration getMaxIdleConnectionTimeout() { } public static Duration getConnectionAcquireTimeout() { - int timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + int timeoutInMs = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS; - String valueFromSystemProperty = System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + String valueFromSystemProperty = System.getProperty(CONNECTION_ACQUIRE_TIMEOUT_IN_MS); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { try { - timeoutInSeconds = Integer.parseInt(valueFromSystemProperty); + timeoutInMs = Integer.parseInt(valueFromSystemProperty); } catch (NumberFormatException e) { logger.warn( "Invalid non-numeric value '{}' for system property {}. Falling back to environment variable or default.", valueFromSystemProperty, - CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + CONNECTION_ACQUIRE_TIMEOUT_IN_MS); valueFromSystemProperty = null; } } if (valueFromSystemProperty == null || valueFromSystemProperty.isEmpty()) { - String valueFromEnvVariable = System.getenv(CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE); + String valueFromEnvVariable = System.getenv(CONNECTION_ACQUIRE_TIMEOUT_IN_MS_VARIABLE); if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { try { - timeoutInSeconds = Integer.parseInt(valueFromEnvVariable); + timeoutInMs = Integer.parseInt(valueFromEnvVariable); } catch (NumberFormatException e) { logger.warn( - "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}s.", + "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}ms.", valueFromEnvVariable, - CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS_VARIABLE, - DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); + CONNECTION_ACQUIRE_TIMEOUT_IN_MS_VARIABLE, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS); } } } - if (timeoutInSeconds <= 0) { + if (timeoutInMs < 500) { logger.warn( - "Invalid connection acquire timeout: {}s. Must be > 0. Falling back to default: {}s.", - timeoutInSeconds, - DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS); - timeoutInSeconds = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS; + "Invalid connection acquire timeout: {}ms. Must be >= 500. Falling back to default: {}ms.", + timeoutInMs, + DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS); + timeoutInMs = DEFAULT_CONNECTION_ACQUIRE_TIMEOUT_IN_MS; } - return Duration.ofSeconds(timeoutInSeconds); + return Duration.ofMillis(timeoutInMs); } /** - * Returns the TCP connect timeout for thin client data plane endpoints. + * Returns the TCP connect timeout for thin client data plane endpoints in milliseconds. * Data plane requests routed via thinclientRegionalEndpoint (from RegionalRoutingContext) * use this aggressive timeout to fail fast when the proxy is unreachable. * Metadata requests on port 443 are unaffected and retain the full 45s timeout. * - * Configurable via system property COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS - * or environment variable COSMOS_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS. - * Default: 5 seconds. + * Configurable via system property COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS + * or environment variable COSMOS_THINCLIENT_CONNECTION_TIMEOUT_IN_MS. + * Default: 5000 milliseconds. */ - public static int getThinClientConnectionTimeoutInSeconds() { - int value = DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS; + public static int getThinClientConnectionTimeoutInMs() { + int value = DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS; - String valueFromSystemProperty = System.getProperty(THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS); + String valueFromSystemProperty = System.getProperty(THINCLIENT_CONNECTION_TIMEOUT_IN_MS); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { try { value = Integer.parseInt(valueFromSystemProperty); @@ -658,33 +658,33 @@ public static int getThinClientConnectionTimeoutInSeconds() { logger.warn( "Invalid non-numeric value '{}' for system property {}. Falling back to environment variable or default.", valueFromSystemProperty, - THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS); + THINCLIENT_CONNECTION_TIMEOUT_IN_MS); valueFromSystemProperty = null; } } if (valueFromSystemProperty == null || valueFromSystemProperty.isEmpty()) { - String valueFromEnvVariable = System.getenv(THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS_VARIABLE); + String valueFromEnvVariable = System.getenv(THINCLIENT_CONNECTION_TIMEOUT_IN_MS_VARIABLE); if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { try { value = Integer.parseInt(valueFromEnvVariable); } catch (NumberFormatException e) { logger.warn( - "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}s.", + "Invalid non-numeric value '{}' for environment variable {}. Falling back to default: {}ms.", valueFromEnvVariable, - THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS_VARIABLE, - DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS); + THINCLIENT_CONNECTION_TIMEOUT_IN_MS_VARIABLE, + DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS); } } } - // Guard against invalid values — timeout must be at least 1 second - if (value <= 0) { + // Guard against invalid values — timeout must be at least 500ms + if (value < 500) { logger.warn( - "Invalid thin client connection timeout: {}s. Must be > 0. Falling back to default: {}s.", + "Invalid thin client connection timeout: {}ms. Must be >= 500. Falling back to default: {}ms.", value, - DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS); - return DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS; + DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS); + return DEFAULT_THINCLIENT_CONNECTION_TIMEOUT_IN_MS; } return value; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpClientConfig.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpClientConfig.java index 193abf876afc..7fa6ba8d3315 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpClientConfig.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpClientConfig.java @@ -40,7 +40,7 @@ public class HttpClientConfig { public HttpClientConfig(Configs configs) { this.configs = configs; this.http2ConnectionConfig = new Http2ConnectionConfig(); - this.thinClientConnectTimeoutMs = Configs.getThinClientConnectionTimeoutInSeconds() * 1000; + this.thinClientConnectTimeoutMs = Configs.getThinClientConnectionTimeoutInMs(); } public HttpClientConfig withMaxHeaderSize(int maxHeaderSize) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/ReactorNettyClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/ReactorNettyClient.java index 3e7f763caeb8..04ee87d22594 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/ReactorNettyClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/ReactorNettyClient.java @@ -257,7 +257,7 @@ public void shutdown() { * Resolves the TCP connect timeout (CONNECT_TIMEOUT_MILLIS) based on the request type. * * Thin client requests (identified by {@link HttpRequest#isThinClientRequest()}) use the thin - * client connection timeout configured via {@link Configs#getThinClientConnectionTimeoutInSeconds()} + * client connection timeout configured via {@link Configs#getThinClientConnectionTimeoutInMs()} * (default 5s) to fail fast when the thin client proxy is unreachable. * Standard gateway requests use the configured connection acquire timeout (default 45s). * From 05e46497e5442af4a869908eec4712caefd77bb4 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 26 Mar 2026 08:34:28 -0700 Subject: [PATCH 6/6] Update CHANGELOG for connection timeout millisecond refactor Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index a077e312c349..6153b2aacf74 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -21,7 +21,8 @@ * Added `appendUserAgentSuffix` method to `AsyncDocumentClient` to allow downstream libraries to append to the user agent after client construction. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505) * Added aggressive HTTP timeout policies for document operations routed to Gateway V2. - [PR 47879](https://github.com/Azure/azure-sdk-for-java/pull/47879) * Added a default connect timeout of 5s for Gateway V2 (thin client) data-plane endpoints. - See [PR 48174](https://github.com/Azure/azure-sdk-for-java/pull/48174) -* Added system property `COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS` and environment variable `COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_SECONDS` to allow overriding the gateway connection acquire timeout (default 45s). - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) +* Added system property `COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS` and environment variable `COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_MS` to allow overriding the gateway connection acquire timeout in milliseconds (default 45000ms). Minimum accepted value is 500ms. Replaces the previous `_IN_SECONDS` variants. - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) +* Changed system property for thin client connection timeout from `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS` to `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS` (default 5000ms, minimum 500ms). - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) ### 4.78.0 (2026-02-10)