Add suspended cache validation; Add possibility to remove old variables from cache; some improvements

Author: nikolay-egorov

src/main/kotlin/org/jetbrains/kotlinx/jupyter/message_types.kt

@@ -560,7 +560,8 @@ class ListErrorsReply(
 class SerializationRequest(
     val cellId: Int,
     val descriptorsState: Map<String, SerializedVariablesState>,
-    val topLevelDescriptorName: String = ""
+    val topLevelDescriptorName: String = "",
+    val pathToDescriptor: List<String> = emptyList()
 ) : MessageContent()
 
 @Serializable

src/main/kotlin/org/jetbrains/kotlinx/jupyter/protocol.kt

@@ -317,7 +317,11 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
 
             val messageContent = getVariablesDescriptorsFromJson(data)
             GlobalScope.launch(Dispatchers.Default) {
-                repl.serializeVariables(messageContent.topLevelDescriptorName, messageContent.descriptorsState) { result ->
+                repl.serializeVariables(
+                    messageContent.topLevelDescriptorName,
+                    messageContent.descriptorsState,
+                    messageContent.pathToDescriptor
+                ) { result ->
                     sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_OPEN, content = result))
                 }
             }
@@ -339,7 +343,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
         is SerializationRequest -> {
             GlobalScope.launch(Dispatchers.Default) {
                 if (content.topLevelDescriptorName.isNotEmpty()) {
-                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState) { result ->
+                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState, content.pathToDescriptor) { result ->
                         sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
                     }
                 } else {

src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt

@@ -5,6 +5,9 @@ import jupyter.kotlin.DependsOn
 import jupyter.kotlin.KotlinContext
 import jupyter.kotlin.KotlinKernelHostProvider
 import jupyter.kotlin.Repository
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.launch
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import org.jetbrains.kotlinx.jupyter.api.Code
@@ -139,7 +142,8 @@ interface ReplForJupyter {
 
     suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
 
-    suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
+    suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String> = emptyList(),
+                                   callback: (SerializationReply) -> Unit)
 
     val homeDir: File?
 
@@ -212,7 +216,7 @@ class ReplForJupyterImpl(
 
     override val variablesSerializer = VariablesSerializer()
 
-    private val librariesScanner = LibrariesScanner(notebook)
+    val librariesScanner = LibrariesScanner(notebook)
     private val resourcesProcessor = LibraryResourcesProcessorImpl()
 
     override var outputConfig
@@ -459,8 +463,10 @@ class ReplForJupyterImpl(
             // printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
             val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState, notebook.unchangedVariables())
 
+            GlobalScope.launch(Dispatchers.Default) {
+                variablesSerializer.tryValidateCache(jupyterId - 1, notebook.cellVariables)
+            }
 
-            val variablesStateUpdate = notebook.variablesState.mapValues { "" }
             EvalResultEx(
                 result.result.value,
                 rendered,
@@ -567,8 +573,9 @@ class ReplForJupyterImpl(
         doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
     }
 
-    override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
-        doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(), ::doSerializeVariables)
+    override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String>,
+                                            callback: (SerializationReply) -> Unit) {
+        doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback, pathToDescriptor = pathToDescriptor), serializationQueue, SerializationReply(), ::doSerializeVariables)
     }
 
     private fun doSerializeVariables(args: SerializationArgs): SerializationReply {
@@ -579,7 +586,7 @@ class ReplForJupyterImpl(
             finalAns
         }
         args.descriptorsState.forEach { (name, state) ->
-            resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state)
+            resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state, args.pathToDescriptor)
         }
         log.debug("Serialization cellID: $cellId")
         log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")
@@ -623,6 +630,7 @@ class ReplForJupyterImpl(
         val descriptorsState: Map<String, SerializedVariablesState>,
         var cellId: Int = -1,
         val topLevelVarName: String = "",
+        val pathToDescriptor: List<String> = emptyList(),
         override val callback: (SerializationReply) -> Unit
     ) : LockQueueArgs<SerializationReply>
 

src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/InternalEvaluatorImpl.kt

@@ -155,7 +155,6 @@ internal class InternalEvaluatorImpl(
 
     private fun updateVariablesState(cellId: Int) {
         variablesWatcher.removeOldUsages(cellId)
-
         variablesHolder.forEach {
             val state = it.value as VariableStateImpl
 
@@ -189,6 +188,8 @@ internal class InternalEvaluatorImpl(
 
                 put(property.name, state)
             }
+            // remove old
+            variablesWatcher.removeOldDeclarations(cellId, addedDeclarations)
         }
     }
 
@@ -199,7 +200,7 @@ internal class InternalEvaluatorImpl(
     private fun updateDataAfterExecution(lastExecutionCellId: Int, resultValue: ResultValue) {
         variablesWatcher.ensureStorageCreation(lastExecutionCellId)
         variablesHolder += getVisibleVariables(resultValue, lastExecutionCellId)
-
+        // remove unreached variables
         updateVariablesState(lastExecutionCellId)
     }
 }