Improved direct transport configuration in Spark to reduce number of transient failures under very high CPU/memory pressure on Spark Executors (Azure#27021)

* Release for Cosmos spark connector 4.6.1

* Fixing regression in TransientIOErrorsRetryingIterator possibly resulting in returning incomplete query results

* Fixing build issue

* Handling empty pages correctly

* Prettifying code

* Improve RNTBD config in Spark to reduce transient failures under high CPU/memory pressure

* Ignoring StructuredStreamingTests and fixing changelog

* Forcing to switch off of netty I/O thread immediately

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-spark_3-1_2-12/CHANGELOG.md

@@ -1,12 +1,13 @@
 ## Release History
-### 4.6.1 (2022-02-10)
+### 4.6.1 (2022-02-11)
 #### Key Bug Fixes
 * Fixed a regression introduced in 4.5.0 that could result in returning incomplete query results when Executors are under high CPU load. - See [PR 26991](https://github.com/Azure/azure-sdk-for-java/pull/26991)
 
 #### New Features
 * Added support for reading from a Cosmos table without schema inference and maintaining system properties `_ts` and `_etag` when writing this data to another Cosmos container. This is helpful when moving data from one container to another without always consistent schema of the documents in the source container. - See [PR 26820](https://github.com/Azure/azure-sdk-for-java/pull/26820)
 * Added support for new `spark.cosmos.write.strategy` value  `ItemOverwriteIfNotModified`, which will allow only updating documents that haven't been modified since reading them (optimistic concurrency). - See [PR 26847](https://github.com/Azure/azure-sdk-for-java/pull/26847)
 * Added support for correlating queries executed via the Cosmos Spark connector with service-telemetry based on the `correlationActivityId`. - See [PR 26908](https://github.com/Azure/azure-sdk-for-java/pull/26908)
+* Improved direct transport configuration in Spark to reduce number of transient failures under very high CPU/memory pressure on Spark Executors. - See [PR 27021](https://github.com/Azure/azure-sdk-for-java/pull/27021)
 
 ### 4.6.0 (2022-01-25)
 #### Key Bug Fixes

sdk/cosmos/azure-cosmos-spark_3-2_2-12/CHANGELOG.md

@@ -1,7 +1,8 @@
 ## Release History
-### 4.6.1 (2022-02-10)
+### 4.6.1 (2022-02-11)
 #### Key Bug Fixes
 * Fixed a regression introduced in 4.5.0 that could result in returning incomplete query results when Executors are under high CPU load. - See [PR 26991](https://github.com/Azure/azure-sdk-for-java/pull/26991)
+* Improved direct transport configuration in Spark to reduce number of transient failures under very high CPU/memory pressure on Spark Executors. - See [PR 27021](https://github.com/Azure/azure-sdk-for-java/pull/27021)
 
 #### New Features
 * Added support for reading from a Cosmos table without schema inference and maintaining system properties `_ts` and `_etag` when writing this data to another Cosmos container. This is helpful when moving data from one container to another without always consistent schema of the documents in the source container. - See [PR 26820](https://github.com/Azure/azure-sdk-for-java/pull/26820)

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala

@@ -172,6 +172,18 @@ private[cosmos] object SparkBridgeImplementationInternal {
       .setIoThreadCountPerCoreFactor(config, ioThreadCountPerCoreFactor)
   }
 
+  def setIoThreadPriority
+  (
+    config: DirectConnectionConfig,
+    ioThreadPriority: Int
+  ): DirectConnectionConfig = {
+
+    ImplementationBridgeHelpers
+      .DirectConnectionConfigHelper
+      .getDirectConnectionConfigAccessor
+      .setIoThreadPriority(config, ioThreadPriority)
+  }
+
   def setUserAgentWithSnapshotInsteadOfBeta() = {
     HttpConstants.Versions.useSnapshotInsteadOfBeta();
   }

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/CosmosClientCache.scala

@@ -23,6 +23,7 @@ import scala.collection.JavaConverters._
 private[spark] object CosmosClientCache extends BasicLoggingTrait {
 
   SparkBridgeImplementationInternal.setUserAgentWithSnapshotInsteadOfBeta()
+  System.setProperty("COSMOS.SWITCH_OFF_IO_THREAD_FOR_RESPONSE", "true")
 
   // removing clients from the cache after 15 minutes
   // The clients won't be disposed - so any still running task can still keep using it
@@ -114,17 +115,26 @@ private[spark] object CosmosClientCache extends BasicLoggingTrait {
         if (cosmosClientConfiguration.useGatewayMode){
           builder = builder.gatewayMode()
         } else {
-          val directConfig = new DirectConnectionConfig()
+          var directConfig = new DirectConnectionConfig()
             .setConnectTimeout(Duration.ofSeconds(CosmosConstants.defaultDirectRequestTimeoutInSeconds))
             .setNetworkRequestTimeout(Duration.ofSeconds(CosmosConstants.defaultDirectRequestTimeoutInSeconds))
 
-          builder = builder.directMode(
+          directConfig =
             // Duplicate the default number of I/O threads per core
             // We know that Spark often works with large payloads and we have seen
             // indicators that the default number of I/O threads can be too low
             // for workloads with large payloads
             SparkBridgeImplementationInternal
-              .setIoThreadCountPerCoreFactor(directConfig, CosmosConstants.defaultIoThreadCountFactorPerCore))
+              .setIoThreadCountPerCoreFactor(directConfig, CosmosConstants.defaultIoThreadCountFactorPerCore)
+
+          directConfig =
+          // Spark workloads often result in very high CPU load
+          // We have seen indicators that increasing Thread priority for I/O threads
+          // can reduce transient I/O errors/timeouts in this case
+            SparkBridgeImplementationInternal
+              .setIoThreadPriority(directConfig, Thread.MAX_PRIORITY)
+
+          builder = builder.directMode(directConfig)
         }
 
         if (cosmosClientConfiguration.preferredRegionsList.isDefined) {

sdk/cosmos/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/SparkE2EStructuredStreamingITest.scala

@@ -37,7 +37,7 @@ class SparkE2EStructuredStreamingITest
   //scalastyle:off multiple.string.literals
   //scalastyle:off magic.number
   "spark change feed micro batch (incremental)" can
-    "be used to copy data to another container" taggedAs(Retryable) in {
+    "be used to copy data to another container" taggedAs(Retryable) ignore {
 
     val processedRecordCount = new AtomicLong()
     var spark = this.createSparkSession(processedRecordCount)
@@ -177,7 +177,7 @@ class SparkE2EStructuredStreamingITest
   //scalastyle:off multiple.string.literals
   //scalastyle:off magic.number
   "spark change feed micro batch (incremental)" can
-    "be used to copy data to another container capturing origin TS and etag" taggedAs(Retryable) in {
+    "be used to copy data to another container capturing origin TS and etag" taggedAs(Retryable) ignore {
 
     val processedRecordCount = new AtomicLong()
     var spark = this.createSparkSession(processedRecordCount)
@@ -327,7 +327,7 @@ class SparkE2EStructuredStreamingITest
   }
 
   "spark change feed micro batch (incremental)" can
-    "be used to copy data to another container with limit" taggedAs(Retryable)  in {
+    "be used to copy data to another container with limit" taggedAs(Retryable) ignore {
 
     val processedRecordCount = new AtomicLong()
     var spark = this.createSparkSession(processedRecordCount)

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/DirectConnectionConfig.java

@@ -28,6 +28,7 @@ public final class DirectConnectionConfig {
     private static final int DEFAULT_MAX_CONNECTIONS_PER_ENDPOINT = 130;
     private static final int DEFAULT_MAX_REQUESTS_PER_CONNECTION = 30;
     private static final int DEFAULT_IO_THREAD_COUNT_PER_CORE_FACTOR = 2;
+    private static final int DEFAULT_IO_THREAD_PRIORITY = Thread.NORM_PRIORITY;
 
     private boolean connectionEndpointRediscoveryEnabled;
     private Duration connectTimeout;
@@ -37,6 +38,7 @@ public final class DirectConnectionConfig {
     private int maxConnectionsPerEndpoint;
     private int maxRequestsPerConnection;
     private int ioThreadCountPerCoreFactor;
+    private int ioThreadPriority;
 
     /**
      * Constructor
@@ -50,6 +52,7 @@ public DirectConnectionConfig() {
         this.maxRequestsPerConnection = DEFAULT_MAX_REQUESTS_PER_CONNECTION;
         this.networkRequestTimeout = DEFAULT_NETWORK_REQUEST_TIMEOUT;
         this.ioThreadCountPerCoreFactor = DEFAULT_IO_THREAD_COUNT_PER_CORE_FACTOR;
+        this.ioThreadPriority = DEFAULT_IO_THREAD_PRIORITY;
     }
 
     /**
@@ -286,6 +289,15 @@ DirectConnectionConfig setIoThreadCountPerCoreFactor(int ioThreadCountPerCoreFac
         return this;
     }
 
+    int getIoThreadPriority() {
+        return ioThreadPriority;
+    }
+
+    DirectConnectionConfig setIoThreadPriority(int ioThreadPriority) {
+        this.ioThreadPriority = ioThreadPriority;
+        return this;
+    }
+
     @Override
     public String toString() {
         return "DirectConnectionConfig{" +
@@ -296,6 +308,7 @@ public String toString() {
             ", maxRequestsPerConnection=" + maxRequestsPerConnection +
             ", networkRequestTimeout=" + networkRequestTimeout +
             ", ioThreadCountPerCoreFactor=" + ioThreadCountPerCoreFactor +
+            ", ioThreadPriority=" + ioThreadPriority +
             '}';
     }
 
@@ -316,6 +329,17 @@ public DirectConnectionConfig setIoThreadCountPerCoreFactor(DirectConnectionConf
                                                                             int ioThreadCountPerCoreFactor) {
                     return config.setIoThreadCountPerCoreFactor(ioThreadCountPerCoreFactor);
                 }
+
+                @Override
+                public int getIoThreadPriority(DirectConnectionConfig config) {
+                    return config.getIoThreadPriority();
+                }
+
+                @Override
+                public DirectConnectionConfig setIoThreadPriority(DirectConnectionConfig config,
+                                                                  int ioThreadPriority) {
+                    return config.setIoThreadPriority(ioThreadPriority);
+                }
             });
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ConnectionPolicy.java

@@ -45,7 +45,7 @@ public final class ConnectionPolicy {
     private Duration tcpNetworkRequestTimeout;
     private boolean tcpConnectionEndpointRediscoveryEnabled;
     private int ioThreadCountPerCoreFactor;
-
+    private int ioThreadPriority;
 
     private boolean clientTelemetryEnabled;
 
@@ -74,6 +74,10 @@ public ConnectionPolicy(DirectConnectionConfig directConnectionConfig) {
             .DirectConnectionConfigHelper
             .getDirectConnectionConfigAccessor()
             .getIoThreadCountPerCoreFactor(directConnectionConfig);
+        this.ioThreadPriority = ImplementationBridgeHelpers
+            .DirectConnectionConfigHelper
+            .getDirectConnectionConfigAccessor()
+            .getIoThreadPriority(directConnectionConfig);
     }
 
     private ConnectionPolicy(ConnectionMode connectionMode) {
@@ -85,6 +89,7 @@ private ConnectionPolicy(ConnectionMode connectionMode) {
         this.readRequestsFallbackEnabled = true;
         this.throttlingRetryOptions = new ThrottlingRetryOptions();
         this.userAgentSuffix = "";
+        this.ioThreadPriority = Thread.NORM_PRIORITY;
     }
 
     /**
@@ -547,11 +552,18 @@ public void setClientTelemetryEnabled(boolean clientTelemetryEnabled) {
 
     public int getIoThreadCountPerCoreFactor() { return this.ioThreadCountPerCoreFactor; }
 
+    public int getIoThreadPriority() { return this.ioThreadPriority; }
+
     public ConnectionPolicy setIoThreadCountPerCoreFactor(int ioThreadCountPerCoreFactor) {
         this.ioThreadCountPerCoreFactor = ioThreadCountPerCoreFactor;
         return this;
     }
 
+    public ConnectionPolicy setIoThreadPriority(int ioThreadPriority) {
+        this.ioThreadPriority = ioThreadPriority;
+        return this;
+    }
+
     @Override
     public String toString() {
         return "ConnectionPolicy{" +

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java

@@ -148,6 +148,9 @@ public interface DirectConnectionConfigAccessor {
             int getIoThreadCountPerCoreFactor(DirectConnectionConfig config);
             DirectConnectionConfig setIoThreadCountPerCoreFactor(
                 DirectConnectionConfig config, int ioThreadCountPerCoreFactor);
+            int getIoThreadPriority(DirectConnectionConfig config);
+            DirectConnectionConfig setIoThreadPriority(
+                DirectConnectionConfig config, int ioThreadPriority);
         }
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java

@@ -507,7 +507,7 @@ private Options(final ConnectionPolicy connectionPolicy) {
                 Runtime.getRuntime().availableProcessors();
             this.userAgent = new UserAgentContainer();
             this.channelAcquisitionContextEnabled = false;
-            this.ioThreadPriority = Thread.NORM_PRIORITY;
+            this.ioThreadPriority = connectionPolicy.getIoThreadPriority();
             this.tcpKeepIntvl = 1; // Configuration for EpollChannelOption.TCP_KEEPINTVL
             this.tcpKeepIdle = 30; // Configuration for EpollChannelOption.TCP_KEEPIDLE
             this.preferTcpNative = true;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdClientChannelHealthChecker.java

@@ -34,7 +34,9 @@ public final class RntbdClientChannelHealthChecker implements ChannelHealthCheck
     // Guidance: The grace period should be large enough to accommodate the round trip time of the slowest server
     // request. Assuming 1s of network RTT, a 2 MB request, a 2 MB response, a connection that can sustain 1 MB/s
     // both ways, and a 5-second deadline at the server, 10 seconds should be enough.
-    private static final long readHangGracePeriodInNanos = 10L * 1_000_000_000L;
+    // Adding an additional 45 seconds grace period because of relatively high number of
+    // false negatives here under high CPU load (in Spark for example)
+    private static final long readHangGracePeriodInNanos = (45L + 10L) * 1_000_000_000L;
 
     // A channel will not be declared unhealthy if a write was attempted recently. As such gaps between
     // Timestamps.lastChannelWriteAttempt and Timestamps.lastChannelWrite lower than this value are ignored.