From 551779eae8268b3fa4539b69ea114029ef189b4c Mon Sep 17 00:00:00 2001
From: Sergei Vorobyov <sergey.vorobyov@jetbrains.com>
Date: Fri, 22 Nov 2024 14:27:45 +0100
Subject: [PATCH] [Gradle|Exec] fix: reuse already collected Gradle build
 environment for configuring operation for execution

Issues IDEA-363359, IDEA-344664

GitOrigin-RevId: 8da1caff1aeafeb44f3ff25f60eaa4996b320882
---
 .../gradle/service/execution/GradleExecutionHelper.java  | 6 ++----
 .../gradle/service/execution/GradleWrapperHelper.kt      | 9 ++++++++-
 .../service/modelAction/GradleModelFetchActionRunner.kt  | 4 ++--
 .../plugins/gradle/service/task/GradleTaskManager.java   | 2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
index 69644e6aefc9..7ef24d61f486 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
@@ -116,19 +116,17 @@ public final class GradleExecutionHelper {
 
   @ApiStatus.Internal
   public static void prepareForExecution(
-    @NotNull ProjectConnection connection,
     @NotNull LongRunningOperation operation,
     @NotNull CancellationToken cancellationToken,
     @NotNull ExternalSystemTaskId id,
     @NotNull GradleExecutionSettings settings,
-    @NotNull ExternalSystemTaskNotificationListener listener
+    @NotNull ExternalSystemTaskNotificationListener listener,
+    @Nullable BuildEnvironment buildEnvironment
   ) {
     clearSystemProperties(operation);
 
     applyIdeaParameters(settings);
 
-    BuildEnvironment buildEnvironment = getBuildEnvironment(connection, id, listener, null, settings);
-
     setupJvmArguments(operation, settings, buildEnvironment);
 
     setupLogging(settings, buildEnvironment);
diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleWrapperHelper.kt b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleWrapperHelper.kt
index da4c2f00b371..e9cd5f288cf1 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleWrapperHelper.kt
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleWrapperHelper.kt
@@ -133,11 +133,18 @@ object GradleWrapperHelper {
     cancellationToken: CancellationToken
   ) {
     SystemPropertiesAdjuster.executeAdjusted(projectPath) {
+
+      /**
+       * Don't reuse this build environment for the main execution process, because the wrapper task changes used Gradle distribution.
+       * It affects [org.gradle.tooling.model.build.GradleEnvironment] in [org.gradle.tooling.model.build.BuildEnvironment].
+       */
+      val buildEnvironment = GradleExecutionHelper.getBuildEnvironment(connection, id, listener, cancellationToken, settings)
+
       val launcher = connection.newBuild()
       val wrapperSettings = GradleExecutionSettings(settings).apply {
         tasks = listOf("wrapper")
       }
-      GradleExecutionHelper.prepareForExecution(connection, launcher, cancellationToken, id, wrapperSettings, listener)
+      GradleExecutionHelper.prepareForExecution(launcher, cancellationToken, id, wrapperSettings, listener, buildEnvironment)
       ExternalSystemTelemetryUtil.getTracer(GradleConstants.SYSTEM_ID)
         .spanBuilder("ExecuteWrapperTask")
         .use { launcher.run() }
diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/modelAction/GradleModelFetchActionRunner.kt b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/modelAction/GradleModelFetchActionRunner.kt
index 84814298e562..86daf9d5289e 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/modelAction/GradleModelFetchActionRunner.kt
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/modelAction/GradleModelFetchActionRunner.kt
@@ -88,12 +88,12 @@ class GradleModelFetchActionRunner private constructor(
 
   private fun <T : LongRunningOperation> T.prepareOperationForSync(): T {
     GradleExecutionHelper.prepareForExecution(
-      resolverContext.connection,
       this,
       resolverContext.cancellationToken,
       resolverContext.externalSystemTaskId,
       resolverContext.settings,
-      resolverContext.listener
+      resolverContext.listener,
+      resolverContext.buildEnvironment
     )
     return this
   }
diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/task/GradleTaskManager.java b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/task/GradleTaskManager.java
index f60e46b37794..8fdadb5eb729 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/task/GradleTaskManager.java
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/task/GradleTaskManager.java
@@ -160,7 +160,7 @@ public class GradleTaskManager implements ExternalSystemTaskManager<GradleExecut
       var operation = isApplicableTestLauncher(id, projectPath, settings, gradleVersion)
                       ? connection.newTestLauncher()
                       : connection.newBuild();
-      GradleExecutionHelper.prepareForExecution(connection, operation, cancellationToken, id, settings, listener);
+      GradleExecutionHelper.prepareForExecution(operation, cancellationToken, id, settings, listener, buildEnvironment);
       if (operation instanceof BuildLauncher) {
         ((BuildLauncher)operation).run();
       }