Use sessions with connection to eliminate async shutdown problems

(cherry picked from commit 78d334b)

compiler/daemon/daemon-client/src/org/jetbrains/kotlin/daemon/client/KotlinCompilerClient.kt

@@ -30,6 +30,7 @@ import java.io.OutputStream
 import java.io.PrintStream
 import java.net.SocketException
 import java.rmi.ConnectException
+import java.rmi.ConnectIOException
 import java.rmi.UnmarshalException
 import java.rmi.server.UnicastRemoteObject
 import java.util.concurrent.Semaphore
@@ -73,51 +74,48 @@ object KotlinCompilerClient {
                                 daemonOptions: DaemonOptions,
                                 reportingTargets: DaemonReportingTargets,
                                 autostart: Boolean = true
-    ): CompileService? {
+    ): CompileService? =
+            connectAndLease(compilerId,
+                            clientAliveFlagFile,
+                            daemonJVMOptions,
+                            daemonOptions,
+                            reportingTargets,
+                            autostart,
+                            leaseSession = false,
+                            sessionAliveFlagFile = null)?.first
 
-        var attempts = 0
 
-        fun retryOrRethrow(e: Exception) {
+    fun connectAndLease(compilerId: CompilerId,
-            if (attempts < DAEMON_CONNECT_CYCLE_ATTEMPTS) {
+                        clientAliveFlagFile: File,
-                reportingTargets.report(DaemonReportCategory.INFO, "retrying($attempts) on:" + e.toString())
+                        daemonJVMOptions: DaemonJVMOptions,
-            }
+                        daemonOptions: DaemonOptions,
-            else throw e
+                        reportingTargets: DaemonReportingTargets,
-        }
+                        autostart: Boolean,
-
+                        leaseSession: Boolean,
-        try {
+                        sessionAliveFlagFile: File? = null
-            while (attempts < DAEMON_CONNECT_CYCLE_ATTEMPTS) {
+    ): Pair<CompileService, Int>? = connectLoop(reportingTargets) {
-                try {
-                    attempts += 1
         val (service, newJVMOptions) = tryFindSuitableDaemonOrNewOpts(File(daemonOptions.runFilesPath), compilerId, daemonJVMOptions, { cat, msg -> reportingTargets.report(cat, msg) })
         if (service != null) {
             // the newJVMOptions could be checked here for additional parameters, if needed
             service.registerClient(clientAliveFlagFile.absolutePath)
             reportingTargets.report(DaemonReportCategory.DEBUG, "connected to the daemon")
-                        return service
+            if (!leaseSession) service to CompileService.NO_SESSION
+            else {
+                val sessionId = service.leaseCompileSession(sessionAliveFlagFile?.absolutePath)
+                if (sessionId is CompileService.CallResult.Dying)
+                    null
+                else
+                    service to sessionId.get()
             }
+        } else {
             reportingTargets.report(DaemonReportCategory.DEBUG, "no suitable daemon found")
             if (autostart) {
                 startDaemon(compilerId, newJVMOptions, daemonOptions, reportingTargets)
                 reportingTargets.report(DaemonReportCategory.DEBUG, "new daemon started, trying to find it")
             }
-                }
+            null
-                catch (e: SocketException) {
-                    retryOrRethrow(e)
-                }
-                catch (e: ConnectException) {
-                    retryOrRethrow(e)
-                }
-                catch (e: UnmarshalException) {
-                    retryOrRethrow(e)
         }
     }
-        }
-        catch (e: Throwable) {
-            reportingTargets.report(DaemonReportCategory.EXCEPTION, e.toString())
-        }
-        return null
-    }
-
 
     fun shutdownCompileService(compilerId: CompilerId, daemonOptions: DaemonOptions): Unit {
         connectToCompileService(compilerId, DaemonJVMOptions(), daemonOptions, DaemonReportingTargets(out = System.out), autostart = false, checkId = false)
@@ -130,6 +128,13 @@ object KotlinCompilerClient {
     }
 
 
+    fun leaseCompileSession(compilerService: CompileService, aliveFlagPath: String?): Int =
+            compilerService.leaseCompileSession(aliveFlagPath).get()
+
+    fun releaseCompileSession(compilerService: CompileService, sessionId: Int): Unit {
+        compilerService.releaseCompileSession(sessionId)
+    }
+
     fun compile(compilerService: CompileService,
                 sessionId: Int,
                 targetPlatform: CompileService.TargetPlatform,
@@ -293,6 +298,33 @@ object KotlinCompilerClient {
 
     // --- Implementation ---------------------------------------
 
+    private inline fun <R> connectLoop(reportingTargets: DaemonReportingTargets, body: () -> R?): R? {
+        try {
+            var attempts = 0
+            while (attempts < DAEMON_CONNECT_CYCLE_ATTEMPTS) {
+                attempts += 1
+                val (res, err) = try {
+                    body() to null
+                }
+                catch (e: SocketException) { null to e }
+                catch (e: ConnectException) { null to e }
+                catch (e: ConnectIOException) { null to e }
+                catch (e: UnmarshalException) { null to e }
+
+                when {
+                    res != null -> return res
+                    err != null && attempts >= DAEMON_CONNECT_CYCLE_ATTEMPTS -> throw err
+                    err != null ->
+                        reportingTargets.report(DaemonReportCategory.INFO, "retrying($attempts) on: " + err.toString())
+                }
+            }
+        }
+        catch (e: Throwable) {
+            reportingTargets.report(DaemonReportCategory.EXCEPTION, e.toString())
+        }
+        return null
+    }
+
     private fun DaemonReportingTargets.report(category: DaemonReportCategory, message: String, source: String = "daemon client") {
         if (category == DaemonReportCategory.DEBUG && !verboseReporting) return
         out?.println("[$source] ${category.name}: $message")

compiler/daemon/src/org/jetbrains/kotlin/daemon/CompileServiceImpl.kt

@@ -261,21 +261,30 @@ class CompileServiceImpl(
     }
 
     override fun scheduleShutdown(graceful: Boolean): CompileService.CallResult<Boolean> = ifAlive(minAliveness = Aliveness.Alive) {
-        CompileService.CallResult.Good(
+        when {
-                if (!graceful || state.alive.compareAndSet(Aliveness.Alive.ordinal, Aliveness.LastSession.ordinal)) {
+            !graceful -> {
+                shutdownWithDelay()
+                CompileService.CallResult.Good(true)
+            }
+            state.alive.compareAndSet(Aliveness.Alive.ordinal, Aliveness.LastSession.ordinal) -> {
                 timer.schedule(1) {
                     ifAliveExclusive(minAliveness = Aliveness.LastSession, ignoreCompilerChanged = true) {
                         when {
-                                !graceful -> shutdownImpl()
                             state.sessions.isEmpty() -> shutdownWithDelay()
-                                else -> log.info("Some sessions are active, waiting for them to finish")
+                            else -> {
+                                daemonOptions.autoshutdownIdleSeconds = TimeUnit.MILLISECONDS.toSeconds(daemonOptions.forceShutdownTimeoutMilliseconds).toInt()
+                                daemonOptions.autoshutdownUnusedSeconds = daemonOptions.autoshutdownIdleSeconds
+                                log.info("Graceful shutdown signalled; unused/idle timeouts are set to ${daemonOptions.autoshutdownUnusedSeconds}/${daemonOptions.autoshutdownIdleSeconds}s")
+                                log.info("Some sessions are active, waiting for them to finish")
+                            }
                         }
                         CompileService.CallResult.Ok()
                     }
                 }
-                    true
+                CompileService.CallResult.Good(true)
+            }
+            else -> CompileService.CallResult.Good(false)
         }
-                else false)
     }
 
     override fun remoteCompile(sessionId: Int,
@@ -562,7 +571,7 @@ class CompileServiceImpl(
         val stub = UnicastRemoteObject.exportObject(this, port, LoopbackNetworkInterface.clientLoopbackSocketFactory, LoopbackNetworkInterface.serverLoopbackSocketFactory) as CompileService
         registry.rebind (COMPILER_SERVICE_RMI_NAME, stub)
 
-        timer.schedule(100) {
+        timer.schedule(10) {
             initiateElections()
         }
         timer.schedule(delay = DAEMON_PERIODIC_CHECK_INTERVAL_MS, period = DAEMON_PERIODIC_CHECK_INTERVAL_MS) {

compiler/tests/org/jetbrains/kotlin/daemon/CompilerDaemonTest.kt

@@ -255,7 +255,7 @@ class CompilerDaemonTest : KotlinIntegrationTestBase() {
 
     fun testDaemonGracefulShutdown() {
         withFlagFile(getTestName(true), ".alive") { flagFile ->
-            val daemonOptions = DaemonOptions(autoshutdownIdleSeconds = 1, shutdownDelayMilliseconds = 1, runFilesPath = File(tmpdir, getTestName(true)).absolutePath)
+            val daemonOptions = DaemonOptions(autoshutdownIdleSeconds = 1, shutdownDelayMilliseconds = 1, forceShutdownTimeoutMilliseconds = 60000, runFilesPath = File(tmpdir, getTestName(true)).absolutePath)
             KotlinCompilerClient.shutdownCompileService(compilerId, daemonOptions)
 
             val logFile = createTempFile("kotlin-daemon-test", ".log")
@@ -401,7 +401,7 @@ class CompilerDaemonTest : KotlinIntegrationTestBase() {
 
     fun testParallelDaemonStart() {
 
-        val PARALLEL_THREADS_TO_START = 8
+        val PARALLEL_THREADS_TO_START = 10
 
         val doneLatch = CountDownLatch(PARALLEL_THREADS_TO_START)
 
@@ -412,18 +412,23 @@ class CompilerDaemonTest : KotlinIntegrationTestBase() {
         fun connectThread(threadNo: Int) = thread(name = "daemonConnect$threadNo") {
             try {
                 withFlagFile(getTestName(true), ".alive") { flagFile ->
-                    val daemonOptions = DaemonOptions(shutdownDelayMilliseconds = 10000, runFilesPath = File(tmpdir, getTestName(true)).absolutePath, verbose = true)
+                    withFlagFile(getTestName(true), ".salive") { sessionFlagFile ->
+                        val daemonOptions = DaemonOptions(shutdownDelayMilliseconds = 10000, // attempt to avoid immediate shutdown of the non-elected daemons
+                                                          runFilesPath = File(tmpdir, getTestName(true)).absolutePath,
+                                                          verbose = true)
                         val logFile = createTempFile("kotlin-daemon-test", ".log")
                         val daemonJVMOptions =
                                 configureDaemonJVMOptions(DaemonJVMOptions(maxMemory = "2048m"),
                                                           "D$COMPILE_DAEMON_LOG_PATH_PROPERTY=\"${logFile.loggerCompatiblePath}\"",
                                                           inheritMemoryLimits = false, inheritAdditionalProperties = false)
-                    val daemon = KotlinCompilerClient.connectToCompileService(compilerId, flagFile, daemonJVMOptions, daemonOptions, DaemonReportingTargets(out = System.err), autostart = true)
+                        val daemonWithSession = KotlinCompilerClient.connectAndLease(compilerId, flagFile, daemonJVMOptions, daemonOptions,
-                    assertNotNull("failed to connect daemon", daemon)
+                                                                                     DaemonReportingTargets(out = System.err), autostart = true,
+                                                                                     leaseSession = true, sessionAliveFlagFile = sessionFlagFile)
+                        assertNotNull("failed to connect daemon", daemonWithSession?.first)
                         val jar = tmpdir.absolutePath + File.separator + "hello.$threadNo.jar"
                         val res = KotlinCompilerClient.compile(
-                            daemon!!,
+                                daemonWithSession!!.first,
-                            CompileService.NO_SESSION,
+                                daemonWithSession.second,
                                 CompileService.TargetPlatform.JVM,
                                 arrayOf(File(getHelloAppBaseDir(), "hello.kt").absolutePath, "-d", jar),
                                 outStreams[threadNo])
@@ -431,6 +436,7 @@ class CompilerDaemonTest : KotlinIntegrationTestBase() {
                         logFiles[threadNo] = logFile
                     }
                 }
+            }
             finally {
                 doneLatch.countDown()
             }