Fix external reader process leak in pkl-server (#1885)

libpkl/libpkl.gradle.kts

@@ -48,6 +48,8 @@ val nativeTestRuntimeOnly: Configuration =
     extendsFrom(configurations.testRuntimeOnly.get())
   }
 
+val externalReaderFixtureConfiguration = configurations.create("externalReaderFixture")
+
 dependencies {
   compileOnly(libs.graalSdk)
 
@@ -62,6 +64,8 @@ dependencies {
   nativeTestImplementation(libs.jna)
   nativeTestImplementation(libs.jnaPlatform)
   nativeTestRuntimeOnly("org.junit.platform:junit-platform-launcher")
+
+  externalReaderFixtureConfiguration(project(":pkl-core", "externalReaderFixture"))
 }
 
 tasks.withType(CCompile::class) {
@@ -344,6 +348,7 @@ val processFiles =
 val testNativeJava =
   tasks.register<Test>("testNativeJava") {
     dependsOn(tasks.assembleNative)
+    dependsOn(externalReaderFixtureConfiguration)
 
     description = "Test native libraries from Java"
     group = "verification"
@@ -366,6 +371,11 @@ val testNativeJava =
     environment("LD_LIBRARY_PATH", buildInfo.targetMachine.libraryDir.get().asFile.absolutePath)
 
     useJUnitPlatform()
+
+    systemProperty(
+      "org.pkl.core.testExternalReaderPath",
+      externalReaderFixtureConfiguration.files.single().absolutePath,
+    )
   }
 
 // link to static library

pkl-commons-test/src/main/kotlin/org/pkl/commons/test/Executables.kt

@@ -19,6 +19,7 @@ import java.nio.file.Files
 import java.nio.file.Path
 import kotlin.io.path.exists
 import org.pkl.commons.test.FileTestUtils.rootProjectDir
+import org.pkl.core.util.IoUtils
 
 sealed class ExecutablePaths(protected val gradleProject: String) {
   abstract val allNative: List<Path>
@@ -80,4 +81,20 @@ object Executables {
     override val allNative: List<Path> =
       listOf(macAarch64, linuxAarch64, linuxAmd64, alpineAmd64, windowsAmd64)
   }
+
+  val externalReaderFixture by lazy {
+    val readerPath =
+      System.getProperty("org.pkl.core.testExternalReaderPath")
+        ?: "pkl-core/build/fixtures/externalreader"
+          .let { if (IoUtils.isWindows()) "$it.bat" else it }
+
+    rootProjectDir.resolve(readerPath).also { path ->
+      if (!Files.exists(path)) {
+        throw AssertionError(
+          "Fixture `${readerPath.substringAfterLast('/')}` not found. To fix this problem, first run" +
+            " `./gradlew pkl-core:externalReaderFixture`."
+        )
+      }
+    }
+  }
 }

pkl-commons-test/src/main/kotlin/org/pkl/commons/test/server/AbstractServerTest.kt

@@ -35,6 +35,7 @@ import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.io.TempDir
 import org.msgpack.core.MessagePack
+import org.pkl.commons.test.Executables
 import org.pkl.commons.test.PackageServer
 import org.pkl.commons.test.debugRendering
 import org.pkl.core.messaging.Messages.*
@@ -238,9 +239,6 @@ abstract class AbstractServerTest {
     assertThat(value.asStringValue().asString()).isEqualTo("my bahumbug")
   }
 
-  @Disabled(
-    "Unable to construct ReadResourceResponse with null contents due to Kotlin compiler bug"
-  )
   @Test
   fun `read resource -- null contents and null error`() {
     val reader = ResourceReaderSpec("bahumbug", true, false)
@@ -260,18 +258,10 @@ abstract class AbstractServerTest {
     assertThat(readResourceMsg.uri.toString()).isEqualTo("bahumbug:/foo.pkl")
     assertThat(readResourceMsg.evaluatorId).isEqualTo(evaluatorId)
 
-    client.send(ReadResourceResponse(readResourceMsg.requestId, evaluatorId, byteArrayOf(), null))
+    client.send(ReadResourceResponse(readResourceMsg.requestId, evaluatorId, null, null))
-    // for this test to be correct this should actually be:
-    // client.send(ReadResourceResponse(readResourceMsg.requestId, evaluatorId, null, null))
-    // this should be evaluated again once https://github.com/apple/pkl/issues/698 is addressed
-    // see conversation here https://github.com/apple/pkl/pull/660#discussion_r1819545811
 
     val evaluateResponse = client.receive<EvaluateResponse>()
-    assertThat(evaluateResponse.error).isNull()
+    assertThat(evaluateResponse.error).contains("Cannot find resource `bahumbug:/foo.pkl`.")
-
-    val unpacker = MessagePack.newDefaultUnpacker(evaluateResponse.result)
-    val value = unpacker.unpackValue()
-    assertThat(value.asStringValue().asString()).isEqualTo("")
   }
 
   @Test
@@ -1133,6 +1123,40 @@ abstract class AbstractServerTest {
       .contains("Rewrite rule must end with '/', but was 'https://example.com'")
   }
 
+  @Test
+  fun `test external reader`() {
+    val evaluatorId =
+      client.sendCreateEvaluatorRequest(
+        externalResourceReaders =
+          mapOf(
+            "foo" to ExternalReader(Executables.externalReaderFixture.toString(), emptyList(), null)
+          )
+      )
+    val requestId = 234L
+
+    client.send(
+      EvaluateRequest(
+        requestId,
+        evaluatorId,
+        URI("repl:text"),
+        """
+        res = read("foo:bar")
+        """
+          .trimIndent(),
+        "res.text",
+      )
+    )
+
+    val response = client.receive<EvaluateResponse>()
+    assertThat(response.error).isNull()
+    assertThat(response.result).isNotNull
+    assertThat(response.requestId()).isEqualTo(requestId)
+
+    val unpacker = MessagePack.newDefaultUnpacker(response.result)
+    val value = unpacker.unpackValue()
+    assertThat(value.asStringValue().asString()).isEqualTo("hello")
+  }
+
   protected fun TestTransport.sendCreateEvaluatorRequest(
     requestId: Long = 123,
     resourceReaders: List<ResourceReaderSpec> = listOf(),
@@ -1141,6 +1165,8 @@ abstract class AbstractServerTest {
     project: Project? = null,
     cacheDir: Path? = null,
     http: Http? = null,
+    externalResourceReaders: Map<String, ExternalReader>? = null,
+    externalModuleReaders: Map<String, ExternalReader>? = null,
   ): Long {
     val message =
       CreateEvaluatorRequest(
@@ -1158,8 +1184,8 @@ abstract class AbstractServerTest {
         null,
         project,
         http,
-        null,
+        externalModuleReaders,
-        null,
+        externalResourceReaders,
         null,
       )
 

pkl-core/pkl-core.gradle.kts

@@ -40,6 +40,8 @@ val externalReaderFixtureImplementation: Configuration =
     extendsFrom(configurations.testImplementation.get())
   }
 
+val externalReaderFixtureConfiguration = configurations.create("externalReaderFixture")
+
 idea {
   module {
     // mark generated/truffle as generated source dir
@@ -170,6 +172,13 @@ val externalReaderFixture =
     }
   }
 
+artifacts {
+  add(
+    externalReaderFixtureConfiguration.name,
+    externalReaderFixture.map { it.outputs.files.singleFile },
+  )
+}
+
 tasks.test {
   configureTest()
   dependsOn(externalReaderFixture)
@@ -191,7 +200,7 @@ tasks.test {
 
   systemProperty(
     "org.pkl.core.testExternalReaderPath",
-    externalReaderFixture.map { it.outputs.files.singleFile.absolutePath },
+    externalReaderFixture.map { it.outputs.files.singleFile.absolutePath }.get(),
   )
 }
 

pkl-core/src/test/kotlin/org/pkl/core/module/ModuleKeyFactoriesTest.kt

@@ -17,7 +17,6 @@ package org.pkl.core.module
 
 import java.io.File
 import java.net.URI
-import java.nio.file.Files
 import java.nio.file.Path
 import java.util.regex.Pattern
 import kotlin.io.path.createDirectories
@@ -27,7 +26,7 @@ import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Assumptions.assumeTrue
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.io.TempDir
-import org.pkl.commons.test.FileTestUtils
+import org.pkl.commons.test.Executables
 import org.pkl.commons.toPath
 import org.pkl.commons.writeString
 import org.pkl.core.EvaluatorBuilder
@@ -38,29 +37,14 @@ import org.pkl.core.externalreader.ExternalReaderProcess
 import org.pkl.core.externalreader.TestExternalModuleReader
 import org.pkl.core.externalreader.TestExternalReaderProcess
 import org.pkl.core.resource.ResourceReaders
-import org.pkl.core.util.IoUtils
 
 class ModuleKeyFactoriesTest {
   companion object {
-    private val externalReaderFixture by lazy {
-      val readerPath =
-        "pkl-core/build/fixtures/externalreader".let { if (IoUtils.isWindows()) "$it.bat" else it }
-
-      FileTestUtils.rootProjectDir.resolve(readerPath).also { path ->
-        if (!Files.exists(path)) {
-          throw AssertionError(
-            "Fixture `externalreader` not found. To fix this problem, first run" +
-              " `./gradlew pkl-core:externalReaderFixture`."
-          )
-        }
-      }
-    }
-
     @JvmStatic
     private fun pathEnvIsSet(): Boolean {
       return System.getenv("PATH")
         ?.split(File.pathSeparator)
-        ?.contains(externalReaderFixture.toAbsolutePath().toString()) ?: false
+        ?.contains(Executables.externalReaderFixture.toAbsolutePath().toString()) ?: false
     }
   }
 
@@ -184,7 +168,7 @@ class ModuleKeyFactoriesTest {
 
   @Test
   fun `external process -- spawning an executable using a path`() {
-    testExternalReader(externalReaderFixture.toAbsolutePath().toString())
+    testExternalReader(Executables.externalReaderFixture.toAbsolutePath().toString())
   }
 
   @Test

pkl-server/pkl-server.gradle.kts

@@ -19,6 +19,8 @@ plugins {
   id("pklNativeLifecycle")
 }
 
+val externalReaderFixtureConfiguration = configurations.create("externalReaderFixture")
+
 dependencies {
   implementation(projects.pklCore)
   implementation(libs.msgpack)
@@ -26,14 +28,34 @@ dependencies {
 
   testImplementation(projects.pklCommonsTest)
   testImplementation(libs.wiremock)
+
+  externalReaderFixtureConfiguration(project(":pkl-core", "externalReaderFixture"))
 }
 
-tasks.test { exclude("**/NativeServerTest.*") }
+tasks.test {
+  configureTest()
+  exclude("**/NativeServerTest.*")
+}
+
+private fun Test.configureTest() {
+  dependsOn(externalReaderFixtureConfiguration)
+
+  systemProperty(
+    "org.pkl.core.testExternalReaderPath",
+    externalReaderFixtureConfiguration.files.single().absolutePath,
+  )
+}
 
 private fun Test.configureNativeTest() {
+  configureTest()
   testClassesDirs = files(tasks.test.get().testClassesDirs)
   classpath = tasks.test.get().classpath
   include("**/NativeServerTest.*")
+
+  systemProperty(
+    "org.pkl.core.testExternalReaderPath",
+    externalReaderFixtureConfiguration.files.single().absolutePath,
+  )
 }
 
 val testMacExecutableAarch64 =

pkl-server/src/main/kotlin/org/pkl/server/Server.kt

@@ -89,9 +89,10 @@ class Server(private val transport: MessageTransport) : AutoCloseable {
    */
   override fun close() {
     transport.closeQuietly()
-    for ((_, evaluator) in evaluators) {
+    for ((evaluatorId, evaluator) in evaluators) {
       // if currently in use, blocks until cancellation complete
       evaluator.closeQuietly()
+      externalReaderProcesses[evaluatorId]?.values?.forEach { it.closeQuietly() }
     }
     executor.shutdown()
   }
@@ -145,7 +146,7 @@ class Server(private val transport: MessageTransport) : AutoCloseable {
     evaluator.close()
 
     // close any running ExternalProcess instances for the closed evaluator
-    externalReaderProcesses[message.evaluatorId]?.values?.forEach { it.close() }
+    externalReaderProcesses.remove(message.evaluatorId)?.values?.forEach { it.close() }
   }
 
   private fun buildDeclaredDependencies(