Pool Kryo instances for efficiency. (#352)

Pooled Kryo

node/src/main/kotlin/net/corda/node/serialization/DefaultWhitelist.kt

@@ -11,6 +11,9 @@ import java.time.LocalDate
 import java.time.Period
 import java.util.*
 
+/**
+ * NOTE: We do not whitelist [HashMap] or [HashSet] since they are unstable under serialization.
+ */
 class DefaultWhitelist : CordaPluginRegistry() {
     override fun customizeSerialization(custom: SerializationCustomization): Boolean {
         custom.apply {
@@ -41,7 +44,6 @@ class DefaultWhitelist : CordaPluginRegistry() {
             addToWhitelist(java.time.Instant::class.java)
             addToWhitelist(java.time.LocalDate::class.java)
             addToWhitelist(java.util.Collections.singletonMap("A", "B").javaClass)
-            addToWhitelist(java.util.HashMap::class.java)
             addToWhitelist(java.util.LinkedHashMap::class.java)
             addToWhitelist(BigDecimal::class.java)
             addToWhitelist(LocalDate::class.java)

node/src/main/kotlin/net/corda/node/services/messaging/CordaRPCClientImpl.kt

@@ -113,18 +113,20 @@ class CordaRPCClientImpl(private val session: ClientSession,
     @GuardedBy("sessionLock")
     private val addressToQueuedObservables = CacheBuilder.newBuilder().weakValues().build<String, QueuedObservable>()
     // This is used to hold a reference counted hard reference when we know there are subscribers.
-    private val hardReferencesToQueuedObservables = mutableSetOf<QueuedObservable>()
+    private val hardReferencesToQueuedObservables = Collections.synchronizedSet(mutableSetOf<QueuedObservable>())
 
     private var producer: ClientProducer? = null
 
-    private inner class ObservableDeserializer(private val qName: String,
-                                               private val rpcName: String,
-                                               private val rpcLocation: Throwable) : Serializer<Observable<Any>>() {
+    class ObservableDeserializer() : Serializer<Observable<Any>>() {
         override fun read(kryo: Kryo, input: Input, type: Class<Observable<Any>>): Observable<Any> {
+            val qName = kryo.context[RPCKryoQNameKey] as String
+            val rpcName = kryo.context[RPCKryoMethodNameKey] as String
+            val rpcLocation = kryo.context[RPCKryoLocationKey] as Throwable
+            val rpcClient = kryo.context[RPCKryoClientKey] as CordaRPCClientImpl
             val handle = input.readInt(true)
-            val ob = sessionLock.withLock {
-                addressToQueuedObservables.getIfPresent(qName) ?: QueuedObservable(qName, rpcName, rpcLocation, this).apply {
-                    addressToQueuedObservables.put(qName, this)
+            val ob = rpcClient.sessionLock.withLock {
+                rpcClient.addressToQueuedObservables.getIfPresent(qName) ?: rpcClient.QueuedObservable(qName, rpcName, rpcLocation).apply {
+                    rpcClient.addressToQueuedObservables.put(qName, this)
                 }
             }
             val result = ob.getForHandle(handle)
@@ -182,9 +184,17 @@ class CordaRPCClientImpl(private val session: ClientSession,
 
             checkMethodVersion(method)
 
-            // sendRequest may return a reconfigured Kryo if the method returns observables.
-            val kryo: Kryo = sendRequest(args, location, method) ?: createRPCKryo()
-            val next: ErrorOr<*> = receiveResponse(kryo, method, timeout)
+            val msg: ClientMessage = createMessage(method)
+            // We could of course also check the return type of the method to see if it's Observable, but I'd
+            // rather haved the annotation be used consistently.
+            val returnsObservables = method.isAnnotationPresent(RPCReturnsObservables::class.java)
+            val kryo = if (returnsObservables) maybePrepareForObservables(location, method, msg) else createRPCKryoForDeserialization(this@CordaRPCClientImpl)
+            val next: ErrorOr<*> = try {
+                sendRequest(args, msg)
+                receiveResponse(kryo, method, timeout)
+            } finally {
+                releaseRPCKryoForDeserialization(kryo)
+            }
             rpcLog.debug { "<- RPC <- ${method.name} = $next" }
             return unwrapOrThrow(next)
         }
@@ -215,22 +225,18 @@ class CordaRPCClientImpl(private val session: ClientSession,
             return next
         }
 
-        private fun sendRequest(args: Array<out Any>?, location: Throwable, method: Method): Kryo? {
-            // We could of course also check the return type of the method to see if it's Observable, but I'd
-            // rather haved the annotation be used consistently.
-            val returnsObservables = method.isAnnotationPresent(RPCReturnsObservables::class.java)
-
+        private fun sendRequest(args: Array<out Any>?, msg: ClientMessage) {
             sessionLock.withLock {
-                val msg: ClientMessage = createMessage(method)
-                val kryo = if (returnsObservables) maybePrepareForObservables(location, method, msg) else null
+                val argsKryo = createRPCKryoForDeserialization(this@CordaRPCClientImpl)
                 val serializedArgs = try {
-                    (args ?: emptyArray<Any?>()).serialize(createRPCKryo())
+                    (args ?: emptyArray<Any?>()).serialize(argsKryo)
                 } catch (e: KryoException) {
                     throw RPCException("Could not serialize RPC arguments", e)
+                } finally {
+                    releaseRPCKryoForDeserialization(argsKryo)
                 }
                 msg.writeBodyBufferBytes(serializedArgs.bytes)
                 producer!!.send(ArtemisMessagingComponent.RPC_REQUESTS_QUEUE, msg)
-                return kryo
             }
         }
 
@@ -242,7 +248,7 @@ class CordaRPCClientImpl(private val session: ClientSession,
             msg.putLongProperty(ClientRPCRequestMessage.OBSERVATIONS_TO, observationsId)
             // And make sure that we deserialise observable handles so that they're linked to the right
             // queue. Also record a bit of metadata for debugging purposes.
-            return createRPCKryo(observableSerializer = ObservableDeserializer(observationsQueueName, method.name, location))
+            return createRPCKryoForDeserialization(this@CordaRPCClientImpl, observationsQueueName, method.name, location)
         }
 
         private fun createMessage(method: Method): ClientMessage {
@@ -278,8 +284,7 @@ class CordaRPCClientImpl(private val session: ClientSession,
     @ThreadSafe
     private inner class QueuedObservable(private val qName: String,
                                          private val rpcName: String,
-                                         private val rpcLocation: Throwable,
-                                         private val observableDeserializer: ObservableDeserializer) {
+                                         private val rpcLocation: Throwable) {
         private val root = PublishSubject.create<MarshalledObservation>()
         private val rootShared = root.doOnUnsubscribe { close() }.share()
 
@@ -345,8 +350,10 @@ class CordaRPCClientImpl(private val session: ClientSession,
 
         private fun deliver(msg: ClientMessage) {
             msg.acknowledge()
-            val kryo = createRPCKryo(observableSerializer = observableDeserializer)
-            val received: MarshalledObservation = msg.deserialize(kryo)
+            val kryo = createRPCKryoForDeserialization(this@CordaRPCClientImpl, qName, rpcName, rpcLocation)
+            val received: MarshalledObservation = try { msg.deserialize(kryo) } finally {
+                releaseRPCKryoForDeserialization(kryo)
+            }
             rpcLog.debug { "<- Observable [$rpcName] <- Received $received" }
             synchronized(observables) {
                 // Force creation of the buffer if it doesn't already exist.

node/src/main/kotlin/net/corda/node/services/messaging/RPCDispatcher.kt

@@ -42,6 +42,8 @@ abstract class RPCDispatcher(val ops: RPCOps, val userService: RPCUserService, v
 
     private val queueToSubscription = HashMultimap.create<String, Subscription>()
 
+    private val handleCounter = AtomicInteger()
+
     // Created afresh for every RPC that is annotated as returning observables. Every time an observable is
     // encountered either in the RPC response or in an object graph that is being emitted by one of those
     // observables, the handle counter is incremented and the server-side observable is subscribed to. The
@@ -49,41 +51,48 @@ abstract class RPCDispatcher(val ops: RPCOps, val userService: RPCUserService, v
     //
     // When the observables are deserialised on the client side, the handle is read from the byte stream and
     // the queue is filtered to extract just those observations.
-    private inner class ObservableSerializer(private val toQName: String) : Serializer<Observable<Any>>() {
-        private val handleCounter = AtomicInteger()
+    class ObservableSerializer() : Serializer<Observable<Any>>() {
+        private fun toQName(kryo: Kryo): String = kryo.context[RPCKryoQNameKey] as String
+        private fun toDispatcher(kryo: Kryo): RPCDispatcher = kryo.context[RPCKryoDispatcherKey] as RPCDispatcher
 
         override fun read(kryo: Kryo, input: Input, type: Class<Observable<Any>>): Observable<Any> {
             throw UnsupportedOperationException("not implemented")
         }
 
         override fun write(kryo: Kryo, output: Output, obj: Observable<Any>) {
-            val handle = handleCounter.andIncrement
+            val qName = toQName(kryo)
+            val dispatcher = toDispatcher(kryo)
+            val handle = dispatcher.handleCounter.andIncrement
             output.writeInt(handle, true)
             // Observables can do three kinds of callback: "next" with a content object, "completed" and "error".
             // Materializing the observable converts these three kinds of callback into a single stream of objects
             // representing what happened, which is useful for us to send over the wire.
             val subscription = obj.materialize().subscribe { materialised: Notification<out Any> ->
-                val newKryo = createRPCKryo(observableSerializer = this@ObservableSerializer)
-                val bits = MarshalledObservation(handle, materialised).serialize(newKryo)
+                val newKryo = createRPCKryoForSerialization(qName, dispatcher)
+                val bits = try { MarshalledObservation(handle, materialised).serialize(newKryo) } finally {
+                    releaseRPCKryoForSerialization(newKryo)
+                }
                 rpcLog.debug("RPC sending observation: $materialised")
-                send(bits, toQName)
+                dispatcher.send(bits, qName)
             }
-            synchronized(queueToSubscription) {
-                queueToSubscription.put(toQName, subscription)
+            synchronized(dispatcher.queueToSubscription) {
+                dispatcher.queueToSubscription.put(qName, subscription)
             }
         }
     }
 
     fun dispatch(msg: ClientRPCRequestMessage) {
         val (argsBytes, replyTo, observationsTo, methodName) = msg
-        val kryo = createRPCKryo(observableSerializer = if (observationsTo != null) ObservableSerializer(observationsTo) else null)
 
         val response: ErrorOr<Any> = ErrorOr.catch {
             val method = methodTable[methodName] ?: throw RPCException("Received RPC for unknown method $methodName - possible client/server version skew?")
             if (method.isAnnotationPresent(RPCReturnsObservables::class.java) && observationsTo == null)
                 throw RPCException("Received RPC without any destination for observations, but the RPC returns observables")
 
-            val args = argsBytes.deserialize(kryo)
+            val kryo = createRPCKryoForSerialization(observationsTo, this)
+            val args = try { argsBytes.deserialize(kryo) } finally {
+                releaseRPCKryoForSerialization(kryo)
+            }
 
             rpcLog.debug { "-> RPC -> $methodName(${args.joinToString()})    [reply to $replyTo]" }
 
@@ -95,13 +104,15 @@ abstract class RPCDispatcher(val ops: RPCOps, val userService: RPCUserService, v
         }
         rpcLog.debug { "<- RPC <- $methodName = $response " }
 
-
         // Serialise, or send back a simple serialised ErrorOr structure if we couldn't do it.
+        val kryo = createRPCKryoForSerialization(observationsTo, this)
         val responseBits = try {
             response.serialize(kryo)
         } catch (e: KryoException) {
             rpcLog.error("Failed to respond to inbound RPC $methodName", e)
             ErrorOr.of(e).serialize(kryo)
+        } finally {
+            releaseRPCKryoForSerialization(kryo)
         }
         send(responseBits, replyTo)
     }

node/src/main/kotlin/net/corda/node/services/messaging/RPCStructures.kt

@@ -7,6 +7,7 @@ import com.esotericsoftware.kryo.Registration
 import com.esotericsoftware.kryo.Serializer
 import com.esotericsoftware.kryo.io.Input
 import com.esotericsoftware.kryo.io.Output
+import com.esotericsoftware.kryo.pool.KryoPool
 import com.google.common.util.concurrent.ListenableFuture
 import net.corda.core.flows.FlowException
 import net.corda.core.serialization.*
@@ -88,10 +89,16 @@ object ClassSerializer : Serializer<Class<*>>() {
 @CordaSerializable
 class PermissionException(msg: String) : RuntimeException(msg)
 
+object RPCKryoClientKey
+object RPCKryoDispatcherKey
+object RPCKryoQNameKey
+object RPCKryoMethodNameKey
+object RPCKryoLocationKey
+
 // The Kryo used for the RPC wire protocol. Every type in the wire protocol is listed here explicitly.
 // This is annoying to write out, but will make it easier to formalise the wire protocol when the time comes,
 // because we can see everything we're using in one place.
-private class RPCKryo(observableSerializer: Serializer<Observable<Any>>? = null) : CordaKryo(makeStandardClassResolver()) {
+private class RPCKryo(observableSerializer: Serializer<Observable<Any>>) : CordaKryo(makeStandardClassResolver()) {
     init {
         DefaultKryoCustomizer.customize(this)
 
@@ -99,49 +106,68 @@ private class RPCKryo(observableSerializer: Serializer<Observable<Any>>? = null)
         register(Class::class.java, ClassSerializer)
         register(MultipartStream.ItemInputStream::class.java, InputStreamSerializer)
         register(MarshalledObservation::class.java, ImmutableClassSerializer(MarshalledObservation::class))
-    }
-
-    // TODO: workaround to prevent Observable registration conflict when using plugin registered kyro classes
-    private val observableRegistration: Registration? = observableSerializer?.let { register(Observable::class.java, it, 10000) }
-
-    private val listenableFutureRegistration: Registration? = observableSerializer?.let {
-        // Register ListenableFuture by making use of Observable serialisation.
-        // TODO Serialisation could be made more efficient as a future can only emit one value (or exception)
+        register(Observable::class.java, observableSerializer)
         @Suppress("UNCHECKED_CAST")
         register(ListenableFuture::class,
-                read = { kryo, input -> it.read(kryo, input, Observable::class.java as Class<Observable<Any>>).toFuture() },
-                write = { kryo, output, obj -> it.write(kryo, output, obj.toObservable()) }
+                read = { kryo, input -> observableSerializer.read(kryo, input, Observable::class.java as Class<Observable<Any>>).toFuture() },
+                write = { kryo, output, obj -> observableSerializer.write(kryo, output, obj.toObservable()) }
+        )
+        register(
+                FlowException::class,
+                read = { kryo, input ->
+                    val message = input.readString()
+                    val cause = kryo.readObjectOrNull(input, Throwable::class.java)
+                    FlowException(message, cause)
+                },
+                write = { kryo, output, obj ->
+                    // The subclass may have overridden toString so we use that
+                    val message = if (obj.javaClass != FlowException::class.java) obj.toString() else obj.message
+                    output.writeString(message)
+                    kryo.writeObjectOrNull(output, obj.cause, Throwable::class.java)
+                }
         )
     }
 
-    // Avoid having to worry about the subtypes of FlowException by converting all of them to just FlowException.
-    // This is a temporary hack until a proper serialisation mechanism is in place.
-    private val flowExceptionRegistration: Registration = register(
-            FlowException::class,
-            read = { kryo, input ->
-                val message = input.readString()
-                val cause = kryo.readObjectOrNull(input, Throwable::class.java)
-                FlowException(message, cause)
-            },
-            write = { kryo, output, obj ->
-                // The subclass may have overridden toString so we use that
-                val message = if (obj.javaClass != FlowException::class.java) obj.toString() else obj.message
-                output.writeString(message)
-                kryo.writeObjectOrNull(output, obj.cause, Throwable::class.java)
-            }
-    )
-
     override fun getRegistration(type: Class<*>): Registration {
-        if (Observable::class.java.isAssignableFrom(type))
-            return observableRegistration ?:
-                    throw IllegalStateException("This RPC was not annotated with @RPCReturnsObservables")
-        if (ListenableFuture::class.java.isAssignableFrom(type))
-            return listenableFutureRegistration ?:
-                    throw IllegalStateException("This RPC was not annotated with @RPCReturnsObservables")
+        val annotated = context[RPCKryoQNameKey] != null
+        if (Observable::class.java.isAssignableFrom(type)) {
+            return if (annotated) super.getRegistration(Observable::class.java)
+            else throw IllegalStateException("This RPC was not annotated with @RPCReturnsObservables")
+        }
+        if (ListenableFuture::class.java.isAssignableFrom(type)) {
+            return if (annotated) super.getRegistration(ListenableFuture::class.java)
+            else throw IllegalStateException("This RPC was not annotated with @RPCReturnsObservables")
+        }
         if (FlowException::class.java.isAssignableFrom(type))
-            return flowExceptionRegistration
+            return super.getRegistration(FlowException::class.java)
         return super.getRegistration(type)
     }
 }
 
-fun createRPCKryo(observableSerializer: Serializer<Observable<Any>>? = null): Kryo = RPCKryo(observableSerializer)
+private val rpcSerKryoPool = KryoPool.Builder { RPCKryo(RPCDispatcher.ObservableSerializer()) }.build()
+
+fun createRPCKryoForSerialization(qName: String? = null, dispatcher: RPCDispatcher? = null): Kryo {
+    val kryo = rpcSerKryoPool.borrow()
+    kryo.context.put(RPCKryoQNameKey, qName)
+    kryo.context.put(RPCKryoDispatcherKey, dispatcher)
+    return kryo
+}
+
+fun releaseRPCKryoForSerialization(kryo: Kryo) {
+    rpcSerKryoPool.release(kryo)
+}
+
+private val rpcDesKryoPool = KryoPool.Builder { RPCKryo(CordaRPCClientImpl.ObservableDeserializer()) }.build()
+
+fun createRPCKryoForDeserialization(rpcClient: CordaRPCClientImpl, qName: String? = null, rpcName: String? = null, rpcLocation: Throwable? = null): Kryo {
+    val kryo = rpcDesKryoPool.borrow()
+    kryo.context.put(RPCKryoClientKey, rpcClient)
+    kryo.context.put(RPCKryoQNameKey, qName)
+    kryo.context.put(RPCKryoMethodNameKey, rpcName)
+    kryo.context.put(RPCKryoLocationKey, rpcLocation)
+    return kryo
+}
+
+fun releaseRPCKryoForDeserialization(kryo: Kryo) {
+    rpcDesKryoPool.release(kryo)
+}

node/src/main/kotlin/net/corda/node/services/persistence/DBCheckpointStorage.kt

@@ -4,7 +4,7 @@ import net.corda.core.crypto.SecureHash
 import net.corda.core.serialization.SerializedBytes
 import net.corda.core.serialization.deserialize
 import net.corda.core.serialization.serialize
-import net.corda.core.serialization.threadLocalStorageKryo
+import net.corda.core.serialization.storageKryo
 import net.corda.node.services.api.Checkpoint
 import net.corda.node.services.api.CheckpointStorage
 import net.corda.node.utilities.*
@@ -39,7 +39,7 @@ class DBCheckpointStorage : CheckpointStorage {
     private val checkpointStorage = synchronizedMap(CheckpointMap())
 
     override fun addCheckpoint(checkpoint: Checkpoint) {
-        checkpointStorage.put(checkpoint.id, checkpoint.serialize(threadLocalStorageKryo(), true))
+        checkpointStorage.put(checkpoint.id, checkpoint.serialize(storageKryo(), true))
     }
 
     override fun removeCheckpoint(checkpoint: Checkpoint) {

node/src/main/kotlin/net/corda/node/services/statemachine/StateMachineManager.kt

@@ -6,6 +6,7 @@ import co.paralleluniverse.io.serialization.kryo.KryoSerializer
 import co.paralleluniverse.strands.Strand
 import com.codahale.metrics.Gauge
 import com.esotericsoftware.kryo.Kryo
+import com.esotericsoftware.kryo.pool.KryoPool
 import com.google.common.collect.HashMultimap
 import com.google.common.util.concurrent.ListenableFuture
 import kotlinx.support.jdk8.collections.removeIf
@@ -71,6 +72,11 @@ class StateMachineManager(val serviceHub: ServiceHubInternal,
 
     inner class FiberScheduler : FiberExecutorScheduler("Same thread scheduler", executor)
 
+    private val quasarKryoPool = KryoPool.Builder {
+        val serializer = Fiber.getFiberSerializer(false) as KryoSerializer
+        DefaultKryoCustomizer.customize(serializer.kryo)
+    }.build()
+
     companion object {
         private val logger = loggerFor<StateMachineManager>()
         internal val sessionTopic = TopicSession("platform.session")
@@ -354,32 +360,23 @@ class StateMachineManager(val serviceHub: ServiceHubInternal,
     }
 
     private fun serializeFiber(fiber: FlowStateMachineImpl<*>): SerializedBytes<FlowStateMachineImpl<*>> {
-        val kryo = quasarKryo()
-        // add the map of tokens -> tokenizedServices to the kyro context
-        SerializeAsTokenSerializer.setContext(kryo, serializationContext)
-        return fiber.serialize(kryo)
+        return quasarKryo().run { kryo ->
+            // add the map of tokens -> tokenizedServices to the kyro context
+            SerializeAsTokenSerializer.setContext(kryo, serializationContext)
+            fiber.serialize(kryo)
+        }
     }
 
     private fun deserializeFiber(checkpoint: Checkpoint): FlowStateMachineImpl<*> {
-        val kryo = quasarKryo()
-        // put the map of token -> tokenized into the kryo context
-        SerializeAsTokenSerializer.setContext(kryo, serializationContext)
-        return checkpoint.serializedFiber.deserialize(kryo).apply { fromCheckpoint = true }
-    }
-
-    private fun quasarKryo(): Kryo {
-        val serializer = Fiber.getFiberSerializer(false) as KryoSerializer
-        return createKryo(serializer.kryo).apply {
-            // Because we like to stick a Kryo object in a ThreadLocal to speed things up a bit, we can end up trying to
-            // serialise the Kryo object itself when suspending a fiber. That's dumb, useless AND can cause crashes, so
-            // we avoid it here.  This is checkpointing specific.
-            register(Kryo::class,
-                    read = { kryo, input -> createKryo((Fiber.getFiberSerializer() as KryoSerializer).kryo) },
-                    write = { kryo, output, obj -> }
-            )
+        return quasarKryo().run { kryo ->
+            // put the map of token -> tokenized into the kryo context
+            SerializeAsTokenSerializer.setContext(kryo, serializationContext)
+            checkpoint.serializedFiber.deserialize(kryo).apply { fromCheckpoint = true }
         }
     }
 
+    private fun quasarKryo(): KryoPool = quasarKryoPool
+
     private fun <T> createFiber(logic: FlowLogic<T>): FlowStateMachineImpl<T> {
         val id = StateMachineRunId.createRandom()
         return FlowStateMachineImpl(id, logic, scheduler).apply { initFiber(this) }

node/src/main/kotlin/net/corda/node/services/vault/NodeVaultService.kt

@@ -8,7 +8,6 @@ import net.corda.core.ThreadBox
 import net.corda.core.bufferUntilSubscribed
 import net.corda.core.contracts.*
 import net.corda.core.crypto.AbstractParty
-import net.corda.core.crypto.AnonymousParty
 import net.corda.core.crypto.CompositeKey
 import net.corda.core.crypto.SecureHash
 import net.corda.core.node.ServiceHub
@@ -16,9 +15,9 @@ import net.corda.core.node.services.Vault
 import net.corda.core.node.services.VaultService
 import net.corda.core.node.services.unconsumedStates
 import net.corda.core.serialization.SingletonSerializeAsToken
-import net.corda.core.serialization.createKryo
 import net.corda.core.serialization.deserialize
 import net.corda.core.serialization.serialize
+import net.corda.core.serialization.storageKryo
 import net.corda.core.tee
 import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.transactions.WireTransaction
@@ -76,8 +75,7 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
                             index = it.key.index
                             stateStatus = Vault.StateStatus.UNCONSUMED
                             contractStateClassName = it.value.state.data.javaClass.name
-                            // TODO: revisit Kryo bug when using THREAD_LOCAL_KYRO
-                            contractState = it.value.state.serialize(createKryo()).bytes
+                            contractState = it.value.state.serialize(storageKryo()).bytes
                             notaryName = it.value.state.notary.name
                             notaryKey = it.value.state.notary.owningKey.toBase58String()
                             recordedTime = services.clock.instant()
@@ -165,8 +163,7 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
                 Sequence{iterator}
                         .map { it ->
                             val stateRef = StateRef(SecureHash.parse(it.txId), it.index)
-                            // TODO: revisit Kryo bug when using THREAD_LOCAL_KRYO
-                            val state = it.contractState.deserialize<TransactionState<T>>(createKryo())
+                            val state = it.contractState.deserialize<TransactionState<T>>(storageKryo())
                             StateAndRef(state, stateRef)
                         }
             }
@@ -184,7 +181,7 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
                                 .and(VaultSchema.VaultStates::index eq it.index)
                         result.get()?.each {
                             val stateRef = StateRef(SecureHash.parse(it.txId), it.index)
-                            val state = it.contractState.deserialize<TransactionState<*>>()
+                            val state = it.contractState.deserialize<TransactionState<*>>(storageKryo())
                             results += StateAndRef(state, stateRef)
                         }
                     }
@@ -353,7 +350,7 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
             while (rs.next()) {
                 val txHash = SecureHash.parse(rs.getString(1))
                 val index = rs.getInt(2)
-                val state = rs.getBytes(3).deserialize<TransactionState<ContractState>>(createKryo())
+                val state = rs.getBytes(3).deserialize<TransactionState<ContractState>>(storageKryo())
                 consumedStates.add(StateAndRef(state, StateRef(txHash, index)))
             }
         }

node/src/main/kotlin/net/corda/node/utilities/JDBCHashMap.kt

@@ -3,7 +3,7 @@ package net.corda.node.utilities
 import net.corda.core.serialization.SerializedBytes
 import net.corda.core.serialization.deserialize
 import net.corda.core.serialization.serialize
-import net.corda.core.serialization.threadLocalStorageKryo
+import net.corda.core.serialization.storageKryo
 import net.corda.core.utilities.loggerFor
 import net.corda.core.utilities.trace
 import org.jetbrains.exposed.sql.*
@@ -65,7 +65,7 @@ fun bytesToBlob(value: SerializedBytes<*>, finalizables: MutableList<() -> Unit>
     return blob
 }
 
-fun serializeToBlob(value: Any, finalizables: MutableList<() -> Unit>): Blob = bytesToBlob(value.serialize(threadLocalStorageKryo(), true), finalizables)
+fun serializeToBlob(value: Any, finalizables: MutableList<() -> Unit>): Blob = bytesToBlob(value.serialize(storageKryo(), true), finalizables)
 
 fun <T : Any> bytesFromBlob(blob: Blob): SerializedBytes<T> {
     try {

node/src/test/kotlin/net/corda/node/services/database/RequeryConfigurationTest.kt

@@ -10,8 +10,8 @@ import net.corda.core.crypto.DigitalSignature
 import net.corda.core.crypto.NullPublicKey
 import net.corda.core.crypto.SecureHash
 import net.corda.core.node.services.Vault
-import net.corda.core.serialization.createKryo
 import net.corda.core.serialization.serialize
+import net.corda.core.serialization.storageKryo
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.WireTransaction
 import net.corda.core.utilities.DUMMY_NOTARY
@@ -128,7 +128,7 @@ class RequeryConfigurationTest {
             index = txnState.index
             stateStatus = Vault.StateStatus.UNCONSUMED
             contractStateClassName = DummyContract.SingleOwnerState::class.java.name
-            contractState = DummyContract.SingleOwnerState(owner = DUMMY_PUBKEY_1).serialize(createKryo()).bytes
+            contractState = DummyContract.SingleOwnerState(owner = DUMMY_PUBKEY_1).serialize(storageKryo()).bytes
             notaryName = txn.tx.notary!!.name
             notaryKey = txn.tx.notary!!.owningKey.toBase58String()
             recordedTime = Instant.now()