Support multiple Truffle Engines incl. independent settings. (#369)

Co-authored-by: Alex Zerntev <alexzerntev@gmail.com>

client/src/main/scala/raw/client/api/CompilerService.scala

@@ -34,59 +34,48 @@ final class CompilerServiceException(
 
 object CompilerService {
 
-  private var engine: Engine = _
-  private var engineCount = 0
-  private val engineLock = new Object
-
-  private def getEngine(implicit settings: RawSettings): Engine = {
-    engineLock.synchronized {
-      if (engine == null) {
-        val options = new java.util.HashMap[String, String]()
-        if (settings.onTrainingWheels) {
-//          options.put("engine.CompileImmediately", "true")
-          //          options.put("engine.TraceCompilation", "true")
-          //          options.put("engine.BackgroundCompilation", "false")
-//          options.put("engine.CompilationFailureAction", "Throw")
-          //      options.put("engine.CompilationFailureAction", "Diagnose")
-          //      options.put("compiler.LogInlinedTargets", "true")
-          //  "-Dpolyglotimpl.CompilationFailureAction=Throw",
-          //  "-Dpolyglotimpl.TreatPerformanceWarningsAsErrors=false",
-          //  "-Dpolyglotimpl.CompilationExceptionsAreFatal=true",
-          //  "-Dpolyglotimpl.BackgroundCompilation=false",
-          //  "-Dpolyglotimpl.TraceCompilationDetails=true",
-          //  "-Dpolyglotimpl.TraceInlining=true"
-          //  "-Dgraal.Dump=Truffle:2",
-          //  "-Dgraal.DumpPath=/tmp/graal_dumps",
-          //  "-Dgraal.PrintGraph=Network",
-          // https://www.graalvm.org/latest/graalvm-as-a-platform/language-implementation-framework/Options/
-        }
-        engine = Engine.newBuilder().allowExperimentalOptions(true).options(options).build()
+  private val enginesLock = new Object
+  private val enginesCache = mutable.HashMap[RawSettings, Engine]()
+
+  // Return engine and a flag indicating if the engine was created.
+  def getEngine()(implicit settings: RawSettings): (Engine, Boolean) = {
+    enginesLock.synchronized {
+      enginesCache.get(settings) match {
+        case Some(engine) =>
+          // Re-using an engine someone else create before. This typically happens on staged compilation.
+          (engine, false)
+        case None =>
+          // First time creating an engine for this settings.
+          val options = new java.util.HashMap[String, String]()
+          if (settings.onTrainingWheels) {
+            // Refer to settings at:
+            // https://www.graalvm.org/latest/graalvm-as-a-platform/language-implementation-framework/Options/
+            //          options.put("engine.CompileImmediately", "true")
+            //          options.put("engine.TraceCompilation", "true")
+            //          options.put("engine.BackgroundCompilation", "false")
+            //          options.put("engine.CompilationFailureAction", "Throw")
+            //          options.put("engine.CompilationFailureAction", "Diagnose")
+            //          options.put("compiler.LogInlinedTargets", "true")
+          }
+          val engine = Engine.newBuilder().allowExperimentalOptions(true).options(options).build()
+          enginesCache.put(settings, engine)
+          (engine, true)
       }
-      engineCount += 1
-      engine
     }
   }
 
-  private def releaseEngine(): Unit = {
-    engineLock.synchronized {
-      engineCount -= 1
-      if (engineCount == 0) {
-        engine.close(true)
-        engine = null
-      }
+  def releaseEngine()(implicit settings: RawSettings): Unit = {
+    enginesLock.synchronized {
+      enginesCache.remove(settings).foreach(engine => engine.close(true))
     }
   }
 
 }
 
 trait CompilerService extends RawService {
 
-  import CompilerService._
-
   implicit protected def settings: RawSettings
 
-  protected lazy val engine: Engine = getEngine
-
   def language: Set[String]
 
   // Get the description of a source program.
@@ -290,10 +279,6 @@ trait CompilerService extends RawService {
     }
   }
 
-  override def doStop(): Unit = {
-    releaseEngine()
-  }
-
 }
 
 final case class Pos(line: Int, column: Int)

client/src/main/scala/raw/client/api/CompilerServiceProvider.scala

@@ -44,7 +44,9 @@ object CompilerServiceProvider {
     }
   }
 
-  def apply(language: String, classLoader: ClassLoader)(implicit settings: RawSettings): CompilerService = {
+  def apply(language: String, classLoader: ClassLoader)(
+      implicit settings: RawSettings
+  ): CompilerService = {
     instanceLock.synchronized {
       instanceMap.collectFirst { case ((l, Some(cl)), i) if cl == classLoader && l.contains(language) => i } match {
         case Some(instance) => instance
@@ -60,7 +62,12 @@ object CompilerServiceProvider {
   private[raw] def set(language: Set[String], instance: CompilerService): Unit = {
     instanceLock.synchronized {
       if (instance == null) {
-        instanceMap.remove((language, None))
+        // Stop and remove entries that match the `language`, regardless the class loader.
+        instanceMap.filterKeys(_._1 == language).foreach {
+          case (key, compiler) =>
+            compiler.stop()
+            instanceMap.remove(key)
+        }
       } else {
         instanceMap.put((language, None), instance)
       }

python-client/src/main/scala/raw/client/python/PythonCompilerService.scala

@@ -12,7 +12,7 @@
 
 package raw.client.python
 
-import org.graalvm.polyglot.{Context, PolyglotAccess, PolyglotException, Source, Value}
+import org.graalvm.polyglot.{Context, Engine, PolyglotAccess, PolyglotException, Source, Value}
 import raw.client.api.{
   AutoCompleteResponse,
   CompilerService,
@@ -53,8 +53,22 @@ import raw.utils.{RawSettings, RawUtils}
 
 import java.io.{IOException, OutputStream}
 
-class PythonCompilerService(maybeClassLoader: Option[ClassLoader] = None)(implicit protected val settings: RawSettings)
-    extends CompilerService {
+class PythonCompilerService(engineDefinition: (Engine, Boolean), maybeClassLoader: Option[ClassLoader] = None)(
+    implicit protected val settings: RawSettings
+) extends CompilerService {
+
+  private val (engine, initedEngine) = engineDefinition
+
+  // The default constructor allows an Engine to be specified, plus a flag to indicate whether it was created here
+  // or externally. That's necessary for the test framework.
+  // This is actually the "default constructor" which obtains a new engine or reuses an existing one.
+  // Note that the engine will be released when the service is stopped only IF this auxiliary constructor created it.
+  // Otherwise, we expect the external party - e.g. the test framework - to close it.
+  // Refer to Rql2TruffleCompilerServiceTestContext to see the engine being created and released from the test
+  // framework, so that every test suite instance has a fresh engine.
+  def this(maybeClassLoader: Option[ClassLoader] = None)(implicit settings: RawSettings) = {
+    this(CompilerService.getEngine, maybeClassLoader)
+  }
 
   override def language: Set[String] = Set("python")
 
@@ -294,6 +308,12 @@ class PythonCompilerService(maybeClassLoader: Option[ClassLoader] = None)(implic
     ValidateResponse(List.empty)
   }
 
+  override def doStop(): Unit = {
+    if (initedEngine) {
+      CompilerService.releaseEngine
+    }
+  }
+
   private def buildTruffleContext(
       environment: ProgramEnvironment,
       maybeOutputStream: Option[OutputStream] = None

snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2TruffleCompilerService.scala

@@ -42,11 +42,24 @@ object Rql2TruffleCompilerService {
   val language: Set[String] = Set("rql2", "rql2-truffle", "snapi")
 }
 
-class Rql2TruffleCompilerService(maybeClassLoader: Option[ClassLoader] = None)(
+class Rql2TruffleCompilerService(engineDefinition: (Engine, Boolean), maybeClassLoader: Option[ClassLoader])(
     implicit protected val settings: RawSettings
 ) extends Rql2CompilerService
     with Rql2TypeUtils {
 
+  private val (engine, initedEngine) = engineDefinition
+
+  // The default constructor allows an Engine to be specified, plus a flag to indicate whether it was created here
+  // or externally. That's necessary for the test framework.
+  // This is actually the "default constructor" which obtains a new engine or reuses an existing one.
+  // Note that the engine will be released when the service is stopped only IF this auxiliary constructor created it.
+  // Otherwise, we expect the external party - e.g. the test framework - to close it.
+  // Refer to Rql2TruffleCompilerServiceTestContext to see the engine being created and released from the test
+  // framework, so that every test suite instance has a fresh engine.
+  def this(maybeClassLoader: Option[ClassLoader] = None)(implicit settings: RawSettings) = {
+    this(CompilerService.getEngine, maybeClassLoader)
+  }
+
   override def language: Set[String] = Rql2TruffleCompilerService.language
 
   private val credentials = CredentialsServiceProvider(maybeClassLoader)
@@ -71,7 +84,7 @@ class Rql2TruffleCompilerService(maybeClassLoader: Option[ClassLoader] = None)(
 
   private def getProgramContext(user: AuthenticatedUser, environment: ProgramEnvironment): ProgramContext = {
     val compilerContext = getCompilerContext(user)
-    val runtimeContext = new RuntimeContext(compilerContext.sourceContext, settings, environment)
+    val runtimeContext = new RuntimeContext(compilerContext.sourceContext, environment)
     new Rql2ProgramContext(runtimeContext, compilerContext)
   }
 
@@ -658,9 +671,11 @@ class Rql2TruffleCompilerService(maybeClassLoader: Option[ClassLoader] = None)(
   }
 
   override def doStop(): Unit = {
-    super.doStop()
     compilerContextCaches.values.foreach(compilerContext => compilerContext.inferrer.stop())
     credentials.stop()
+    if (initedEngine) {
+      CompilerService.releaseEngine
+    }
   }
 
   private def rawValueToPolyglotValue(rawValue: RawValue, ctx: Context): Value = {

snapi-client/src/test/scala/raw/compiler/rql2/truffle/Rql2TruffleCompilerServiceTestContext.scala

@@ -12,6 +12,7 @@
 
 package raw.compiler.rql2.truffle
 
+import org.graalvm.polyglot.Engine
 import raw.client.rql2.truffle.Rql2TruffleCompilerService
 import raw.compiler.rql2.api.Rql2CompilerServiceTestContext
 import raw.utils.{RawTestSuite, RawUtils, SettingsTestContext}
@@ -21,11 +22,23 @@ trait Rql2TruffleCompilerServiceTestContext extends Rql2CompilerServiceTestConte
 
   var rql2TruffleCompilerService: Rql2TruffleCompilerService = _
 
+  var engine: Engine = _
+
   override def beforeAll(): Unit = {
     super.beforeAll()
+
+    // Create an isolated Truffle Engine
+    val options = new java.util.HashMap[String, String]()
+    options.put("rql.settings", settings.renderAsString)
+    engine = Engine
+      .newBuilder()
+      .allowExperimentalOptions(true)
+      .options(options)
+      .build()
+
     property("raw.compiler.impl", "rql2-truffle")
 
-    rql2TruffleCompilerService = new Rql2TruffleCompilerService
+    rql2TruffleCompilerService = new Rql2TruffleCompilerService((engine, false), None)
     setCompilerService(rql2TruffleCompilerService)
   }
 
@@ -34,6 +47,10 @@ trait Rql2TruffleCompilerServiceTestContext extends Rql2CompilerServiceTestConte
       RawUtils.withSuppressNonFatalException(rql2TruffleCompilerService.stop())
       rql2TruffleCompilerService = null
     }
+    if (engine != null) {
+      RawUtils.withSuppressNonFatalException(engine.close())
+      engine = null
+    }
     super.afterAll()
   }
 

snapi-client/src/test/scala/raw/compiler/rql2/truffle/TruffleCompilerTestContext.scala

@@ -14,23 +14,8 @@ package raw.compiler.rql2.truffle
 
 import raw.compiler.rql2.api.Rql2OutputTestContext
 import raw.compiler.rql2.tests.CompilerTestContext
-import raw.runtime.truffle.RawLanguage
 
 class TruffleCompilerTestContext
     extends CompilerTestContext
     with Rql2OutputTestContext
-    with Rql2TruffleCompilerServiceTestContext {
-
-  override def beforeAll(): Unit = {
-    super.beforeAll()
-    // Forcibly drop all language caches before each test suite.
-    RawLanguage.dropCaches()
-  }
-
-  override def afterAll(): Unit = {
-    // Forcibly drop all language caches after each test suite.
-    RawLanguage.dropCaches()
-    super.afterAll()
-  }
-
-}
+    with Rql2TruffleCompilerServiceTestContext

snapi-frontend/src/main/scala/raw/runtime/RuntimeContext.scala

@@ -24,12 +24,13 @@ import scala.collection.mutable
  */
 class RuntimeContext(
     val sourceContext: SourceContext,
-    val settings: RawSettings,
     val environment: ProgramEnvironment
 ) {
 
   private val zoneID = ZoneId.of("UTC")
 
+  final val settings: RawSettings = sourceContext.settings
+
   final val currentTimestamp: LocalDateTime = LocalDateTime.now(zoneID)
 
   final val paramsFromTemplating: mutable.Map[String, RawValue] = mutable.HashMap[String, RawValue]()
@@ -42,12 +43,10 @@ class RuntimeContext(
 
   final def cloneWith(
       newSourceContext: SourceContext = sourceContext,
-      newSettings: RawSettings = settings,
       newEnvironment: ProgramEnvironment = environment
   ): RuntimeContext = {
     new RuntimeContext(
       newSourceContext,
-      newSettings,
       newEnvironment
     )
   }

snapi-frontend/src/test/scala/raw/creds/api/CredentialsTestContext.scala

@@ -62,12 +62,6 @@ trait CredentialsTestContext extends BeforeAndAfterAll {
     newHttpCreds.foreach { case (user, (name, cred)) => credentials.registerNewHttpCredential(user, name, cred) }
     dropboxTokens.foreach { case (user, token) => credentials.registerDropboxToken(user, token) }
     secrets.foreach { case (user, secret) => credentials.registerSecret(user, secret) }
-    // Set the system properties for the credentials so that this is overriding what RawLanguage
-    // runs with. This way, the credential server booted by the test framework is used.
-    for (property <- Seq("raw.creds.impl", "raw.creds.client.server-address")) {
-      System.clearProperty(property)
-      settings.getStringOpt(property).foreach(v => System.setProperty(property, v))
-    }
   }
 
   override def afterAll(): Unit = {