Better use of Futures

node/src/main/kotlin/com/r3corda/node/internal/AbstractNode.kt

@@ -9,7 +9,7 @@ import com.r3corda.core.crypto.Party
 import com.r3corda.core.crypto.X509Utilities
 import com.r3corda.core.messaging.SingleMessageRecipient
 import com.r3corda.core.messaging.createMessage
-import com.r3corda.core.messaging.runOnNextMessage
+import com.r3corda.core.messaging.onNext
 import com.r3corda.core.node.CityDatabase
 import com.r3corda.core.node.CordaPluginRegistry
 import com.r3corda.core.node.NodeInfo
@@ -37,6 +37,7 @@ import com.r3corda.node.services.monitor.NodeMonitorService
 import com.r3corda.node.services.network.InMemoryNetworkMapCache
 import com.r3corda.node.services.network.NetworkMapService
 import com.r3corda.node.services.network.NetworkMapService.Companion.REGISTER_PROTOCOL_TOPIC
+import com.r3corda.node.services.network.NetworkMapService.RegistrationResponse
 import com.r3corda.node.services.network.NodeRegistration
 import com.r3corda.node.services.network.PersistentNetworkMapService
 import com.r3corda.node.services.persistence.*
@@ -357,7 +358,7 @@ abstract class AbstractNode(val configuration: NodeConfiguration, val networkMap
                 "has any other map node been configured.")
     }
 
