[python] add concurrency limit for env resolving (PY-60410)

+ change required background thread to explicit withContext(IO) where it is needed.

GitOrigin-RevId: 79948000391765fcf0b8c24d1814a8dbe52d83a2

python/python-hatch/src/com/intellij/python/hatch/hatch.kt

@@ -6,7 +6,6 @@ import com.intellij.openapi.util.NlsSafe
 import com.intellij.platform.eel.fs.EelFsError
 import com.intellij.python.hatch.cli.HatchEnvironment
 import com.intellij.python.hatch.service.CliBasedHatchService
-import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.jetbrains.python.PythonBinary
 import com.jetbrains.python.PythonHomePath
 import com.jetbrains.python.Result
@@ -67,27 +66,22 @@ data class ProjectStructure(
 interface HatchService {
   fun getWorkingDirectoryPath(): Path
 
-  @RequiresBackgroundThread
   suspend fun isHatchManagedProject(): Result<Boolean, PyError>
 
-  @RequiresBackgroundThread
   suspend fun createNewProject(projectName: String): Result<ProjectStructure, PyError>
 
   /**
    * param[basePythonBinaryPath] base python for environment, the one on the PATH should be used if null.
    * param[envName] environment name to create, 'default' should be used if null.
    */
-  @RequiresBackgroundThread
   suspend fun createVirtualEnvironment(basePythonBinaryPath: PythonBinary? = null, envName: String? = null): Result<PythonVirtualEnvironment.Existing, PyError>
 
-  @RequiresBackgroundThread
   suspend fun findVirtualEnvironments(): Result<List<HatchVirtualEnvironment>, PyError>
 }
 
 /**
  * Hatch Service for working directory (where hatch.toml / pyproject.toml is usually placed)
  */
-@RequiresBackgroundThread
 suspend fun getHatchService(workingDirectoryPath: Path, hatchExecutablePath: Path? = null): Result<HatchService, PyError> {
   return CliBasedHatchService(hatchExecutablePath = hatchExecutablePath, workingDirectoryPath = workingDirectoryPath)
 }
@@ -96,7 +90,6 @@ suspend fun getHatchService(workingDirectoryPath: Path, hatchExecutablePath: Pat
  * Hatch Service for Module.
  * Working directory considered as the module base path.
  */
-@RequiresBackgroundThread
 suspend fun Module.getHatchService(hatchExecutablePath: Path? = null): Result<HatchService, PyError> {
   val workingDirectoryPath = basePath?.let { Path.of(it) }
                              ?: return Result.failure(WorkingDirectoryNotFoundHatchError(basePath))

python/python-hatch/src/com/intellij/python/hatch/runtime/HatchRuntime.kt

@@ -3,20 +3,20 @@ package com.intellij.python.hatch.runtime
 import com.intellij.platform.eel.EelApi
 import com.intellij.platform.eel.provider.localEel
 import com.intellij.python.community.execService.*
-import com.intellij.python.hatch.BasePythonExecutableNotFoundHatchError
-import com.intellij.python.hatch.HatchConfiguration
-import com.intellij.python.hatch.HatchError
-import com.intellij.python.hatch.WorkingDirectoryNotFoundHatchError
+import com.intellij.python.community.execService.WhatToExec.Binary
+import com.intellij.python.hatch.*
 import com.intellij.python.hatch.cli.HatchCli
 import com.jetbrains.python.PythonBinary
+import com.jetbrains.python.PythonHomePath
 import com.jetbrains.python.Result
 import com.jetbrains.python.errorProcessing.PyError
+import com.jetbrains.python.resolvePythonBinary
 import java.nio.file.Path
 import kotlin.io.path.isDirectory
 import kotlin.io.path.isExecutable
 
 class HatchRuntime(
-  val hatchBinary: WhatToExec.Binary,
+  val hatchBinary: Binary,
   val execOptions: ExecOptions,
   private val execService: ExecService = ExecService(),
 ) {
@@ -58,6 +58,17 @@ class HatchRuntime(
   internal suspend fun <T> executeInteractive(vararg arguments: String, eelProcessInteractiveHandler: EelProcessInteractiveHandler<T>): Result<T, PyError.ExecException> {
     return execService.executeInteractive(hatchBinary, arguments.toList(), execOptions, eelProcessInteractiveHandler)
   }
+
+  internal suspend fun resolvePythonVirtualEnvironment(pythonHomePath: PythonHomePath): Result<PythonVirtualEnvironment, PyError> {
+    val pythonVersion = pythonHomePath.takeIf { it.isDirectory() }?.resolvePythonBinary()?.let { pythonBinaryPath ->
+      execService.execGetStdout(Binary(pythonBinaryPath), listOf("--version")).getOr { return it }.trim()
+    }
+    val pythonVirtualEnvironment = when {
+      pythonVersion == null -> PythonVirtualEnvironment.NotExisting(pythonHomePath)
+      else -> PythonVirtualEnvironment.Existing(pythonHomePath, pythonVersion)
+    }
+    return Result.success(pythonVirtualEnvironment)
+  }
 }
 
 
@@ -84,7 +95,7 @@ suspend fun createHatchRuntime(
   val actualEnvVars = defaultVariables + envVars
 
   val runtime = HatchRuntime(
-    hatchBinary = WhatToExec.Binary(actualHatchExecutable),
+    hatchBinary = Binary(actualHatchExecutable),
     execOptions = ExecOptions(
       env = actualEnvVars,
       workingDirectory = workingDirectoryPath

python/python-hatch/src/com/intellij/python/hatch/service/CliBasedHatchService.kt

@@ -7,23 +7,18 @@ import com.intellij.platform.eel.getOr
 import com.intellij.platform.eel.provider.asEelPath
 import com.intellij.platform.eel.provider.asNioPath
 import com.intellij.platform.eel.provider.getEelDescriptor
-import com.intellij.python.community.execService.ExecService
-import com.intellij.python.community.execService.WhatToExec.Binary
 import com.intellij.python.hatch.*
 import com.intellij.python.hatch.cli.ENV_TYPE_VIRTUAL
 import com.intellij.python.hatch.cli.HatchEnvironment
 import com.intellij.python.hatch.cli.HatchEnvironments
 import com.intellij.python.hatch.runtime.HatchRuntime
 import com.intellij.python.hatch.runtime.createHatchRuntime
-import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.jetbrains.python.PythonBinary
-import com.jetbrains.python.PythonHomePath
 import com.jetbrains.python.Result
 import com.jetbrains.python.errorProcessing.PyError
-import com.jetbrains.python.resolvePythonBinary
-import kotlinx.coroutines.async
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.*
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
 import java.nio.file.Path
 import kotlin.io.path.exists
 import kotlin.io.path.isDirectory
@@ -42,36 +37,43 @@ internal class CliBasedHatchService private constructor(
       ).getOr { return it }
       return Result.success(CliBasedHatchService(workingDirectoryPath, hatchRuntime))
     }
+
+    private val concurrencyLimit = Semaphore(permits = 5)
+
+    private suspend fun <A, B> Iterable<A>.concurrentMap(f: suspend (A) -> B): List<B> = coroutineScope {
+      map {
+        async {
+          concurrencyLimit.withPermit { f(it) }
+        }
+      }.awaitAll()
+    }
   }
 
   override fun getWorkingDirectoryPath(): Path = workingDirectoryPath
 
-  @RequiresBackgroundThread
   override suspend fun isHatchManagedProject(): Result<Boolean, PyError> {
-    val isHatchManaged = when {
-      workingDirectoryPath.resolve("hatch.toml").exists() -> true
-      else -> {
-        val pyProjectTomlPath = workingDirectoryPath.resolve("pyproject.toml").takeIf { it.isRegularFile() }
-        val hatchRegex = """^\[tool\.hatch\..+]$""".toRegex(RegexOption.MULTILINE)
-        pyProjectTomlPath?.readText()?.contains(hatchRegex) == true
+    val isHatchManaged = withContext(Dispatchers.IO) {
+      when {
+        workingDirectoryPath.resolve("hatch.toml").exists() -> true
+        else -> {
+          val pyProjectTomlPath = workingDirectoryPath.resolve("pyproject.toml").takeIf { it.isRegularFile() }
+          val hatchRegex = """^\[tool\.hatch\..+]$""".toRegex(RegexOption.MULTILINE)
+          pyProjectTomlPath?.readText()?.contains(hatchRegex) == true
+        }
       }
     }
     return Result.success(isHatchManaged)
   }
 
-  suspend fun <A, B> Iterable<A>.parallelMap(f: suspend (A) -> B): List<B> = coroutineScope {
-    map { async { f(it) } }.awaitAll()
-  }
 
-  @RequiresBackgroundThread
   override suspend fun findVirtualEnvironments(): Result<List<HatchVirtualEnvironment>, PyError> {
     val hatchEnv = hatchRuntime.hatchCli().env()
     val environments: HatchEnvironments = hatchEnv.show().getOr { return it }
     val virtualEnvironments = environments.getAvailableVirtualHatchEnvironments()
 
-    val available = virtualEnvironments.parallelMap { env ->
-      val pythonHomePath = hatchEnv.find(env.name).getOr { return@parallelMap null } ?: return@parallelMap null
-      val pythonVirtualEnvironment = pythonHomePath.toPythonVirtualEnvironment().getOr { return@parallelMap null }
+    val available = virtualEnvironments.concurrentMap { env ->
+      val pythonHomePath = hatchEnv.find(env.name).getOr { return@concurrentMap null } ?: return@concurrentMap null
+      val pythonVirtualEnvironment = hatchRuntime.resolvePythonVirtualEnvironment(pythonHomePath).getOr { return@concurrentMap null }
       HatchVirtualEnvironment(
         hatchEnvironment = env,
         pythonVirtualEnvironment = pythonVirtualEnvironment
@@ -82,7 +84,6 @@ internal class CliBasedHatchService private constructor(
   }
 
 
-  @RequiresBackgroundThread
   override suspend fun createNewProject(projectName: String): Result<ProjectStructure, PyError> {
     val eelApi = workingDirectoryPath.getEelDescriptor().upgrade()
     val tempDir = eelApi.fs.createTemporaryDirectory(EelFileSystemApi.CreateTemporaryEntryOptions.Builder().build()).getOr { failure ->
@@ -101,7 +102,6 @@ internal class CliBasedHatchService private constructor(
     ))
   }
 
-  @RequiresBackgroundThread
   override suspend fun createVirtualEnvironment(basePythonBinaryPath: PythonBinary?, envName: String?): Result<PythonVirtualEnvironment.Existing, PyError> {
     val pythonBasedRuntime = basePythonBinaryPath?.let { path ->
       hatchRuntime.withBasePythonBinaryPath(path).getOr { return it }
@@ -111,7 +111,7 @@ internal class CliBasedHatchService private constructor(
 
     hatchEnv.create(envName).getOr { return it }
     val pythonHomePath = hatchEnv.find(envName).getOr { return it }
-    val pythonVirtualEnvironment = pythonHomePath?.toPythonVirtualEnvironment()?.getOr { return it }
+    val pythonVirtualEnvironment = pythonHomePath?.let { hatchRuntime.resolvePythonVirtualEnvironment(it) }?.getOr { return it }
 
     val result = when (pythonVirtualEnvironment) {
       is PythonVirtualEnvironment.Existing -> Result.success(pythonVirtualEnvironment)
@@ -121,18 +121,6 @@ internal class CliBasedHatchService private constructor(
   }
 }
 
-@RequiresBackgroundThread
-internal suspend fun PythonHomePath.toPythonVirtualEnvironment(): Result<PythonVirtualEnvironment, PyError> {
-  val pythonVersion = this.takeIf { it.isDirectory() }?.resolvePythonBinary()?.let { pythonBinaryPath ->
-    ExecService().execGetStdout(Binary(pythonBinaryPath), listOf("--version")).getOr { return it }.trim()
-  }
-  val pythonVirtualEnvironment = when {
-    pythonVersion == null -> PythonVirtualEnvironment.NotExisting(this)
-    else -> PythonVirtualEnvironment.Existing(this, pythonVersion)
-  }
-  return Result.success(pythonVirtualEnvironment)
-}
-
 private fun HatchEnvironments.getAvailableVirtualHatchEnvironments(): List<HatchEnvironment> {
   val matricesFlatted = matrices.flatMap { matrixEnvironment ->
     matrixEnvironment.envs.map { envName ->