You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2026/10/15 09:20:43 UTC

(spark) branch branch-4.x updated: [SPARK-59938][CONNECT] Improve cancellation for server-side Maven artifact resolution

This is an automated email from the ASF dual-hosted git repository.

HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new d0530d7639d7 [SPARK-59938][CONNECT] Improve cancellation for server-side Maven artifact resolution
d0530d7639d7 is described below

commit d0530d7639d74b5a2765c04623106835387b72e8
Author: qianlan717 <qi...@databricks.com>
AuthorDate: Thu Oct 15 18:20:15 2026 +0900

    [SPARK-59938][CONNECT] Improve cancellation for server-side Maven artifact resolution
    
    ### What changes were proposed in this pull request?
    
    Add best-effort cancellation while runtime Maven dependency resolution is running for
    `SparkSession.addArtifact`, including server-side Spark Connect resolution.
    
    - Run Ivy resolution on a worker and poll the caller's cancellation state every 100 ms.
    - Record cancellation in a per-resolution flag, interrupt the worker directly, and check
      cancellation around resolution and retrieval.
    - Wait for the worker's cleanup and Ivy lock release before returning a cancellation error.
    - Preserve Ivy's existing URL handler and authentication behavior. Do not replace the HTTP
      transport or call `Ivy.interrupt()`.
    - Remove the HTTP client dependency that was needed only by the earlier transport-replacement
      implementation.
    
    The existing synchronous resolution entry points for `SparkContext.addJar` and Hive loading
    remain unchanged. No new timeout configuration or overall resolution deadline is introduced.
    
    ### Why are the changes needed?
    
    Master already has cancellation checks before and after runtime resolution and while waiting
    for the Ivy lock. It also applies `spark.jars.ivyConnectTimeout` and `spark.jars.ivyReadTimeout`
    to runtime repository requests. Those checks do not signal cancellation to an Ivy resolution
    that is already running.
    
    This change adds active interruption without reimplementing Ivy's downloader or risking
    authentication regressions from replacing its transport. Ivy operations that observe
    interruption, such as the artifact copy loop, can stop before completing the remaining work.
    The persistent cancellation flag also prevents a completed resolve or retrieve phase from
    continuing into the next phase after cancellation.
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes. Cancelling a runtime Maven dependency request now requests interruption of the active
    Ivy worker rather than only checking cancellation at resolution boundaries.
    
    Cancellation remains best-effort: a blocked read may not unblock on interruption, and some
    Ivy operations may ignore or consume the interrupt. The existing connection timeout
    (30 seconds by default) and read timeout (5 minutes by default) are reused. The read timeout
    is not a whole-download or whole-resolution deadline, and worker cleanup is not covered by
    those network timeouts. The caller waits for the worker, so prompt cancellation or termination
    within one read timeout is not guaranteed.
    
    ### How was this patch tested?
    
    Added or updated `MavenUtilsSuite` coverage for:
    
    - Cancellation before resolution and while waiting for the Ivy lock.
    - A stalled HTTP request with a short read timeout, followed by a successful resolution.
    - Cancellation during a progressing, multi-chunk download without replacing a custom handler.
    - Waiting for an Ivy resolver that temporarily ignores interruption.
    - Preserving the registered URL handler during successful and cancelled resolutions.
    
    Fixture cleanup now releases blockers and awaits executor termination before resource
    teardown, while preserving an original failure if cleanup also fails.
    
    The PR also adds Spark Connect coverage for server-side Maven resolution and invalid Maven
    coordinates following an upload. The latter uses mocked registration and does not establish
    real-registry atomicity during an actual Ivy failure or cancellation.
    
    Local validation:
    
    - `git diff --check`, POM XML parsing, ASCII checks, and the changed-source line-length check
      passed.
    - `common-utils/Test/compile` passed with Java 17 and invocation-only overrides for the
      cached compiler bridge and artifact proxy. The initial attempt failed at dependency
      resolution because the default Maven hosts could not be resolved.
    - Runtime test suites have not been run on the PR branch.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5.2)
    
    Closes #59196 from qianlan717/qianlan-chen_data/SPARK-59938-maven-resolution.
    
    Authored-by: qianlan717 <qi...@databricks.com>
    Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
    (cherry picked from commit ebd1ef3964b0370a0a3e3dc231ff7deda1ef7fe2)
    Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
---
 .../scala/org/apache/spark/util/MavenUtils.scala   | 130 +++++-
 .../org/apache/spark/util/MavenUtilsSuite.scala    | 481 +++++++++++++++++++--
 sql/connect/server/pom.xml                         |   7 +
 .../connect/service/AddArtifactsHandlerSuite.scala | 141 +++++-
 4 files changed, 701 insertions(+), 58 deletions(-)

diff --git a/common/utils/src/main/scala/org/apache/spark/util/MavenUtils.scala b/common/utils/src/main/scala/org/apache/spark/util/MavenUtils.scala
index f7e37cc3ab68..7c655c2b5f79 100644
--- a/common/utils/src/main/scala/org/apache/spark/util/MavenUtils.scala
+++ b/common/utils/src/main/scala/org/apache/spark/util/MavenUtils.scala
@@ -21,7 +21,8 @@ import java.io.{File, IOException, PrintStream}
 import java.net.URI
 import java.text.ParseException
 import java.util.UUID
-import java.util.concurrent.{CancellationException, TimeUnit}
+import java.util.concurrent.{Callable, CancellationException, ExecutionException, FutureTask, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.locks.ReentrantLock
 
 import scala.jdk.CollectionConverters._
@@ -62,17 +63,28 @@ private[spark] object MavenUtils extends Logging {
     }
   }
 