-    private fun updateRegistration(networkMapAddr: SingleMessageRecipient, type: AddOrRemove): ListenableFuture<NetworkMapService.RegistrationResponse> {
+    private fun updateRegistration(networkMapAddr: SingleMessageRecipient, type: AddOrRemove): ListenableFuture<RegistrationResponse> {
         // Register this node against the network
         val instant = platformClock.instant()
         val expires = instant + NetworkMapService.DEFAULT_EXPIRATION_PERIOD
@@ -365,13 +366,8 @@ abstract class AbstractNode(val configuration: NodeConfiguration, val networkMap
         val sessionID = random63BitValue()
         val request = NetworkMapService.RegistrationRequest(reg.toWire(storage.myLegalIdentityKey.private), net.myAddress, sessionID)
         val message = net.createMessage(REGISTER_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, request.serialize().bits)
-        val future = SettableFuture.create<NetworkMapService.RegistrationResponse>()
-
-        net.runOnNextMessage(REGISTER_PROTOCOL_TOPIC, sessionID, RunOnCallerThread) { message ->
-            future.set(message.data.deserialize())
-        }
+        val future = net.onNext<RegistrationResponse>(REGISTER_PROTOCOL_TOPIC, sessionID, RunOnCallerThread)
         net.send(message, networkMapAddr)
-
         return future
     }
 

node/src/main/kotlin/com/r3corda/node/services/network/InMemoryNetworkMapCache.kt

@@ -2,10 +2,11 @@ package com.r3corda.node.services.network
 
 import com.google.common.annotations.VisibleForTesting
 import com.google.common.util.concurrent.ListenableFuture
-import com.google.common.util.concurrent.MoreExecutors
 import com.google.common.util.concurrent.SettableFuture
+import com.r3corda.core.RunOnCallerThread
 import com.r3corda.core.contracts.Contract
 import com.r3corda.core.crypto.Party
+import com.r3corda.core.map
 import com.r3corda.core.messaging.*
 import com.r3corda.core.node.NodeInfo
 import com.r3corda.core.node.services.DEFAULT_SESSION_ID
@@ -19,6 +20,10 @@ import com.r3corda.core.serialization.SingletonSerializeAsToken
 import com.r3corda.core.serialization.deserialize
 import com.r3corda.core.serialization.serialize
 import com.r3corda.node.services.api.RegulatorService
+import com.r3corda.node.services.network.NetworkMapService.Companion.FETCH_PROTOCOL_TOPIC
+import com.r3corda.node.services.network.NetworkMapService.Companion.SUBSCRIPTION_PROTOCOL_TOPIC
+import com.r3corda.node.services.network.NetworkMapService.FetchMapResponse
+import com.r3corda.node.services.network.NetworkMapService.SubscribeResponse
 import com.r3corda.node.services.transactions.NotaryService
 import com.r3corda.node.utilities.AddOrRemove
 import rx.Observable
@@ -82,15 +87,13 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
 
         // Add a message handler for the response, and prepare a future to put the data into.
         // Note that the message handler will run on the network thread (not this one).
-        val future = SettableFuture.create<Unit>()
-        _registrationFuture.setFuture(future)
-        net.runOnNextMessage(NetworkMapService.FETCH_PROTOCOL_TOPIC, sessionID, MoreExecutors.directExecutor()) { message ->
-            val resp = message.data.deserialize<NetworkMapService.FetchMapResponse>()
+        val future = net.onNext<FetchMapResponse>(FETCH_PROTOCOL_TOPIC, sessionID, RunOnCallerThread).map { resp ->
             // We may not receive any nodes back, if the map hasn't changed since the version specified
             resp.nodes?.forEach { processRegistration(it) }
-            future.set(Unit)
+            Unit
         }
-        net.send(NetworkMapService.FETCH_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, req, networkMapAddress)
+        net.send(FETCH_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, req, networkMapAddress)
+        _registrationFuture.setFuture(future)
 
         return future
     }
@@ -121,29 +124,24 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
 
         // Add a message handler for the response, and prepare a future to put the data into.
         // Note that the message handler will run on the network thread (not this one).
-        val future = SettableFuture.create<Unit>()
-        _registrationFuture.setFuture(future)
-        net.runOnNextMessage(NetworkMapService.SUBSCRIPTION_PROTOCOL_TOPIC, sessionID, MoreExecutors.directExecutor()) { message ->
-            val resp = message.data.deserialize<NetworkMapService.SubscribeResponse>()
-            if (resp.confirmed) {
-                future.set(Unit)
-            } else {
-                future.setException(NetworkCacheError.DeregistrationFailed())
-            }
+
+        val future = net.onNext<SubscribeResponse>(SUBSCRIPTION_PROTOCOL_TOPIC, sessionID, RunOnCallerThread).map {
+            if (it.confirmed) Unit else throw NetworkCacheError.DeregistrationFailed()
         }
-        net.send(NetworkMapService.SUBSCRIPTION_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, req, service.address)
+        _registrationFuture.setFuture(future)
+
+        net.send(SUBSCRIPTION_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, req, service.address)
 
         return future
     }
 
     fun processUpdatePush(req: NetworkMapService.Update) {
-        val reg: NodeRegistration
         try {
-            reg = req.wireReg.verified()
-        } catch(e: SignatureException) {
+            val reg = req.wireReg.verified()
+            processRegistration(reg)
+        } catch (e: SignatureException) {
             throw NodeMapError.InvalidSignature()
         }
-        processRegistration(reg)
     }
 
     private fun processRegistration(reg: NodeRegistration) {

node/src/test/kotlin/com/r3corda/node/services/persistence/PerFileTransactionStorageTests.kt

@@ -1,25 +1,27 @@
 package com.r3corda.node.services.persistence
 
-import co.paralleluniverse.strands.SettableFuture
 import com.google.common.jimfs.Configuration.unix
 import com.google.common.jimfs.Jimfs
 import com.google.common.primitives.Ints
-import com.r3corda.core.transactions.SignedTransaction
+import com.google.common.util.concurrent.SettableFuture
 import com.r3corda.core.crypto.DigitalSignature
 import com.r3corda.core.crypto.NullPublicKey
 import com.r3corda.core.serialization.SerializedBytes
+import com.r3corda.core.transactions.SignedTransaction
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
+import java.nio.file.FileSystem
 import java.nio.file.Files
+import java.nio.file.Path
 import java.util.concurrent.TimeUnit
 import kotlin.test.assertEquals
 
 class PerFileTransactionStorageTests {
 
-    val fileSystem = Jimfs.newFileSystem(unix())
-    val storeDir = fileSystem.getPath("store")
+    val fileSystem: FileSystem = Jimfs.newFileSystem(unix())
+    val storeDir: Path = fileSystem.getPath("store")
     lateinit var transactionStorage: PerFileTransactionStorage
 
     @Before
@@ -74,7 +76,7 @@ class PerFileTransactionStorageTests {
 
     @Test
     fun `updates are fired`() {
-        val future = SettableFuture<SignedTransaction>()
+        val future = SettableFuture.create<SignedTransaction>()
         transactionStorage.updates.subscribe { tx -> future.set(tx) }
         val expected = newTransaction()
         transactionStorage.addTransaction(expected)