Make DescriptorSchemaCache in Json thread-local on Native

formats/json/commonMain/src/kotlinx/serialization/json/Json.kt

@@ -56,7 +56,12 @@ public sealed class Json(
     override val serializersModule: SerializersModule
 ) : StringFormat {
 
-    internal val schemaCache: DescriptorSchemaCache = DescriptorSchemaCache()
+    @Deprecated(
+        "Should not be accessed directly, use Json.schemaCache accessor instead",
+        ReplaceWith("schemaCache"),
+        DeprecationLevel.ERROR
+    )
+    internal val _schemaCache: DescriptorSchemaCache = DescriptorSchemaCache()
 
     /**
      * The default instance of [Json] with default configuration.
@@ -291,5 +296,10 @@ private class JsonImpl(configuration: JsonConfiguration, module: SerializersModu
     }
 }
 
+/**
+ * This accessor should be used to workaround for freezing problems in Native, see Native source set
+ */
+internal expect val Json.schemaCache: DescriptorSchemaCache
+
 private const val defaultIndent = "    "
 private const val defaultDiscriminator = "type"

formats/json/jsMain/src/kotlinx/serialization/json/JsonSchemaCache.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2017-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.serialization.json
+
+import kotlinx.serialization.json.internal.*
+
+@Suppress("DEPRECATION_ERROR")
+internal actual val Json.schemaCache: DescriptorSchemaCache get() = this._schemaCache

formats/json/jvmMain/src/kotlinx/serialization/json/JsonSchemaCache.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2017-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.serialization.json
+
+import kotlinx.serialization.json.internal.*
+
+@Suppress("DEPRECATION_ERROR")
+internal actual val Json.schemaCache: DescriptorSchemaCache get() = this._schemaCache

formats/json/nativeMain/src/kotlinx/serialization/json/JsonSchemaCache.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2017-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.serialization.json
+
+import kotlinx.serialization.json.internal.*
+import kotlin.native.ref.*
+import kotlin.random.*
+
+/**
+ * This maps emulate thread-locality of DescriptorSchemaCache for Native.
+ *
+ * Custom JSON instances are considered thread-safe (in JVM) and can be frozen and transferred to different workers (in Native).
+ * Therefore, DescriptorSchemaCache should be either a concurrent freeze-aware map or thread local.
+ * Each JSON instance have it's own schemaCache, and it's impossible to use @ThreadLocal on non-global vals.
+ * Thus we make @ThreadLocal this special map: it provides schemaCache for a particular Json instance
+ * and should be used instead of a member `_schemaCache` on Native.
+ *
+ * To avoid memory leaks (when Json instance is no longer in use), WeakReference is used with periodical self-cleaning.
+ */
+@ThreadLocal
+private val jsonToCache: MutableMap<WeakJson, DescriptorSchemaCache> = mutableMapOf()
+
+/**
+ * Because WeakReference itself does not have proper equals/hashCode
+ */
+private class WeakJson(json: Json) {
+    private val ref = WeakReference(json)
+    private val initialHashCode = json.hashCode()
+
+    val isDead: Boolean get() = ref.get() == null
+
+    override fun equals(other: Any?): Boolean {
+        if (other !is WeakJson) return false
+        val thiz = this.ref.get() ?: return false
+        val that = other.ref.get() ?: return false
+        return thiz == that
+    }
+
+    override fun hashCode(): Int = initialHashCode
+}
+
+/**
+ * To maintain O(1) access, we cleanup the table from dead references with 1/size probability
+ */
+private fun cleanUpWeakMap() {
+    val size = jsonToCache.size
+    if (size <= 10) return // 10 is arbitrary small number to ignore polluting
+    // Roll 1/size probability
+    if (Random.nextInt(0, size) == 0) {
+        val iter = jsonToCache.iterator()
+        while (iter.hasNext()) {
+            if (iter.next().key.isDead) iter.remove()
+        }
+    }
+}
+
+/**
+ * Accessor for DescriptorSchemaCache
+ */
+internal actual val Json.schemaCache: DescriptorSchemaCache
+    get() = jsonToCache.getOrPut(WeakJson(this)) { DescriptorSchemaCache() }.also { cleanUpWeakMap() }

formats/json/nativeTest/src/kotlinx/serialization/json/MultiWorkerJsonTest.kt

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2017-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.serialization.json
+
+import kotlinx.serialization.*
+import kotlin.native.concurrent.*
+import kotlin.test.*
+
+class MultiWorkerJsonTest {
+    @Serializable
+    data class PlainOne(val one: Int)
+
+    @Serializable
+    data class PlainTwo(val two: Int)
+
+    private fun doTest(json: () -> Json) {
+        val worker = Worker.start()
+        val operation = {
+            for (i in 0..999) {
+                assertEquals(PlainOne(42), json().decodeFromString("""{"one":42,"two":239}"""))
+            }
+        }
+        worker.executeAfter(1000, operation.freeze())
+        for (i in 0..999) {
+            assertEquals(PlainTwo(239), json().decodeFromString("""{"one":42,"two":239}"""))
+        }
+        worker.requestTermination()
+    }
+
+
+    @Test
+    fun testJsonIsFreezeSafe() {
+        val json = Json {
+            isLenient = true
+            ignoreUnknownKeys = true
+            useAlternativeNames = true
+        }
+        // reuse instance
+        doTest { json }
+    }
+
+    @Test
+    fun testJsonInstantiation() {
+        // create new instanceEveryTime
+        doTest {
+            Json {
+                isLenient = true
+                ignoreUnknownKeys = true
+                useAlternativeNames = true
+            }
+        }
+    }
+}