Ensure toolchain is properly configured.

User specified toolchain maybe configured after initial task
configuration action will happen.

^KT-43095 In Progress

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/dsl/ToolchainDsl.kt

@@ -10,11 +10,10 @@ import org.gradle.api.GradleException
 import org.gradle.api.Project
 import org.gradle.api.plugins.ExtensionContainer
 import org.gradle.api.plugins.JavaPluginExtension
-import org.gradle.internal.jvm.Jvm
+import org.gradle.api.plugins.PluginContainer
-import org.gradle.jvm.toolchain.JavaLanguageVersion
+import org.gradle.api.tasks.TaskContainer
 import org.gradle.jvm.toolchain.JavaToolchainService
 import org.gradle.jvm.toolchain.JavaToolchainSpec
-import org.gradle.jvm.toolchain.internal.DefaultToolchainSpec
 import org.gradle.util.GradleVersion
 import org.jetbrains.kotlin.gradle.tasks.KotlinJavaToolchain.Companion.TOOLCHAIN_SUPPORTED_VERSION
 import org.jetbrains.kotlin.gradle.tasks.UsesKotlinJavaToolchain
@@ -32,33 +31,11 @@ internal interface ToolchainSupport {
             val currentVersion = GradleVersion.version(project.gradle.gradleVersion)
             return when {
                 currentVersion < TOOLCHAIN_SUPPORTED_VERSION -> project.objects.newInstance<NonExistingToolchainSupport>()
-                else -> project.objects.newInstance<DefaultToolchainSupport>(project.extensions).also {
+                else -> project.objects.newInstance<DefaultToolchainSupport>(
-                    configureDefaultToolchain(project)
+                    project.extensions,
-                }
+                    project.tasks,
-            }
+                    project.plugins
-        }
+                )
-
-        private fun configureDefaultToolchain(project: Project) {
-            project.plugins.withId("org.gradle.java-base") {
-                project
-                    .tasks
-                    .withType<UsesKotlinJavaToolchain>()
-                    .configureEach {
-                        val javaToolchainService = project
-                            .extensions
-                            .findByType(JavaToolchainService::class.java)!!
-                        val javaToolchainSpec = project
-                            .extensions
-                            .getByType(JavaPluginExtension::class.java)
-                            .toolchain
-
-                        // Only set when toolchain is configured
-                        if (javaToolchainSpec.languageVersion.isPresent) {
-                            it.kotlinJavaToolchain.toolchain.use(
-                                javaToolchainService.launcherFor(javaToolchainSpec)
-                            )
-                        }
-                    }
             }
         }
     }
@@ -73,16 +50,39 @@ internal abstract class NonExistingToolchainSupport : ToolchainSupport {
 }
 
 internal abstract class DefaultToolchainSupport @Inject constructor(
-    private val extensions: ExtensionContainer
+    private val extensions: ExtensionContainer,
+    private val tasks: TaskContainer,
+    private val plugins: PluginContainer
 ) : ToolchainSupport {
     private val toolchainSpec: JavaToolchainSpec
         get() = extensions
             .getByType(JavaPluginExtension::class.java)
             .toolchain
 
+    init {
+        configureToolchain()
+    }
+
     override fun applyToolchain(
         action: Action<Any>
     ) {
         action.execute(toolchainSpec)
+        configureToolchain()
+    }
+
+    private fun configureToolchain() {
+        plugins.withId("org.gradle.java-base") {
+            tasks
+                .withType<UsesKotlinJavaToolchain>()
+                .configureEach {
+                    // Only set when toolchain is configured
+                    if (toolchainSpec.languageVersion.isPresent) {
+                        val toolchainService = extensions.findByType(JavaToolchainService::class.java)!!
+                        it.kotlinJavaToolchain.toolchain.use(
+                            toolchainService.launcherFor(toolchainSpec)
+                        )
+                    }
+                }
+        }
     }
 }