From 57e40ddd7cf258361a5dde827256eaa972baf27d Mon Sep 17 00:00:00 2001 From: Vitaly Legchilkin Date: Mon, 10 Mar 2025 13:52:21 +0100 Subject: [PATCH] [python] add concurrency limit for env resolving (PY-60410) + change required background thread to explicit withContext(IO) where it is needed. GitOrigin-RevId: 79948000391765fcf0b8c24d1814a8dbe52d83a2 --- .../src/com/intellij/python/hatch/hatch.kt | 7 --- .../python/hatch/runtime/HatchRuntime.kt | 23 +++++-- .../hatch/service/CliBasedHatchService.kt | 62 ++++++++----------- 3 files changed, 42 insertions(+), 50 deletions(-) diff --git a/python/python-hatch/src/com/intellij/python/hatch/hatch.kt b/python/python-hatch/src/com/intellij/python/hatch/hatch.kt index 1629832e2f33..3a32c1b2ff9f 100644 --- a/python/python-hatch/src/com/intellij/python/hatch/hatch.kt +++ b/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 - @RequiresBackgroundThread suspend fun createNewProject(projectName: String): Result /** * 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 - @RequiresBackgroundThread suspend fun findVirtualEnvironments(): Result, PyError> } /** * Hatch Service for working directory (where hatch.toml / pyproject.toml is usually placed) */ -@RequiresBackgroundThread suspend fun getHatchService(workingDirectoryPath: Path, hatchExecutablePath: Path? = null): Result { 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 { val workingDirectoryPath = basePath?.let { Path.of(it) } ?: return Result.failure(WorkingDirectoryNotFoundHatchError(basePath)) diff --git a/python/python-hatch/src/com/intellij/python/hatch/runtime/HatchRuntime.kt b/python/python-hatch/src/com/intellij/python/hatch/runtime/HatchRuntime.kt index 6bd65ee8e298..8ee5bb0a91f0 100644 --- a/python/python-hatch/src/com/intellij/python/hatch/runtime/HatchRuntime.kt +++ b/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 executeInteractive(vararg arguments: String, eelProcessInteractiveHandler: EelProcessInteractiveHandler): Result { return execService.executeInteractive(hatchBinary, arguments.toList(), execOptions, eelProcessInteractiveHandler) } + + internal suspend fun resolvePythonVirtualEnvironment(pythonHomePath: PythonHomePath): Result { + 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 diff --git a/python/python-hatch/src/com/intellij/python/hatch/service/CliBasedHatchService.kt b/python/python-hatch/src/com/intellij/python/hatch/service/CliBasedHatchService.kt index efaec4aa7f32..8c56261244a4 100644 --- a/python/python-hatch/src/com/intellij/python/hatch/service/CliBasedHatchService.kt +++ b/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 Iterable.concurrentMap(f: suspend (A) -> B): List = coroutineScope { + map { + async { + concurrencyLimit.withPermit { f(it) } + } + }.awaitAll() + } } override fun getWorkingDirectoryPath(): Path = workingDirectoryPath - @RequiresBackgroundThread override suspend fun isHatchManagedProject(): Result { - 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 Iterable.parallelMap(f: suspend (A) -> B): List = coroutineScope { - map { async { f(it) } }.awaitAll() - } - @RequiresBackgroundThread override suspend fun findVirtualEnvironments(): Result, 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 { 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 { 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 { - 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 { val matricesFlatted = matrices.flatMap { matrixEnvironment -> matrixEnvironment.envs.map { envName ->