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

@@ -578,7 +578,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

@@ -315,7 +315,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))
                 }
             }
@@ -337,7 +341,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
@@ -122,7 +125,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?
 
@@ -193,7 +197,7 @@ class ReplForJupyterImpl(
 
     override val variablesSerializer = VariablesSerializer()
 
-    private val librariesScanner = LibrariesScanner(notebook)
+    val librariesScanner = LibrariesScanner(notebook)
     private val resourcesProcessor = LibraryResourcesProcessorImpl()
 
     override var outputConfig
@@ -428,8 +432,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,
@@ -532,8 +538,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 {
@@ -544,7 +551,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}")
@@ -588,6 +595,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

@@ -151,7 +151,6 @@ internal class InternalEvaluatorImpl(
 
     private fun updateVariablesState(cellId: Int) {
         variablesWatcher.removeOldUsages(cellId)
-
         variablesHolder.forEach {
             val state = it.value as VariableStateImpl
 
@@ -171,11 +170,23 @@ internal class InternalEvaluatorImpl(
             it.name
         }.toHashSet()
         val ans = mutableMapOf<String, VariableStateImpl>()
+        // maybe remove known declarations
+        val addedDeclarations = mutableSetOf<String>()
+
         fields.forEach { property ->
             if (!memberKPropertiesNames.contains(property.name)) return@forEach
 
             val state = VariableStateImpl(property, cellClassInstance)
             variablesWatcher.addDeclaration(cellId, property.name)
+            addedDeclarations.add(property.name)
+
+            // try check values
+            if (variablesHolder.containsKey(property.name)) {
+                val seenState = variablesHolder[property.name]
+                if (seenState?.value?.equals(state.value) == true) {
+                    addedDeclarations.remove(property.name)
+                }
+            }
 
             // it was val, now it's var
             if (isValField(property)) {
@@ -187,6 +198,9 @@ internal class InternalEvaluatorImpl(
 
             ans[property.name] = state
         }
+        // remove old
+        variablesWatcher.removeOldDeclarations(cellId, addedDeclarations)
+
         return ans
     }
 
@@ -197,7 +211,7 @@ internal class InternalEvaluatorImpl(
     private fun updateDataAfterExecution(lastExecutionCellId: Int, resultValue: ResultValue) {
         variablesWatcher.ensureStorageCreation(lastExecutionCellId)
         variablesHolder += getVisibleVariables(resultValue, lastExecutionCellId)
-
+        // remove unreached variables
         updateVariablesState(lastExecutionCellId)
     }
 }