-  private def acquireIvyLock(isCancelled: () => Boolean): Unit = {
-    def checkCancelled(): Unit = {
-      if (isCancelled() || Thread.currentThread().isInterrupted) {
-        throw new CancellationException("Maven dependency resolution was cancelled")
-      }
+  private class IvyResolutionCancellation {
+    private val requested = new AtomicBoolean(false)
+
+    def isCancelled: Boolean = requested.get()
+
+    def cancel(resolutionThread: Thread): Unit = {
+      requested.set(true)
+      resolutionThread.interrupt()
+    }
+  }
+
+  private def checkCancelled(isCancelled: () => Boolean): Unit = {
+    if (isCancelled() || Thread.currentThread().isInterrupted) {
+      throw new CancellationException("Maven dependency resolution was cancelled")
     }
+  }
 
-    checkCancelled()
+  private def acquireIvyLock(isCancelled: () => Boolean): Unit = {
+    checkCancelled(isCancelled)
     try {
       while (!ivyLock.tryLock(ivyLockPollIntervalMs, TimeUnit.MILLISECONDS)) {
-        checkCancelled()
+        checkCancelled(isCancelled)
       }
     } catch {
       case _: InterruptedException =>
@@ -80,7 +92,7 @@ private[spark] object MavenUtils extends Logging {
         throw new CancellationException("Maven dependency resolution was cancelled")
     }
     try {
-      checkCancelled()
+      checkCancelled(isCancelled)
     } catch {
       case e: Throwable =>
         ivyLock.unlock()
@@ -88,6 +100,65 @@ private[spark] object MavenUtils extends Logging {
     }
   }
 
+  private def runWithCancellation[T](isCancelled: () => Boolean)(
+      body: IvyResolutionCancellation => T): T = {
+    checkCancelled(isCancelled)
+    val cancellation = new IvyResolutionCancellation
+    val task = new FutureTask[T](new Callable[T] {
+      override def call(): T = body(cancellation)
+    })
+    val resolutionThread = new Thread(task, "spark-ivy-resolution")
+    resolutionThread.setDaemon(true)
+    resolutionThread.start()
+
+    def cancelAndWait(): Unit = {
+      var wasInterrupted = Thread.interrupted()
+      try {
+        cancellation.cancel(resolutionThread)
+        while (resolutionThread.isAlive) {
+          try {
+            resolutionThread.join()
+          } catch {
+            case _: InterruptedException => wasInterrupted = true
+          }
+        }
+      } finally {
+        if (wasInterrupted) {
+          Thread.currentThread().interrupt()
+        }
+      }
+    }
+
+    try {
+      var result: T = null.asInstanceOf[T]
+      var completed = false
+      while (!completed) {
+        checkCancelled(isCancelled)
+        try {
+          result = task.get(ivyLockPollIntervalMs, TimeUnit.MILLISECONDS)
+          completed = true
+        } catch {
+          case _: TimeoutException =>
+        }
+      }
+      checkCancelled(isCancelled)
+      result
+    } catch {
+      case _: InterruptedException =>
+        cancelAndWait()
+        Thread.currentThread().interrupt()
+        throw new CancellationException("Maven dependency resolution was cancelled")
+      case e: CancellationException =>
+        cancelAndWait()
+        throw e
+      case e: ExecutionException =>
+        throw e.getCause
+      case e: Throwable =>
+        cancelAndWait()
+        throw e
+    }
+  }
+
   // Exposed for testing
   // var printStream = SparkSubmit.printStream
 
@@ -524,14 +595,14 @@ private[spark] object MavenUtils extends Logging {
       exclusions: Seq[String] = Nil,
       isTest: Boolean = false)(
       implicit printStream: PrintStream): Seq[String] = {
-    resolveMavenCoordinatesWithCancellation(
+    resolveMavenCoordinatesInternal(
       coordinates,
       ivySettings,
       noCacheIvySettings,
       transitive,
       exclusions,
       isTest,
-      () => false)
+      None)
   }
 
   private[spark] def resolveMavenCoordinatesWithCancellation(
@@ -546,6 +617,32 @@ private[spark] object MavenUtils extends Logging {
     if (coordinates == null || coordinates.trim.isEmpty) {
       Nil
     } else {
+      runWithCancellation(isCancelled) { cancellation =>
+        resolveMavenCoordinatesInternal(
+          coordinates,
+          ivySettings,
+          noCacheIvySettings,
+          transitive,
+          exclusions,
+          isTest,
+          Some(cancellation))
+      }
+    }
+  }
+
+  private def resolveMavenCoordinatesInternal(
+      coordinates: String,
+      ivySettings: IvySettings,
+      noCacheIvySettings: Option[IvySettings],
+      transitive: Boolean,
+      exclusions: Seq[String],
+      isTest: Boolean,
+      cancellation: Option[IvyResolutionCancellation])(
+      implicit printStream: PrintStream): Seq[String] = {
+    if (coordinates == null || coordinates.trim.isEmpty) {
+      Nil
+    } else {
+      val isCancelled = () => cancellation.exists(_.isCancelled)
       acquireIvyLock(isCancelled)
       val sysOut = System.out
       // Default configuration name for ivy
@@ -573,6 +670,9 @@ private[spark] object MavenUtils extends Logging {
           val resolveOptions = new ResolveOptions
           resolveOptions.setTransitive(transitive)
           val retrieveOptions = new RetrieveOptions
+          if (cancellation.isDefined) {
+            retrieveOptions.setOverwriteMode(RetrieveOptions.OVERWRITEMODE_DIFFERENT)
+          }
           // Turn downloading and logging off for testing
           if (isTest) {
             resolveOptions.setDownload(false)
@@ -595,7 +695,9 @@ private[spark] object MavenUtils extends Logging {
             md.addExcludeRule(createExclusion(e + ":*", ivySettings, ivyConfName))
           }
           // resolve dependencies
+          checkCancelled(isCancelled)
           val rr: ResolveReport = ivy.resolve(md, resolveOptions)
+          checkCancelled(isCancelled)
           if (rr.hasError) {
             // SPARK-46302: When there are some corrupted jars in the local maven repo,
             // we try to continue without the cache
@@ -612,19 +714,25 @@ private[spark] object MavenUtils extends Logging {
 
               val noCacheIvy = Ivy.newInstance(noCacheIvySettings.get)
               withIvyContext(noCacheIvy) {
+                checkCancelled(isCancelled)
                 val noCacheRr = noCacheIvy.resolve(md, resolveOptions)
+                checkCancelled(isCancelled)
                 if (noCacheRr.hasError) {
                   throw new RuntimeException(noCacheRr.getAllProblemMessages.toString)
                 }
+                checkCancelled(isCancelled)
                 noCacheIvy.retrieve(
                   noCacheRr.getModuleDescriptor.getModuleRevisionId, retrieveOptions)
+                checkCancelled(isCancelled)
                 resolveDependencyPaths(noCacheRr.getArtifacts.toArray, packagesDirectory)
               }
             } else {
               throw new RuntimeException(rr.getAllProblemMessages.toString)
             }
           } else {
+            checkCancelled(isCancelled)
             ivy.retrieve(rr.getModuleDescriptor.getModuleRevisionId, retrieveOptions)
+            checkCancelled(isCancelled)
             resolveDependencyPaths(rr.getArtifacts.toArray, packagesDirectory)
           }
         }
diff --git a/common/utils/src/test/scala/org/apache/spark/util/MavenUtilsSuite.scala b/common/utils/src/test/scala/org/apache/spark/util/MavenUtilsSuite.scala
index b0e249b5a9f1..318221f61f73 100644
--- a/common/utils/src/test/scala/org/apache/spark/util/MavenUtilsSuite.scala
+++ b/common/utils/src/test/scala/org/apache/spark/util/MavenUtilsSuite.scala
@@ -17,23 +17,26 @@
 
 package org.apache.spark.util
 
-import java.io.{File, OutputStream, PrintStream}
-import java.net.URI
+import java.io.{File, IOException, OutputStream, PrintStream}
+import java.net.{InetAddress, InetSocketAddress, ServerSocket, Socket, SocketTimeoutException, URI}
 import java.nio.charset.StandardCharsets
-import java.nio.file.{Files, Paths}
-import java.util.concurrent.{CancellationException, CountDownLatch, Executors, TimeUnit}
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+import java.nio.file.{Files, Path, Paths}
+import java.nio.file.attribute.FileTime
+import java.util.concurrent.{CancellationException, CountDownLatch, Executors, ExecutorService, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference}
 
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.duration._
 import scala.jdk.CollectionConverters._
 
+import com.sun.net.httpserver.HttpServer
 import org.apache.ivy.core.IvyContext
 import org.apache.ivy.core.module.descriptor.{DependencyDescriptor, MDArtifact}
 import org.apache.ivy.core.resolve.{ResolveData, ResolvedModuleRevision}
 import org.apache.ivy.core.settings.IvySettings
 import org.apache.ivy.plugins.resolver.{AbstractResolver, ChainResolver, FileSystemResolver, IBiblioResolver}
+import org.apache.ivy.util.url.{URLHandlerDispatcher, URLHandlerRegistry}
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite // scalastyle:ignore funsuite
 
@@ -85,6 +88,12 @@ class MavenUtilsSuite
     settings
   }
 
+  private def shutdownExecutor(executor: ExecutorService): Unit = {
+    executor.shutdownNow()
+    assert(executor.awaitTermination(10, TimeUnit.SECONDS),
+      "Ivy resolution executor did not terminate")
+  }
+
   override def beforeEach(): Unit = {
     super.beforeEach()
     tempIvyPath = SparkFileUtils.createTempDir(namePrefix = "ivy").getAbsolutePath()
@@ -130,19 +139,31 @@ class MavenUtilsSuite
     }
   }
 
-  test("runtime resolver configures Ivy network timeouts") {
-    val settings = MavenUtils.buildIvySettings(None, Some(tempIvyPath))
+  test("runtime resolver configures network timeouts for default and custom Ivy settings") {
+    val settingsFile = Paths.get(tempIvyPath, "timeout-ivysettings.xml")
+    Files.writeString(
+      settingsFile,
+      s"""<ivysettings>
+         |  <settings defaultResolver="custom"/>
+         |  <resolvers>
+         |    <ibiblio name="custom" m2compatible="true"
+         |      root="${Paths.get(tempIvyPath).toUri}"/>
+         |  </resolvers>
+         |</ivysettings>""".stripMargin)
+    val settingsToTest = Seq(
+      MavenUtils.buildIvySettings(None, Some(tempIvyPath)),
+      MavenUtils.loadIvySettings(settingsFile.toString, None, Some(tempIvyPath)))
 
-    MavenUtils.setResolverTimeouts(settings, connectTimeoutMs = 1234, readTimeoutMs = 5678)
+    settingsToTest.foreach { settings =>
+      MavenUtils.setResolverTimeouts(settings, connectTimeoutMs = 1234, readTimeoutMs = 5678)
 
-    val resolvers = settings.getDefaultResolver
-      .asInstanceOf[ChainResolver]
-      .getResolvers
-      .asScala
-      .collect { case resolver: AbstractResolver => resolver }
-    assert(resolvers.nonEmpty)
-    assert(resolvers.forall(_.getTimeoutConstraint.getConnectionTimeout == 1234))
-    assert(resolvers.forall(_.getTimeoutConstraint.getReadTimeout == 5678))
+      val resolvers = settings.getResolvers.asScala.collect {
+        case resolver: AbstractResolver => resolver
+      }
+      assert(resolvers.nonEmpty)
+      assert(resolvers.forall(_.getTimeoutConstraint.getConnectionTimeout == 1234))
+      assert(resolvers.forall(_.getTimeoutConstraint.getReadTimeout == 5678))
+    }
   }
 
   test("add dependencies works correctly") {
@@ -205,6 +226,7 @@ class MavenUtilsSuite
       val resolver = new RuntimeDependencyResolver(
         ivySettingsPath = Some(settings.toString),
         ivyPath = Some(tempIvyPath))
+      val previousUrlHandler = URLHandlerRegistry.getDefault
 
       val resolved = resolver.resolve(
         URI.create(s"ivy://${main.toString}"),
@@ -213,6 +235,7 @@ class MavenUtilsSuite
 
       assert(resolved.exists(_.getFileName.toString.contains("my.runtime.lib_mylib-0.1")))
       assert(resolved.forall(_.startsWith(Paths.get(tempIvyPath))))
+      assert(URLHandlerRegistry.getDefault eq previousUrlHandler)
     }
   }
 
@@ -229,6 +252,24 @@ class MavenUtilsSuite
     assert(error.getMessage.contains("do not allow repositories"))
   }
 
+  test("cancellable Maven resolution unwraps worker failures") {
+    val missingRepo = new File(tempIvyPath, "missing-repo")
+    assert(missingRepo.mkdirs())
+
+    val error = intercept[RuntimeException] {
+      MavenUtils.resolveMavenCoordinatesWithCancellation(
+        "my.runtime.missing:mylib:0.1",
+        buildTestIvySettings(missingRepo, Paths.get(tempIvyPath, "missing-cache").toString),
+        noCacheIvySettings = None,
+        transitive = false,
+        exclusions = Nil,
+        isTest = true,
+        isCancelled = () => false)
+    }
+    assert(error.getMessage.contains("unresolved dependency"))
+    assert(error.getMessage.contains("my.runtime.missing#mylib;0.1"))
+  }
+
   test("runtime dependency resolver honors cancellation before resolution") {
     val resolver = new RuntimeDependencyResolver(None, Some(tempIvyPath))
 
@@ -242,14 +283,15 @@ class MavenUtilsSuite
   }
 
   test("runtime dependency resolver honors cancellation while waiting for Ivy") {
-    val executor = Executors.newFixedThreadPool(2)
-    implicit val executionContext: ExecutionContext =
-      ExecutionContext.fromExecutorService(executor)
-    try {
-      val main = MavenCoordinate("my.runtime.waiting", "mylib", "0.1")
-      IvyTestUtils.withRepository(main, None, None) { repo =>
+    val main = MavenCoordinate("my.runtime.waiting", "mylib", "0.1")
+    IvyTestUtils.withRepository(main, None, None) { repo =>
+      val executor = Executors.newFixedThreadPool(2)
+      implicit val executionContext: ExecutionContext =
+        ExecutionContext.fromExecutorService(executor)
+      val releaseFirst = new CountDownLatch(1)
+      val cancelled = new AtomicBoolean(false)
+      SparkErrorUtils.tryWithSafeFinally {
         val firstEnteredIvy = new CountDownLatch(1)
-        val releaseFirst = new CountDownLatch(1)
         val blockingPrintStream = new PrintStream(noOpOutputStream) {
           private val blocked = new AtomicBoolean(false)
           // scalastyle:off println
@@ -273,8 +315,6 @@ class MavenUtilsSuite
         }
 
         assert(firstEnteredIvy.await(5, TimeUnit.SECONDS))
-        val cancelled = new AtomicBoolean(false)
-        val cancellationChecks = new AtomicInteger(0)
         val waitingForIvy = new CountDownLatch(1)
         val resolver = new RuntimeDependencyResolver(
           ivySettingsPath = None,
@@ -285,7 +325,15 @@ class MavenUtilsSuite
             resolverTimeoutMs,
             resolverTimeoutMs,
             isCancelled = () => {
-              if (cancellationChecks.incrementAndGet() >= 5) {
+              val workerWaiting = Thread.getAllStackTraces.asScala.exists { case (worker, trace) =>
+                worker.getName == "spark-ivy-resolution" &&
+                  worker.getState == Thread.State.TIMED_WAITING &&
+                  trace.exists { frame =>
+                    frame.getClassName == MavenUtils.getClass.getName &&
+                      frame.getMethodName == "acquireIvyLock"
+                  }
+              }
+              if (workerWaiting) {
                 waitingForIvy.countDown()
               }
               cancelled.get()
@@ -302,9 +350,348 @@ class MavenUtilsSuite
           releaseFirst.countDown()
         }
         SparkThreadUtils.awaitResultNoSparkExceptionConversion(first, 5.seconds)
+      } {
+        cancelled.set(true)
+        releaseFirst.countDown()
+        shutdownExecutor(executor)
+      }
+    }
+  }
+
+  test("runtime dependency resolver cancels stalled reads using the read timeout") {
+    val executor = Executors.newFixedThreadPool(2)
+    implicit val executionContext: ExecutionContext =
+      ExecutionContext.fromExecutorService(executor)
+    val server = new ServerSocket(0, 1, InetAddress.getByName("127.0.0.1"))
+    server.setSoTimeout(100)
+    val requestAccepted = new CountDownLatch(1)
+    val releaseRequest = new CountDownLatch(1)
+    val cancelled = new AtomicBoolean(false)
+    val serverFuture = Future {
+      val sockets = ArrayBuffer.empty[Socket]
+      try {
+        while (releaseRequest.getCount > 0) {
+          try {
+            sockets += server.accept()
+            requestAccepted.countDown()
+          } catch {
+            case _: SocketTimeoutException =>
+          }
+        }
+      } finally {
+        sockets.foreach(_.close())
+      }
+    }
+    SparkErrorUtils.tryWithSafeFinally {
+      val settings = Paths.get(tempIvyPath, "cancellable-ivysettings.xml")
+      Files.writeString(
+        settings,
+        s"""<ivysettings>
+           |  <settings defaultResolver="blocked"/>
+           |  <resolvers>
+           |    <ibiblio name="blocked" m2compatible="true"
+           |      root="http://127.0.0.1:${server.getLocalPort}"/>
+           |  </resolvers>
+           |</ivysettings>""".stripMargin)
+      val resolver = new RuntimeDependencyResolver(
+        ivySettingsPath = Some(settings.toString),
+        ivyPath = Some(Paths.get(tempIvyPath, "cancelled").toString))
+      val previousUrlHandler = URLHandlerRegistry.getDefault
+      val resolution = Future {
+        resolver.resolve(
+          URI.create("ivy://my.runtime.active:mylib:0.1"),
+          resolverTimeoutMs,
+          readTimeoutMs = 250,
+          isCancelled = () => cancelled.get())
+      }
+
+      assert(requestAccepted.await(5, TimeUnit.SECONDS))
+      assert(URLHandlerRegistry.getDefault eq previousUrlHandler)
+      cancelled.set(true)
+      intercept[CancellationException] {
+        SparkThreadUtils.awaitResultNoSparkExceptionConversion(resolution, 10.seconds)
+      }
+      assert(URLHandlerRegistry.getDefault eq previousUrlHandler)
+
+      val next = MavenCoordinate("my.runtime.after.cancel", "mylib", "0.1")
+      IvyTestUtils.withRepository(next, None, None) { repo =>
+        val resolved = MavenUtils.resolveMavenCoordinates(
+          next.toString,
+          MavenUtils.buildIvySettings(Some(repo), Some(tempIvyPath)),
+          transitive = true,
+          isTest = true)
+        assert(resolved.nonEmpty)
+      }
+    } {
+      cancelled.set(true)
+      releaseRequest.countDown()
+      SparkErrorUtils.tryWithSafeFinally {
+        SparkThreadUtils.awaitResultNoSparkExceptionConversion(serverFuture, 5.seconds)
+      } {
+        SparkErrorUtils.tryWithSafeFinally {
+          server.close()
+        } {
+          shutdownExecutor(executor)
+        }
+      }
+    }
+  }
+
+  test("runtime resolver preserves the HTTP handler when cancelling a progressing download") {
+    val executor = Executors.newSingleThreadExecutor()
+    implicit val executionContext: ExecutionContext =
+      ExecutionContext.fromExecutorService(executor)
+    val server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0)
+    val downloadStarted = new CountDownLatch(1)
+    val releaseDownload = new CountDownLatch(1)
+    val cancelled = new AtomicBoolean(false)
+    val downloadSize = 1024 * 1024
+    val bytesSent = new AtomicInteger(0)
+    val previousUrlHandler = URLHandlerRegistry.getDefault
+    val customUrlHandler = new URLHandlerDispatcher
+    customUrlHandler.setDefault(previousUrlHandler)
+    URLHandlerRegistry.setDefault(customUrlHandler)
+    SparkErrorUtils.tryWithSafeFinally {
+      server.createContext("/", exchange => {
+        try {
+          if (!exchange.getRequestURI.getPath.endsWith(".jar")) {
+            exchange.sendResponseHeaders(404, -1)
+          } else if (exchange.getRequestMethod == "HEAD") {
+            exchange.getResponseHeaders.set("Content-Length", downloadSize.toString)
+            exchange.sendResponseHeaders(200, -1)
+          } else {
+            exchange.sendResponseHeaders(200, downloadSize)
+            val output = exchange.getResponseBody
+            val chunk = new Array[Byte](1024)
+            while (bytesSent.get() < downloadSize && releaseDownload.getCount > 0) {
+              output.write(chunk)
+              output.flush()
+              bytesSent.addAndGet(chunk.length)
+              downloadStarted.countDown()
+              releaseDownload.await(25, TimeUnit.MILLISECONDS)
+            }
+          }
+        } catch {
+          case _: IOException =>
+        } finally {
+          exchange.close()
+        }
+      })
+      server.start()
+      val settings = Paths.get(tempIvyPath, "streaming-ivysettings.xml")
+      Files.writeString(
+        settings,
+        s"""<ivysettings>
+           |  <settings defaultResolver="streaming"/>
+           |  <resolvers>
+           |    <ibiblio name="streaming" m2compatible="true" usepoms="false"
+           |      root="http://127.0.0.1:${server.getAddress.getPort}/"/>
+           |  </resolvers>
+           |</ivysettings>""".stripMargin)
+      val resolver = new RuntimeDependencyResolver(
+        ivySettingsPath = Some(settings.toString),
+        ivyPath = Some(Paths.get(tempIvyPath, "streaming").toString))
+      val resolution = Future {
+        resolver.resolve(
+          URI.create("ivy://my.runtime.streaming:mylib:0.1?transitive=false"),
+          resolverTimeoutMs,
+          resolverTimeoutMs,
+          isCancelled = () => cancelled.get())
+      }
+
+      assert(downloadStarted.await(5, TimeUnit.SECONDS))
+      assert(URLHandlerRegistry.getDefault eq customUrlHandler)
+      cancelled.set(true)
+      intercept[CancellationException] {
+        SparkThreadUtils.awaitResultNoSparkExceptionConversion(resolution, 5.seconds)
+      }
+      assert(bytesSent.get() < downloadSize)
+      assert(URLHandlerRegistry.getDefault eq customUrlHandler)
+    } {
+      cancelled.set(true)
+      releaseDownload.countDown()
+      SparkErrorUtils.tryWithSafeFinally {
+        SparkErrorUtils.tryWithSafeFinally {
+          server.stop(0)
+        } {
+          shutdownExecutor(executor)
+        }
+      } {
+        URLHandlerRegistry.setDefault(previousUrlHandler)
+      }
+    }
+  }
+
+  test("runtime dependency resolver waits for an uncooperative Ivy worker") {
+    val executor = Executors.newFixedThreadPool(2)
+    implicit val executionContext: ExecutionContext =
+      ExecutionContext.fromExecutorService(executor)
+    val resolverEntered = new CountDownLatch(1)
+    val releaseResolver = new CountDownLatch(1)
+    val resolverCompleted = new AtomicBoolean(false)
+    val cancelled = new AtomicBoolean(false)
+    SparkErrorUtils.tryWithSafeFinally {
+      val blockingResolver = new FileSystemResolver {
+        override def getDependency(
+            _descriptor: DependencyDescriptor,
+            _data: ResolveData): ResolvedModuleRevision = {
+          resolverEntered.countDown()
+          var released = false
+          while (!released) {
+            try {
+              released = releaseResolver.await(100, TimeUnit.MILLISECONDS)
+            } catch {
+              case _: InterruptedException =>
+            }
+          }
+          resolverCompleted.set(true)
+          null
+        }
+      }
+      blockingResolver.setName("blocking")
+      val settings = MavenUtils.buildIvySettings(
+        remoteRepos = None,
+        ivyPath = Some(Paths.get(tempIvyPath, "uncooperative").toString))
+      settings.addResolver(blockingResolver)
+      settings.setDefaultResolver(blockingResolver.getName)
+
+      val resolution = Future {
+        MavenUtils.resolveMavenCoordinatesWithCancellation(
+          "my.runtime.uncooperative:mylib:0.1",
+          settings,
+          noCacheIvySettings = None,
+          transitive = true,
+          exclusions = Nil,
+          isTest = true,
+          isCancelled = () => cancelled.get())
+      }
+
+      assert(resolverEntered.await(5, TimeUnit.SECONDS))
+      cancelled.set(true)
+      val delayedRelease = Future {
+        try {
+          Thread.sleep(3000)
+        } finally {
+          releaseResolver.countDown()
+        }
+      }
+
+      intercept[CancellationException] {
+        SparkThreadUtils.awaitResultNoSparkExceptionConversion(resolution, 10.seconds)
+      }
+      assert(resolverCompleted.get())
+      SparkThreadUtils.awaitResultNoSparkExceptionConversion(delayedRelease, 5.seconds)
+    } {
+      cancelled.set(true)
+      releaseResolver.countDown()
+      shutdownExecutor(executor)
+    }
+  }
+
+  for {
+    cancellationTrigger <- Seq("caller interrupt", "interrupt flag", "cancelled predicate")
+    interruptDuringJoin <- Seq(false, true)
+  } {
+    test("cancellable Maven resolution restores caller interruption after joining the worker " +
+        s"(cancellationTrigger=$cancellationTrigger, " +
+        s"interruptDuringJoin=$interruptDuringJoin)") {
+      val resolverEntered = new CountDownLatch(1)
+      val resolverInterrupted = new CountDownLatch(1)
+      val releaseResolver = new CountDownLatch(1)
+      val callerWaiting = new CountDownLatch(1)
+      val releaseCaller = new CountDownLatch(1)
+      val workerThread = new AtomicReference[Thread]()
+      val callerFailure = new AtomicReference[Throwable]()
+      val callerInterrupted = new AtomicBoolean(false)
+      val blockingResolver = new FileSystemResolver {
+        override def getDependency(
+            _descriptor: DependencyDescriptor,
+            _data: ResolveData): ResolvedModuleRevision = {
+          workerThread.set(Thread.currentThread())
+          resolverEntered.countDown()
+          var released = false
+          while (!released) {
+            try {
+              released = releaseResolver.await(100, TimeUnit.MILLISECONDS)
+            } catch {
+              case _: InterruptedException => resolverInterrupted.countDown()
+            }
+          }
+          null
+        }
+      }
+      blockingResolver.setName("blocking")
+      val settings = MavenUtils.buildIvySettings(
+        remoteRepos = None,
+        ivyPath = Some(Paths.get(tempIvyPath, "caller-interrupted").toString))
+      settings.addResolver(blockingResolver)
+      settings.setDefaultResolver(blockingResolver.getName)
+      val caller = new Thread(() => {
+        try {
+          MavenUtils.resolveMavenCoordinatesWithCancellation(
+            "my.runtime.caller.interrupted:mylib:0.1",
+            settings,
+            noCacheIvySettings = None,
+            transitive = false,
+            exclusions = Nil,
+            isTest = true,
+            isCancelled = () => {
+              if (resolverEntered.getCount == 0) {
+                callerWaiting.countDown()
+                releaseCaller.await()
+                if (cancellationTrigger == "interrupt flag") {
+                  Thread.currentThread().interrupt()
+                }
+                cancellationTrigger == "cancelled predicate"
+              } else {
+                false
+              }
+            })
+        } catch {
+          case error: Throwable => callerFailure.set(error)
+        } finally {
+          callerInterrupted.set(Thread.currentThread().isInterrupted)
+        }
+      }, "ivy-interrupted-caller")
+
+      SparkErrorUtils.tryWithSafeFinally {
+        caller.start()
+        assert(resolverEntered.await(5, TimeUnit.SECONDS))
+        assert(callerWaiting.await(5, TimeUnit.SECONDS))
+        if (cancellationTrigger == "caller interrupt") {
+          caller.interrupt()
+        } else {
+          releaseCaller.countDown()
+        }
+        assert(resolverInterrupted.await(5, TimeUnit.SECONDS))
+        if (interruptDuringJoin) {
+          caller.interrupt()
+        }
+        caller.join(TimeUnit.SECONDS.toMillis(1))
+        assert(caller.isAlive, "Caller returned before the Ivy worker was released")
+        releaseResolver.countDown()
+        caller.join(TimeUnit.SECONDS.toMillis(5))
+
+        assert(!caller.isAlive)
+        assert(callerFailure.get().isInstanceOf[CancellationException])
+        val shouldBeInterrupted =
+          cancellationTrigger != "cancelled predicate" || interruptDuringJoin
+        assert(callerInterrupted.get() == shouldBeInterrupted)
+        assert(!workerThread.get().isAlive)
+      } {
+        releaseCaller.countDown()
+        releaseResolver.countDown()
+        caller.interrupt()
+        SparkErrorUtils.tryWithSafeFinally {
+          caller.join(TimeUnit.SECONDS.toMillis(10))
+          assert(!caller.isAlive, "Ivy resolution caller did not terminate")
+        } {
+          Option(workerThread.get()).foreach { worker =>
+            worker.join(TimeUnit.SECONDS.toMillis(10))
+            assert(!worker.isAlive, "Ivy resolution worker did not terminate")
+          }
+        }
       }
-    } finally {
-      executor.shutdownNow()
     }
   }
 
@@ -332,6 +719,44 @@ class MavenUtilsSuite
     }
   }
 
+  test("runtime dependency resolver replaces truncated retrieved jars") {
+    val main = MavenCoordinate("my.runtime.retrieve", "mylib", "0.1")
+    val localIvyHome = Paths.get(tempIvyPath, "local-ivy")
+    val runtimeIvyPath = Paths.get(tempIvyPath, "runtime-ivy")
+    IvyTestUtils.withRepository(
+      main,
+      dependencies = None,
+      rootDir = Some(localIvyHome.resolve("local").toFile),
+      useIvyLayout = true) { _ =>
+      val resolver = new RuntimeDependencyResolver(
+        ivySettingsPath = None,
+        ivyPath = Some(runtimeIvyPath.toString),
+        localIvyPath = Some(localIvyHome.toString))
+      def resolve(): Seq[Path] = {
+        resolver.resolve(
+          URI.create(s"ivy://${main.toString}?transitive=false"),
+          resolverTimeoutMs,
+          resolverTimeoutMs)
+      }
+
+      val resolved = resolve()
+      assert(resolved.size == 1)
+      val jar = resolved.head
+      val completeBytes = Files.readAllBytes(jar)
+      val lastModified = Files.getLastModifiedTime(jar)
+      Files.write(jar, completeBytes.take(completeBytes.length / 2))
+      val newerTime = FileTime.fromMillis(lastModified.toMillis + TimeUnit.MINUTES.toMillis(1))
+      Files.setLastModifiedTime(jar, newerTime)
+      assert(Files.size(jar) < completeBytes.length)
+      assert(Files.getLastModifiedTime(jar).toMillis > lastModified.toMillis)
+
+      assert(resolve() == resolved)
+      assert(Files.readAllBytes(jar).sameElements(completeBytes))
+      assert(resolve() == resolved)
+      assert(Files.readAllBytes(jar).sameElements(completeBytes))
+    }
+  }
+
   test("search for artifact at local repositories") {
     val main = new MavenCoordinate("my.great.lib", "mylib", "0.1")
     val dep = "my.great.dep:mydep:0.5"
diff --git a/sql/connect/server/pom.xml b/sql/connect/server/pom.xml
index 4b6180524be2..f99fa986bcb8 100644
--- a/sql/connect/server/pom.xml
+++ b/sql/connect/server/pom.xml
@@ -69,6 +69,13 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-common-utils_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <classifier>tests</classifier>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.apache.spark</groupId>
       <artifactId>spark-pipelines_${scala.binary.version}</artifactId>
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala
index da0b9ead6f43..587d050d8d3b 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala
@@ -37,13 +37,14 @@ import io.grpc.StatusRuntimeException
 import io.grpc.protobuf.StatusProto
 import io.grpc.stub.StreamObserver
 
-import org.apache.spark.SparkRuntimeException
+import org.apache.spark.{SparkConf, SparkRuntimeException}
 import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.{AddArtifactsRequest, AddArtifactsResponse}
 import org.apache.spark.sql.Artifact
 import org.apache.spark.sql.connect.ResourceHelper
 import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.util.{ThreadUtils, Utils}
+import org.apache.spark.util.{IvyTestUtils, MavenUtils, ThreadUtils, Utils}
+import org.apache.spark.util.MavenUtils.MavenCoordinate
 
 class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
 
@@ -51,6 +52,25 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
 
   private val sessionId = UUID.randomUUID.toString()
   private val sessionKey = SessionKey("c1", sessionId)
+  private val ivySettingsDir = Utils.createTempDir(namePrefix = "connect-ivy-settings")
+  private val mavenCoordinate = MavenCoordinate("my.connect.driver", "mylib", "0.1")
+  private val mavenRepository = IvyTestUtils.createLocalRepositoryForTests(
+    mavenCoordinate,
+    dependencies = None,
+    rootDir = Some(ivySettingsDir.toPath.resolve("repository").toFile))
+  private val ivySettings = ivySettingsDir.toPath.resolve("ivysettings.xml")
+  Files.writeString(
+    ivySettings,
+    s"""<ivysettings>
+       |  <settings defaultResolver="connect"/>
+       |  <resolvers>
+       |    <ibiblio name="connect" m2compatible="true" root="${mavenRepository.toURI}"/>
+       |  </resolvers>
+       |</ivysettings>""".stripMargin)
+
+  override protected def sparkConf: SparkConf = {
+    super.sparkConf.set(MavenUtils.JAR_IVY_SETTING_PATH_KEY, ivySettings.toString)
+  }
 
   override def beforeEach(): Unit = {
     super.beforeEach()
@@ -58,6 +78,14 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
     SparkConnectService.sessionManager.initializeBaseSession(() => spark.newSession())
   }
 
+  override def afterAll(): Unit = {
+    try {
+      super.afterAll()
+    } finally {
+      Utils.deleteRecursively(ivySettingsDir)
+    }
+  }
+
   class DummyStreamObserver(p: Promise[AddArtifactsResponse])
       extends StreamObserver[AddArtifactsResponse] {
     override def onNext(v: AddArtifactsResponse): Unit = p.success(v)
@@ -280,6 +308,23 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
     }
   }
 
+  private def uploadedArtifactEntry(
+      name: String,
+      path: Path): proto.AddArtifactsRequest.ArtifactEntry = {
+    val bytes = ByteString.copyFrom(Files.readAllBytes(path))
+    val crc = new CRC32()
+    crc.update(bytes.toByteArray)
+    val artifact = proto.AddArtifactsRequest.SingleChunkArtifact
+      .newBuilder()
+      .setName(name)
+      .setData(
+        proto.AddArtifactsRequest.ArtifactChunk
+          .newBuilder()
+          .setData(bytes)
+          .setCrc(crc.getValue))
+    proto.AddArtifactsRequest.ArtifactEntry.newBuilder().setArtifact(artifact).build()
+  }
+
   private def addChunkedArtifact(
       handler: SparkConnectAddArtifactsHandler,
       name: String,
@@ -351,21 +396,6 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
       new DummyStreamObserver(promise),
       resolvedMavenArtifacts = Map(ivyUri -> Seq(resolved)))
     try {
-      def uploaded(name: String, path: Path): proto.AddArtifactsRequest.ArtifactEntry = {
-        val bytes = ByteString.copyFrom(Files.readAllBytes(path))
-        val crc = new CRC32()
-        crc.update(bytes.toByteArray)
-        val artifact = proto.AddArtifactsRequest.SingleChunkArtifact
-          .newBuilder()
-          .setName(name)
-          .setData(
-            proto.AddArtifactsRequest.ArtifactChunk
-              .newBuilder()
-              .setData(bytes)
-              .setCrc(crc.getValue))
-        proto.AddArtifactsRequest.ArtifactEntry.newBuilder().setArtifact(artifact).build()
-      }
-
       val request = AddArtifactsRequest
         .newBuilder()
         .setSessionId(sessionId)
@@ -373,7 +403,7 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
         .setBatch(
           proto.AddArtifactsRequest.Batch
             .newBuilder()
-            .addEntries(uploaded(
+            .addEntries(uploadedArtifactEntry(
               "classes/smallClassFile.class",
               inputFilePath.resolve("smallClassFile.class")))
             .addEntries(
@@ -382,7 +412,8 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
                 .setMavenDependency(proto.AddArtifactsRequest.MavenDependency
                   .newBuilder()
                   .setUri(ivyUri.toString)))
-            .addEntries(uploaded("jars/smallJar.jar", inputFilePath.resolve("smallJar.jar")))
+            .addEntries(
+              uploadedArtifactEntry("jars/smallJar.jar", inputFilePath.resolve("smallJar.jar")))
             .build())
         .build()
 
@@ -408,6 +439,78 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper {
     }
   }
 
+  test("server-side Maven dependency uses the driver Ivy settings") {
+    val promise = Promise[AddArtifactsResponse]()
+    val handler = new TestAddArtifactsHandler(new DummyStreamObserver(promise))
+    try {
+      val request = AddArtifactsRequest
+        .newBuilder()
+        .setSessionId(sessionId)
+        .setUserContext(proto.UserContext.newBuilder().setUserId("c1"))
+        .setBatch(
+          proto.AddArtifactsRequest.Batch
+            .newBuilder()
+            .addEntries(
+              proto.AddArtifactsRequest.ArtifactEntry
+                .newBuilder()
+                .setMavenDependency(proto.AddArtifactsRequest.MavenDependency
+                  .newBuilder()
+                  .setUri(s"ivy://${mavenCoordinate.toString}"))))
+        .build()
+
+      handler.onNext(request)
+      handler.onCompleted()
+
+      val response = ThreadUtils.awaitResult(promise.future, 5.seconds)
+      val jarName =
+        s"jars/${mavenCoordinate.groupId}_${mavenCoordinate.artifactId}-" +
+          s"${mavenCoordinate.version}.jar"
+      assert(handler.getFinalArtifacts == Seq(jarName))
+      assert(response.getArtifactsCount == 1)
+      assert(response.getArtifacts(0).getName == jarName)
+      assert(response.getArtifacts(0).getIsCrcSuccessful)
+    } finally {
+      handler.forceCleanUp()
+    }
+  }
+
+  test("failing Maven dependency does not register a preceding upload") {
+    val error = Promise[Throwable]()
+    val observer = new StreamObserver[AddArtifactsResponse] {
+      override def onNext(value: AddArtifactsResponse): Unit = {}
+      override def onError(throwable: Throwable): Unit = error.success(throwable)
+      override def onCompleted(): Unit = {}
+    }
+    val handler = new TestAddArtifactsHandler(observer)
+    try {
+      val request = AddArtifactsRequest
+        .newBuilder()
+        .setSessionId(sessionId)
+        .setUserContext(proto.UserContext.newBuilder().setUserId("c1"))
+        .setBatch(
+          proto.AddArtifactsRequest.Batch
+            .newBuilder()
+            .addEntries(uploadedArtifactEntry(
+              "classes/smallClassFile.class",
+              inputFilePath.resolve("smallClassFile.class")))
+            .addEntries(
+              proto.AddArtifactsRequest.ArtifactEntry
+                .newBuilder()
+                .setMavenDependency(proto.AddArtifactsRequest.MavenDependency
+                  .newBuilder()
+                  .setUri("ivy://invalid-coordinate"))))
+        .build()
+
+      handler.onNext(request)
+      handler.onCompleted()
+
+      ThreadUtils.awaitResult(error.future, 5.seconds)
+      assert(handler.getFinalArtifacts.isEmpty)
+    } finally {
+      handler.forceCleanUp()
+    }
+  }
+
   test("server-side Maven dependencies reject requested repositories") {
     val error = Promise[Throwable]()
     val observer = new StreamObserver[AddArtifactsResponse] {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org