Remove SAME_THREAD executor and it's use in MockNetwork etc.

Remove all traces of unused optional Executor in messaging.

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

@@ -4,7 +4,6 @@ import com.codahale.metrics.MetricRegistry
 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.crypto.Party
 import com.r3corda.core.crypto.X509Utilities
 import com.r3corda.core.messaging.SingleMessageRecipient
@@ -374,7 +373,7 @@ abstract class AbstractNode(open val configuration: NodeConfiguration, val netwo
         val reg = NodeRegistration(info, instant.toEpochMilli(), type, expires)
         val legalIdentityKey = obtainLegalIdentityKey()
         val request = NetworkMapService.RegistrationRequest(reg.toWire(legalIdentityKey.private), net.myAddress)
-        return net.sendRequest(REGISTER_PROTOCOL_TOPIC, request, networkMapAddr, RunOnCallerThread)
+        return net.sendRequest(REGISTER_PROTOCOL_TOPIC, request, networkMapAddr)
     }
 
     protected open fun makeKeyManagementService(): KeyManagementService = PersistentKeyManagementService(partyKeys)

node/src/main/kotlin/com/r3corda/node/services/api/AbstractNodeService.kt

@@ -31,7 +31,7 @@ abstract class AbstractNodeService(val services: ServiceHubInternal) : Singleton
             addMessageHandler(topic: String,
                               crossinline handler: (Q) -> R,
                               crossinline exceptionConsumer: (Message, Exception) -> Unit): MessageHandlerRegistration {
-        return net.addMessageHandler(topic, DEFAULT_SESSION_ID, null) { message, r ->
+        return net.addMessageHandler(topic, DEFAULT_SESSION_ID) { message, r ->
             try {
                 val request = message.data.deserialize<Q>()
                 val response = handler(request)

node/src/main/kotlin/com/r3corda/node/services/messaging/NodeMessagingClient.kt

@@ -21,7 +21,6 @@ import java.time.Instant
 import java.util.*
 import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.CountDownLatch
-import java.util.concurrent.Executor
 import javax.annotation.concurrent.ThreadSafe
 
 // TODO: Stop the wallet explorer and other clients from using this class and get rid of persistentInbox
@@ -92,8 +91,7 @@ class NodeMessagingClient(config: NodeConfiguration,
     }
 
     /** A registration to handle messages of different types */
-    data class Handler(val executor: Executor?,
-                       val topicSession: TopicSession,
+    data class Handler(val topicSession: TopicSession,
                        val callback: (Message, MessageHandlerRegistration) -> Unit) : MessageHandlerRegistration
 
     /**
@@ -369,15 +367,13 @@ class NodeMessagingClient(config: NodeConfiguration,
         }
     }
 
-    override fun addMessageHandler(topic: String, sessionID: Long, executor: Executor?,
-                                   callback: (Message, MessageHandlerRegistration) -> Unit): MessageHandlerRegistration
-            = addMessageHandler(TopicSession(topic, sessionID), executor, callback)
+    override fun addMessageHandler(topic: String, sessionID: Long, callback: (Message, MessageHandlerRegistration) -> Unit): MessageHandlerRegistration
+            = addMessageHandler(TopicSession(topic, sessionID), callback)
 
     override fun addMessageHandler(topicSession: TopicSession,
-                                   executor: Executor?,
                                    callback: (Message, MessageHandlerRegistration) -> Unit): MessageHandlerRegistration {
         require(!topicSession.isBlank()) { "Topic must not be blank, as the empty topic is a special case." }
-        val handler = Handler(executor, topicSession, callback)
+        val handler = Handler(topicSession, callback)
         handlers.add(handler)
         val messagesToRedeliver = state.locked {
             val messagesToRedeliver = undeliveredMessages

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

@@ -3,7 +3,6 @@ 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.SettableFuture
-import com.r3corda.core.RunOnCallerThread
 import com.r3corda.core.contracts.Contract
 import com.r3corda.core.crypto.Party
 import com.r3corda.core.map
@@ -20,12 +19,10 @@ import com.r3corda.core.node.services.ServiceType
 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 com.r3corda.protocols.sendRequest
 import rx.Observable
@@ -70,7 +67,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
                                ifChangedSinceVer: Int?): ListenableFuture<Unit> {
         if (subscribe && !registeredForPush) {
             // Add handler to the network, for updates received from the remote network map service.
-            net.addMessageHandler(NetworkMapService.PUSH_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, null) { message, r ->
+            net.addMessageHandler(NetworkMapService.PUSH_PROTOCOL_TOPIC, DEFAULT_SESSION_ID) { message, r ->
                 try {
                     val req = message.data.deserialize<NetworkMapService.Update>()
                     val ackMessage = net.createMessage(NetworkMapService.PUSH_ACK_PROTOCOL_TOPIC, DEFAULT_SESSION_ID,
@@ -88,7 +85,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
 
         // Fetch the network map and register for updates at the same time
         val req = NetworkMapService.FetchMapRequest(subscribe, ifChangedSinceVer, net.myAddress)
-        val future = net.sendRequest<FetchMapResponse>(FETCH_PROTOCOL_TOPIC, req, networkMapAddress, RunOnCallerThread).map { resp ->
+        val future = net.sendRequest<FetchMapResponse>(FETCH_PROTOCOL_TOPIC, req, networkMapAddress).map { resp ->
             // We may not receive any nodes back, if the map hasn't changed since the version specified
             resp.nodes?.forEach { processRegistration(it) }
             Unit
@@ -120,7 +117,7 @@ open class InMemoryNetworkMapCache : SingletonSerializeAsToken(), NetworkMapCach
     override fun deregisterForUpdates(net: MessagingService, service: NodeInfo): ListenableFuture<Unit> {
         // Fetch the network map and register for updates at the same time
         val req = NetworkMapService.SubscribeRequest(false, net.myAddress)
-        val future = net.sendRequest<SubscribeResponse>(SUBSCRIPTION_PROTOCOL_TOPIC, req, service.address, RunOnCallerThread).map {
+        val future = net.sendRequest<SubscribeResponse>(SUBSCRIPTION_PROTOCOL_TOPIC, req, service.address).map {
             if (it.confirmed) Unit else throw NetworkCacheError.DeregistrationFailed()
         }
         _registrationFuture.setFuture(future)

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

@@ -154,7 +154,7 @@ abstract class AbstractNetworkMapService
         handlers += addMessageHandler(NetworkMapService.SUBSCRIPTION_PROTOCOL_TOPIC,
                 { req: NetworkMapService.SubscribeRequest -> processSubscriptionRequest(req) }
         )
-        handlers += net.addMessageHandler(NetworkMapService.PUSH_ACK_PROTOCOL_TOPIC, DEFAULT_SESSION_ID, null) { message, r ->
+        handlers += net.addMessageHandler(NetworkMapService.PUSH_ACK_PROTOCOL_TOPIC, DEFAULT_SESSION_ID) { message, r ->
             val req = message.data.deserialize<NetworkMapService.UpdateAcknowledge>()
             processAcknowledge(req)
         }

node/src/main/kotlin/com/r3corda/node/services/statemachine/StateMachineManager.kt

@@ -175,7 +175,7 @@ class StateMachineManager(val serviceHub: ServiceHubInternal,
             started = true
             stateMachines.keys.forEach { resumeRestoredFiber(it) }
         }
-        serviceHub.networkService.addMessageHandler(sessionTopic, executor) { message, reg ->
+        serviceHub.networkService.addMessageHandler(sessionTopic) { message, reg ->
             executor.checkOnThread()
             val sessionMessage = message.data.deserialize<SessionMessage>()
             when (sessionMessage) {

node/src/main/kotlin/com/r3corda/node/utilities/AffinityExecutor.kt

@@ -1,5 +1,6 @@
 package com.r3corda.node.utilities
 
+import com.google.common.util.concurrent.SettableFuture
 import com.google.common.util.concurrent.Uninterruptibles
 import com.r3corda.core.utilities.loggerFor
 import java.util.*
@@ -41,19 +42,20 @@ interface AffinityExecutor : Executor {
     }
 
     /**
-     * Posts a no-op task to the executor and blocks this thread waiting for it to complete. This can be useful in
-     * tests when you want to be sure that a previous task submitted via [execute] has completed.
+     * Run the executor until there are no tasks pending and none executing.
      */
-    fun flush() {
-        fetchFrom { }
-    }
+    fun flush()
 
     /**
      * An executor backed by thread pool (which may often have a single thread) which makes it easy to schedule
      * tasks in the future and verify code is running on the executor.
      */
-    class ServiceAffinityExecutor(threadName: String, numThreads: Int) : AffinityExecutor,
+    open class ServiceAffinityExecutor(threadName: String, numThreads: Int) : AffinityExecutor,
             ThreadPoolExecutor(numThreads, numThreads, 0L, TimeUnit.MILLISECONDS, LinkedBlockingQueue<Runnable>()) {
+        companion object {
+            val logger = loggerFor<ServiceAffinityExecutor>()
+        }
+
         private val threads = Collections.synchronizedSet(HashSet<Thread>())
         private val uncaughtExceptionHandler = Thread.currentThread().uncaughtExceptionHandler
 
@@ -82,8 +84,11 @@ interface AffinityExecutor : Executor {
 
         override val isOnThread: Boolean get() = Thread.currentThread() in threads
 
-        companion object {
-            val logger = loggerFor<ServiceAffinityExecutor>()
+        override fun flush() {
+            do {
+                val f = SettableFuture.create<Boolean>()
+                execute { f.set(queue.isEmpty() && activeCount == 1) }
+            } while (!f.get())
         }
     }
 
@@ -110,12 +115,9 @@ interface AffinityExecutor : Executor {
         }
 
         val taskQueueSize: Int get() = commandQ.size
-    }
 
-    companion object {
-        val SAME_THREAD: AffinityExecutor = object : AffinityExecutor {
-            override val isOnThread: Boolean get() = true
-            override fun execute(command: Runnable) = command.run()
+        override fun flush() {
+            throw UnsupportedOperationException()
         }
     }
 }

node/src/test/kotlin/com/r3corda/node/services/ArtemisMessagingTests.kt

@@ -117,7 +117,7 @@ class ArtemisMessagingTests {
     }
 
     private fun createMessagingClient(server: HostAndPort = hostAndPort): NodeMessagingClient {
-        return NodeMessagingClient(config, server, identity.public, AffinityExecutor.SAME_THREAD, false).apply {
+        return NodeMessagingClient(config, server, identity.public, AffinityExecutor.ServiceAffinityExecutor("ArtemisMessagingTests", 1), false).apply {
             configureWithDevSSLCertificate()
             messagingClient = this
         }

node/src/test/kotlin/com/r3corda/node/services/NodeSchedulerServiceTest.kt

@@ -71,7 +71,7 @@ class NodeSchedulerServiceTest : SingletonSerializeAsToken() {
 
     init {
         val kms = MockKeyManagementService(ALICE_KEY)
-        val mockMessagingService = InMemoryMessagingNetwork(false).InMemoryMessaging(false, InMemoryMessagingNetwork.Handle(0, "None"), persistenceTx = { it() })
+        val mockMessagingService = InMemoryMessagingNetwork(false).InMemoryMessaging(false, InMemoryMessagingNetwork.Handle(0, "None"), AffinityExecutor.ServiceAffinityExecutor("test", 1), persistenceTx = { it() })
         services = object : MockServiceHubInternal(overrideClock = testClock, keyManagement = kms, net = mockMessagingService), TestReference {
             override val testReference = this@NodeSchedulerServiceTest
         }

node/src/test/kotlin/com/r3corda/node/services/statemachine/StateMachineManagerTests.kt

@@ -117,6 +117,8 @@ class StateMachineManagerTests {
         val payload = random63BitValue()
         node1.services.registerProtocolInitiator(ReceiveThenSuspendProtocol::class) { SendProtocol(payload, it) }
         node2.smm.add(ReceiveThenSuspendProtocol(node1.info.legalIdentity)) // Prepare checkpointed receive protocol
+        // Make sure the add() has finished initial processing.
+        node2.smm.executor.flush()
         node2.stop() // kill receiver
         val restoredProtocol = node2.restartAndGetRestoredProtocol<ReceiveThenSuspendProtocol>(node1)
         assertThat(restoredProtocol.receivedPayloads[0]).isEqualTo(payload)
@@ -137,6 +139,8 @@ class StateMachineManagerTests {
         // Kick off first send and receive
         node2.smm.add(PingPongProtocol(node3.info.legalIdentity, payload))
         assertEquals(1, node2.checkpointStorage.checkpoints().size)
+        // Make sure the add() has finished initial processing.
+        node2.smm.executor.flush()
         // Restart node and thus reload the checkpoint and resend the message with same UUID
         node2.stop()
         val node2b = net.createNode(node1.info.address, node2.id, advertisedServices = *node2.advertisedServices.toTypedArray())

node/src/test/kotlin/com/r3corda/node/utilities/AffinityExecutorTests.kt

@@ -3,7 +3,6 @@ package com.r3corda.node.utilities
 import org.junit.After
 import org.junit.Test
 import java.util.*
-import java.util.concurrent.CompletableFuture
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.atomic.AtomicReference
 import kotlin.concurrent.thread
@@ -12,30 +11,30 @@ import kotlin.test.assertFails
 import kotlin.test.assertNotEquals
 
 class AffinityExecutorTests {
-    @Test fun `AffinityExecutor SAME_THREAD executes on calling thread`() {
-        assert(AffinityExecutor.SAME_THREAD.isOnThread)
-
-        run {
-            val thatThread = CompletableFuture<Thread>()
-            AffinityExecutor.SAME_THREAD.execute { thatThread.complete(Thread.currentThread()) }
-            assertEquals(Thread.currentThread(), thatThread.get())
-        }
-        run {
-            val thatThread = CompletableFuture<Thread>()
-            AffinityExecutor.SAME_THREAD.executeASAP { thatThread.complete(Thread.currentThread()) }
-            assertEquals(Thread.currentThread(), thatThread.get())
-        }
-    }
-
-    var executor: AffinityExecutor.ServiceAffinityExecutor? = null
+    var _executor: AffinityExecutor.ServiceAffinityExecutor? = null
+    val executor: AffinityExecutor.ServiceAffinityExecutor get() = _executor!!
 
     @After fun shutdown() {
-        executor?.shutdown()
+        _executor?.shutdown()
+        _executor = null
+    }
+
+    @Test fun `flush handles nested executes`() {
+        _executor = AffinityExecutor.ServiceAffinityExecutor("test4", 1)
+        var nestedRan = false
+        val latch = CountDownLatch(1)
+        executor.execute {
+            latch.await()
+            executor.execute { nestedRan = true }
+        }
+        latch.countDown()
+        executor.flush()
+        assert(nestedRan)
     }
 
     @Test fun `single threaded affinity executor runs on correct thread`() {
         val thisThread = Thread.currentThread()
-        val executor = AffinityExecutor.ServiceAffinityExecutor("test thread", 1)
+        _executor = AffinityExecutor.ServiceAffinityExecutor("test thread", 1)
         assert(!executor.isOnThread)
         assertFails { executor.checkOnThread() }
 
@@ -55,7 +54,7 @@ class AffinityExecutorTests {
     }
 
     @Test fun `pooled executor`() {
-        val executor = AffinityExecutor.ServiceAffinityExecutor("test2", 3)
+        _executor = AffinityExecutor.ServiceAffinityExecutor("test2", 3)
         assert(!executor.isOnThread)
 
         val latch = CountDownLatch(1)
@@ -89,7 +88,7 @@ class AffinityExecutorTests {
         // Run in a separate thread to avoid messing with any default exception handlers in the unit test thread.
         thread {
             Thread.currentThread().setUncaughtExceptionHandler { thread, throwable -> exception.set(throwable) }
-            val executor = AffinityExecutor.ServiceAffinityExecutor("test3", 1)
+            _executor = AffinityExecutor.ServiceAffinityExecutor("test3", 1)
             executor.execute {
                 throw Exception("foo")
             }