Lazy NodeVaultService.states (#349)

Convert NodeVaultService states to return Iterable (backed by Sequence) Vs the old way using a List. Worth noting this relieves memory pressure as the number of vault states grows.

* remove toList in ContractUpgradeFlowTest

core/src/main/kotlin/net/corda/core/node/services/Services.kt

@@ -203,15 +203,15 @@ interface VaultService {
                       onlyFromParties: Set<AbstractParty>? = null): Pair<TransactionBuilder, List<CompositeKey>>
 
     /**
-     * Return [ContractState]s of a given [Contract] type and list of [Vault.StateStatus]
+     * Return [ContractState]s of a given [Contract] type and [Iterable] of [Vault.StateStatus].
      */
-    fun <T : ContractState> states(clazzes: Set<Class<T>>, statuses: EnumSet<Vault.StateStatus>): List<StateAndRef<T>>
+    fun <T : ContractState> states(clazzes: Set<Class<T>>, statuses: EnumSet<Vault.StateStatus>): Iterable<StateAndRef<T>>
 }
 
-inline fun <reified T: ContractState> VaultService.unconsumedStates(): List<StateAndRef<T>> =
+inline fun <reified T: ContractState> VaultService.unconsumedStates(): Iterable<StateAndRef<T>> =
         states(setOf(T::class.java), EnumSet.of(Vault.StateStatus.UNCONSUMED))
 
-inline fun <reified T: ContractState> VaultService.consumedStates(): List<StateAndRef<T>> =
+inline fun <reified T: ContractState> VaultService.consumedStates(): Iterable<StateAndRef<T>> =
         states(setOf(T::class.java), EnumSet.of(Vault.StateStatus.CONSUMED))
 
 /** Returns the [linearState] heads only when the type of the state would be considered an 'instanceof' the given type. */

core/src/test/kotlin/net/corda/core/flows/ContractUpgradeFlowTest.kt

@@ -165,11 +165,11 @@ class ContractUpgradeFlowTest {
         // Starts contract upgrade flow.
         a.services.startFlow(ContractUpgradeFlow.Instigator(stateAndRef, CashV2::class.java))
         mockNet.runNetwork()
-        // Get contract state form the vault.
+        // Get contract state from the vault.
-        val state = databaseTransaction(a.database) { a.vault.unconsumedStates<ContractState>() }
+        val firstState = databaseTransaction(a.database) { a.vault.unconsumedStates<ContractState>().single() }
-        assertTrue(state.single().state.data is CashV2.State, "Contract state is upgraded to the new version.")
+        assertTrue(firstState.state.data is CashV2.State, "Contract state is upgraded to the new version.")
-        assertEquals(Amount(1000000, USD).`issued by`(a.info.legalIdentity.ref(1)), (state.first().state.data as CashV2.State).amount, "Upgraded cash contain the correct amount.")
+        assertEquals(Amount(1000000, USD).`issued by`(a.info.legalIdentity.ref(1)), (firstState.state.data as CashV2.State).amount, "Upgraded cash contain the correct amount.")
-        assertEquals(listOf(a.info.legalIdentity.owningKey), (state.first().state.data as CashV2.State).owners, "Upgraded cash belongs to the right owner.")
+        assertEquals(listOf(a.info.legalIdentity.owningKey), (firstState.state.data as CashV2.State).owners, "Upgraded cash belongs to the right owner.")
     }
 
     class CashV2 : UpgradedContract<Cash.State, CashV2.State> {

finance/src/test/kotlin/net/corda/contracts/asset/CashTests.kt

@@ -78,7 +78,7 @@ class CashTests {
             services.fillWithSomeTestCash(howMuch = 80.SWISS_FRANCS, atLeastThisManyStates = 1, atMostThisManyStates = 1,
                     issuedBy = MINI_CORP.ref(1), issuerKey = MINI_CORP_KEY, ownedBy = OUR_PUBKEY_1)
 
-            vaultStatesUnconsumed = services.vaultService.unconsumedStates<Cash.State>()
+            vaultStatesUnconsumed = services.vaultService.unconsumedStates<Cash.State>().toList()
         }
     }
 

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

