Introduce commHandlers for kernel-provided functionality; support deb…

…ugPortComm messages in kernel
Kotlin · Aug 30, 2022 · db0dc9c · db0dc9c
1 parent a4df334
commit db0dc9c
Show file tree

Hide file tree

Showing 9 changed files with 88 additions and 3 deletions.
diff --git a/jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/debug.kt b/jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/debug.kt
@@ -0,0 +1,12 @@
+package org.jetbrains.kotlinx.jupyter.messaging
+
+import kotlinx.serialization.Serializable
+
+object ProvidedCommMessages {
+    const val OPEN_DEBUG_PORT_TARGET: String = "open_debug_port_target"
+}
+
+@Serializable
+class OpenDebugPortReply(
+    val port: Int?
+) : OkReply()
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/connection.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/connection.kt
@@ -149,6 +149,9 @@ class JupyterConnectionImpl(
 
     override val stdinIn = StdinInputStream()
 
+    override val debugPort: Int?
+        get() = config.debugPort
+
     private var _contextMessage: RawMessage? = null
     override fun setContextMessage(message: RawMessage?) {
         _contextMessage = message

diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/JupyterConnectionInternal.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/JupyterConnectionInternal.kt
@@ -24,6 +24,8 @@ interface JupyterConnectionInternal : JupyterConnection {
     val executor: JupyterExecutor
     val stdinIn: InputStream
 
+    val debugPort: Int?
+
     fun sendStatus(status: KernelStatus, incomingMessage: RawMessage? = null)
     fun doWrappedInBusyIdle(incomingMessage: RawMessage? = null, action: () -> Unit)
     fun updateSessionInfo(message: RawMessage)

diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/defaultCommHandlers.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/defaultCommHandlers.kt
@@ -0,0 +1,36 @@
+package org.jetbrains.kotlinx.jupyter.messaging
+
+import kotlinx.serialization.json.JsonObject
+import org.jetbrains.kotlinx.jupyter.ReplForJupyterImpl
+import org.jetbrains.kotlinx.jupyter.api.libraries.Comm
+import org.jetbrains.kotlinx.jupyter.api.libraries.sendData
+import org.jetbrains.kotlinx.jupyter.messaging.ProvidedCommMessages.OPEN_DEBUG_PORT_TARGET
+
+typealias CommMessageProcessCallback = (Comm, JsonObject, ReplForJupyterImpl) -> Unit
+
+abstract class DefaultCommHandler {
+    abstract val targetId: String
+
+    abstract fun onReceive(): CommMessageProcessCallback
+
+    abstract fun register(repl: ReplForJupyterImpl)
+}
+
+class DebugPortCommHandler : DefaultCommHandler() {
+    override val targetId: String
+        get() = OPEN_DEBUG_PORT_TARGET
+
+    override fun onReceive(): CommMessageProcessCallback {
+        return { comm, _, repl ->
+            comm.sendData(OpenDebugPortReply(repl.debugPort))
+        }
+    }
+
+    override fun register(repl: ReplForJupyterImpl) {
+        repl.notebook.commManager.registerCommTarget(targetId) { comm, _ ->
+            comm.onMessage {
+                onReceive()
+            }
+        }
+    }
+}
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt
@@ -53,6 +53,7 @@ import org.jetbrains.kotlinx.jupyter.magics.CompoundCodePreprocessor
 import org.jetbrains.kotlinx.jupyter.magics.ErrorsMagicsProcessor
 import org.jetbrains.kotlinx.jupyter.magics.FullMagicsHandler
 import org.jetbrains.kotlinx.jupyter.magics.MagicsProcessor
+import org.jetbrains.kotlinx.jupyter.messaging.DefaultCommHandler
 import org.jetbrains.kotlinx.jupyter.messaging.DisplayHandler
 import org.jetbrains.kotlinx.jupyter.messaging.NoOpDisplayHandler
 import org.jetbrains.kotlinx.jupyter.repl.CellExecutor
@@ -124,6 +125,7 @@ interface ReplOptions {
     var executedCodeLogging: ExecutedCodeLogging
     var writeCompiledClasses: Boolean
     var outputConfig: OutputConfig
+    val debugPort: Int?
 }
 
 interface ReplForJupyter {
@@ -188,9 +190,17 @@ class ReplForJupyterImpl(
     private val scriptReceivers: List<Any> = emptyList(),
     override val isEmbedded: Boolean = false,
     override val notebook: MutableNotebook,
-    override val librariesScanner: LibrariesScanner
+    override val librariesScanner: LibrariesScanner,
+    override val debugPort: Int? = null,
+    providedCommHandlers: List<DefaultCommHandler>
 ) : ReplForJupyter, ReplOptions, BaseKernelHost, UserHandlesProvider {
 
+    init {
+        providedCommHandlers.forEach {
+            it.register(this)
+        }
+    }
+
     override val currentBranch: String
         get() = runtimeProperties.currentBranch
     override val librariesDir: File = KERNEL_LIBRARIES.homeLibrariesDir(homeDir)
@@ -418,7 +428,7 @@ class ReplForJupyterImpl(
     @Suppress("unused")
     private fun printUsagesInfo(cellId: Int, usedVariables: Set<String>?) {
         log.debug(buildString {
-            if (usedVariables == null || usedVariables.isEmpty()) {
+            if (usedVariables.isNullOrEmpty()) {
                 append("No usages for cell $cellId")
                 return@buildString
             }

diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/BaseReplFactory.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/BaseReplFactory.kt
@@ -10,6 +10,8 @@ import org.jetbrains.kotlinx.jupyter.libraries.LibrariesScanner
 import org.jetbrains.kotlinx.jupyter.libraries.LibraryResolver
 import org.jetbrains.kotlinx.jupyter.libraries.ResolutionInfoProvider
 import org.jetbrains.kotlinx.jupyter.messaging.CommManagerImpl
+import org.jetbrains.kotlinx.jupyter.messaging.DebugPortCommHandler
+import org.jetbrains.kotlinx.jupyter.messaging.DefaultCommHandler
 import org.jetbrains.kotlinx.jupyter.messaging.DisplayHandler
 import org.jetbrains.kotlinx.jupyter.messaging.NoOpDisplayHandler
 import java.io.File
@@ -28,4 +30,15 @@ abstract class BaseReplFactory : ReplFactory() {
     override fun provideIsEmbedded() = false
     override fun provideLibrariesScanner(): LibrariesScanner = LibrariesScanner(notebook)
     override fun provideCommManager(): CommManager = CommManagerImpl(connection)
+    override fun provideDefaultCommHandlers(): List<DefaultCommHandler> = listOf(
+        DebugPortCommHandler()
+    )
+
+    init {
+        val uniqueTargets = defaultCommHandlers.map { it.targetId }.toSet().size
+        assert(uniqueTargets == defaultCommHandlers.size) {
+            val duplicates = defaultCommHandlers.groupingBy { it }.eachCount().filter { it.value > 1 }.keys
+            "Duplicate default comm targets found! $duplicates"
+        }
+    }
 }
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/MockJupyterConnection.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/MockJupyterConnection.kt
@@ -35,6 +35,8 @@ object MockJupyterConnection : JupyterConnectionInternal {
         get() = throw NotImplementedError()
     override val stdinIn: InputStream
         get() = throw NotImplementedError()
+    override val debugPort: Int?
+        get() = null
 
     override fun sendStatus(status: KernelStatus, incomingMessage: RawMessage?) {
         throw NotImplementedError()

diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/ReplFactory.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/ReplFactory.kt
@@ -8,6 +8,7 @@ import org.jetbrains.kotlinx.jupyter.api.libraries.CommManager
 import org.jetbrains.kotlinx.jupyter.libraries.LibrariesScanner
 import org.jetbrains.kotlinx.jupyter.libraries.LibraryResolver
 import org.jetbrains.kotlinx.jupyter.libraries.ResolutionInfoProvider
+import org.jetbrains.kotlinx.jupyter.messaging.DefaultCommHandler
 import org.jetbrains.kotlinx.jupyter.messaging.DisplayHandler
 import org.jetbrains.kotlinx.jupyter.messaging.JupyterConnectionInternal
 import java.io.File
@@ -26,7 +27,9 @@ abstract class ReplFactory {
             scriptReceivers,
             isEmbedded,
             notebook,
-            librariesScanner
+            librariesScanner,
+            connection.debugPort,
+            defaultCommHandlers
         )
     }
 
@@ -69,6 +72,9 @@ abstract class ReplFactory {
     protected val commManager: CommManager by lazy { provideCommManager() }
     protected abstract fun provideCommManager(): CommManager
 
+    protected val defaultCommHandlers: List<DefaultCommHandler> by lazy { provideDefaultCommHandlers() }
+    protected abstract fun provideDefaultCommHandlers(): List<DefaultCommHandler>
+
     // TODO: add other methods incl. display handler and socket messages listener
     // Inheritors should be constructed of connection (JupyterConnection)
 }
diff --git a/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/parseMagicsTests.kt b/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/parseMagicsTests.kt
@@ -142,6 +142,7 @@ class ParseMagicsTests {
         override var executedCodeLogging = ExecutedCodeLogging.OFF
         override var writeCompiledClasses = false
         override var outputConfig = OutputConfig()
+        override val debugPort: Int? = null
     }
 
     private val options = TestReplOptions()