@@ -153,7 +153,7 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
         }
     }
 
-    override fun <T: ContractState> states(clazzes: Set<Class<T>>, statuses: EnumSet<Vault.StateStatus>): List<StateAndRef<T>> {
+    override fun <T: ContractState> states(clazzes: Set<Class<T>>, statuses: EnumSet<Vault.StateStatus>): Iterable<StateAndRef<T>> {
         val stateAndRefs =
             session.withTransaction(TransactionIsolation.REPEATABLE_READ) {
                 var result = select(VaultSchema.VaultStates::class)
@@ -161,15 +161,16 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
                 // TODO: temporary fix to continue supporting track() function (until becomes Typed)
                 if (!clazzes.map {it.name}.contains(ContractState::class.java.name))
                     result.and (VaultSchema.VaultStates::contractStateClassName `in` (clazzes.map { it.name }))
-                result.get()
+                val iterator = result.get().iterator()
+                Sequence{iterator}
                         .map { it ->
                             val stateRef = StateRef(SecureHash.parse(it.txId), it.index)
-                            // TODO: revisit Kryo bug when using THREAD_LOCAL_KYRO
+                            // TODO: revisit Kryo bug when using THREAD_LOCAL_KRYO
                             val state = it.contractState.deserialize<TransactionState<T>>(createKryo())
                             StateAndRef(state, stateRef)
-                        }.toList()
+                        }
             }
-        return stateAndRefs
+        return stateAndRefs.asIterable()
     }
 
     override fun statesForRefs(refs: List<StateRef>): Map<StateRef, TransactionState<*>?> {

node/src/test/kotlin/net/corda/node/services/NodeVaultServiceTest.kt

@@ -102,7 +102,7 @@ class NodeVaultServiceTest {
             }
             services1.fillWithSomeTestCash(100.DOLLARS, DUMMY_NOTARY, 3, 3, Random(0L))
 
-            val w1 = services1.vaultService.unconsumedStates<Cash.State>()
+            val w1 = services1.vaultService.unconsumedStates<Cash.State>().toList()
             assertThat(w1).hasSize(3)
 
             val stateRefs = listOf(w1[1].ref, w1[2].ref)

node/src/test/kotlin/net/corda/node/services/VaultWithCashTest.kt

@@ -24,6 +24,7 @@ import net.corda.testing.MEGA_CORP
 import net.corda.testing.MEGA_CORP_KEY
 import net.corda.testing.node.MockServices
 import net.corda.testing.node.makeTestDataSourceProperties
+import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.jetbrains.exposed.sql.Database
 import org.junit.After
@@ -76,15 +77,15 @@ class VaultWithCashTest {
             // Fix the PRNG so that we get the same splits every time.
             services.fillWithSomeTestCash(100.DOLLARS, DUMMY_NOTARY, 3, 3, Random(0L))
 
-            val w = vault.unconsumedStates<Cash.State>()
+            val w = vault.unconsumedStates<Cash.State>().toList()
-            assertEquals(3, w.toList().size)
+            assertEquals(3, w.size)
 
-            val state = w.toList()[0].state.data
+            val state = w[0].state.data
             assertEquals(30.45.DOLLARS `issued by` DUMMY_CASH_ISSUER, state.amount)
             assertEquals(services.key.public.composite, state.owner)
 
-            assertEquals(34.70.DOLLARS `issued by` DUMMY_CASH_ISSUER, (w.toList()[2].state.data).amount)
+            assertEquals(34.70.DOLLARS `issued by` DUMMY_CASH_ISSUER, (w[2].state.data).amount)
-            assertEquals(34.85.DOLLARS `issued by` DUMMY_CASH_ISSUER, (w.toList()[1].state.data).amount)
+            assertEquals(34.85.DOLLARS `issued by` DUMMY_CASH_ISSUER, (w[1].state.data).amount)
         }
     }
 
@@ -164,7 +165,7 @@ class VaultWithCashTest {
             dummyIssue.toLedgerTransaction(services).verify()
 
             services.recordTransactions(dummyIssue)
-            assertEquals(1, vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>().size)
+            assertThat(vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>()).hasSize(1)
 
             // Move the same state
             val dummyMove = TransactionType.General.Builder(notary = DUMMY_NOTARY).apply {
@@ -176,7 +177,7 @@ class VaultWithCashTest {
             dummyIssue.toLedgerTransaction(services).verify()
 
             services.recordTransactions(dummyMove)
-            assertEquals(1, vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>().size)
+            assertThat(vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>()).hasSize(1)
         }
     }
 
@@ -220,19 +221,19 @@ class VaultWithCashTest {
         databaseTransaction(database) {
 
             services.fillWithSomeTestDeals(listOf("123","456","789"))
-            val deals = vault.unconsumedStates<net.corda.contracts.testing.DummyDealContract.State>()
+            val deals = vault.unconsumedStates<net.corda.contracts.testing.DummyDealContract.State>().toList()
             deals.forEach { println(it.state.data.ref) }
 
             services.fillWithSomeTestLinearStates(3)
-            val linearStates = vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>()
+            val linearStates = vault.unconsumedStates<net.corda.contracts.testing.DummyLinearContract.State>().toList()
             linearStates.forEach { println(it.state.data.linearId) }
 
             // Create a txn consuming different contract types
             val dummyMove = TransactionType.General.Builder(notary = DUMMY_NOTARY).apply {
                 addOutputState(net.corda.contracts.testing.DummyLinearContract.State(participants = listOf(freshKey.public.composite)))
                 addOutputState(net.corda.contracts.testing.DummyDealContract.State(ref = "999", participants = listOf(freshKey.public.composite)))
-                addInputState(linearStates[0])
+                addInputState(linearStates.first())
-                addInputState(deals[0])
+                addInputState(deals.first())
                 signWith(DUMMY_NOTARY_KEY)
             }.toSignedTransaction()
 

node/src/test/kotlin/net/corda/node/services/persistence/DataVendingServiceTests.kt

@@ -17,6 +17,7 @@ import net.corda.node.utilities.databaseTransaction
 import net.corda.testing.MEGA_CORP
 import net.corda.testing.node.MockNetwork
 import net.corda.testing.node.MockNetwork.MockNode
+import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Test
 import kotlin.test.assertEquals
@@ -48,14 +49,13 @@ class DataVendingServiceTests {
         ptx.signWith(registerKey)
         val tx = ptx.toSignedTransaction()
         databaseTransaction(vaultServiceNode.database) {
-            assertEquals(0, vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>().size)
+            assertThat(vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>()).isEmpty()
 
             registerNode.sendNotifyTx(tx, vaultServiceNode)
 
             // Check the transaction is in the receiving node
             val actual = vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>().singleOrNull()
             val expected = tx.tx.outRef<Cash.State>(0)
-
             assertEquals(expected, actual)
         }
     }
@@ -79,12 +79,12 @@ class DataVendingServiceTests {
         ptx.signWith(registerKey)
         val tx = ptx.toSignedTransaction(false)
         databaseTransaction(vaultServiceNode.database) {
-            assertEquals(0, vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>().size)
+            assertThat(vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>()).isEmpty()
 
             registerNode.sendNotifyTx(tx, vaultServiceNode)
 
             // Check the transaction is not in the receiving node
-            assertEquals(0, vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>().size)
+            assertThat(vaultServiceNode.services.vaultService.unconsumedStates<Cash.State>()).isEmpty()
         }
     }