fix: wait for all flow-graph connections before returning from FlowGraph()

feat(test): improve tests, add one for CloningFlow

Author: silenium-dev

src/main/kotlin/dev/silenium/libs/flows/api/Extensions.kt

@@ -1,10 +1,9 @@
 package dev.silenium.libs.flows.api
 
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
 
 interface FlowGraphBuilder : FlowGraph {
-    infix fun <T, P> Source<T, P>.connectTo(sink: Sink<T, P>): Job =
-        flow.onEach(sink::submit).launchIn(this@FlowGraphBuilder)
+    infix fun <T, P> Source<T, P>.connectTo(sink: Sink<T, P>): Result<Job>
+
+    suspend fun finalize(): Result<Unit>
 }

src/main/kotlin/dev/silenium/libs/flows/impl/FlowGraphImpl.kt

@@ -1,9 +1,7 @@
 package dev.silenium.libs.flows.impl
 
 import dev.silenium.libs.flows.api.*
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.cancel
+import kotlinx.coroutines.*
 import kotlin.coroutines.CoroutineContext
 import kotlin.reflect.KClass
 
@@ -98,16 +96,37 @@ internal class FlowGraphImpl(private val coroutineScope: CoroutineScope) :
     ) : CoroutineContext.Element
 }
 
-internal class FlowGraphBuilderImpl(private val flowGraph: FlowGraph) : FlowGraphBuilder, FlowGraph by flowGraph
+internal class FlowGraphBuilderImpl(private val flowGraph: FlowGraph) : FlowGraphBuilder, FlowGraph by flowGraph {
+    private val connectionStarted = mutableSetOf<Job>()
+
+    override fun <T, P> Source<T, P>.connectTo(sink: Sink<T, P>): Result<Job> {
+        outputMetadata.forEach { (pad, metadata) ->
+            sink.configure(pad, metadata).onFailure {
+                return Result.failure(IllegalStateException("Unable to configure input pad $pad of sink $sink", it))
+            }
+        }
+        val started = CompletableDeferred<Unit>()
+        return launch {
+            started.complete(Unit)
+            flow.collect(sink)
+        }.also {
+            connectionStarted.add(started)
+        }.let { Result.success(it) }
+    }
+
+    override suspend fun finalize(): Result<Unit> = runCatching {
+        connectionStarted.joinAll()
+    }
+}
 
 internal fun FlowGraph.builder() = FlowGraphBuilderImpl(this)
 
-fun FlowGraph(
+suspend fun FlowGraph(
     coroutineContext: CoroutineContext = Dispatchers.Default,
     block: FlowGraphBuilder.() -> Unit,
-): FlowGraph = FlowGraphImpl(coroutineContext).builder().apply(block)
+): FlowGraph = FlowGraphImpl(coroutineContext).builder().apply(block).apply { finalize() }
 
-fun FlowGraph(
+suspend fun FlowGraph(
     coroutineScope: CoroutineScope,
     block: FlowGraphBuilder.() -> Unit,
-): FlowGraph = FlowGraphImpl(coroutineScope).builder().apply(block)
+): FlowGraph = FlowGraphImpl(coroutineScope).builder().apply(block).apply { finalize() }

src/test/kotlin/dev/silenium/libs/flows/base/SourceBaseTest.kt

@@ -30,9 +30,12 @@ class SourceBaseTest : FunSpec({
         val job = CoroutineScope(Dispatchers.Default).launch {
             bufferSource.flow.collect(decoder)
         }
+        val started = CompletableDeferred<Unit>()
         val listAsync = async(Dispatchers.Default) {
+            started.complete(Unit)
             decoder.flow.toList()
         }
+        started.await()
         val bufs = inputs.map { input ->
             val base64 = Base64.encode(input.encodeToByteArray())
             val byteBuffer = ByteBuffer.allocateDirect(base64.length)

src/test/kotlin/dev/silenium/libs/flows/impl/CloningFlowTest.kt

@@ -0,0 +1,69 @@
+package dev.silenium.libs.flows.impl
+
+import dev.silenium.libs.flows.api.ReferenceCounted
+import io.kotest.core.spec.style.FunSpec
+import io.kotest.matchers.nulls.shouldBeNull
+import io.kotest.matchers.nulls.shouldNotBeNull
+import io.kotest.matchers.shouldBe
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.async
+import kotlinx.coroutines.flow.toList
+import java.util.concurrent.atomic.AtomicLong
+import java.util.concurrent.atomic.AtomicReference
+
+class TestData(
+    val buf: AtomicReference<ByteArray?>,
+    private val refCount_: AtomicLong = AtomicLong(0L)
+) : ReferenceCounted<TestData> {
+    constructor(data: ByteArray) : this(AtomicReference(data))
+
+    init {
+        refCount_.incrementAndGet()
+    }
+
+    val refCount: Long
+        get() = refCount_.get()
+
+    override fun clone(): Result<TestData> {
+        return Result.success(TestData(buf, refCount_))
+    }
+
+    override fun close() {
+        if (refCount_.decrementAndGet() == 0L) {
+            buf.set(null)
+        }
+    }
+}
+
+class CloningFlowTest : FunSpec({
+    test("cloning flow clones properly") {
+        val flow = CloningFlow<TestData>()
+
+        val started = CompletableDeferred<Unit>()
+        val items = async {
+            started.complete(Unit)
+            flow.toList()
+        }
+        started.await()
+
+        val inputs = listOf(
+            "Hello, World!",
+            "some text",
+            "Another text",
+        )
+        inputs.forEach { input ->
+            val data = TestData(input.encodeToByteArray())
+            flow.publish(data)
+            data.close()
+        }
+
+        flow.close()
+        items.await().map { item ->
+            item.buf.get().shouldNotBeNull().decodeToString().also {
+                item.close()
+                item.refCount shouldBe 0L
+                item.buf.get().shouldBeNull()
+            }
+        }.toSet() shouldBe inputs.toSet()
+    }
+})

src/test/kotlin/dev/silenium/libs/flows/impl/FlowGraphImplTest.kt

@@ -13,7 +13,7 @@ class FlowGraphImplTest : FunSpec({
     test("FlowGraphBuilder") {
         val graph = FlowGraph(CoroutineScope(Dispatchers.Default)) {
             val source = source(BufferSource<Base64Buffer, DataType>(0u to DataType.BASE64), "buffer-source")
-            val sink = sink(BufferSink<ByteArray, DataType>(0u to DataType.PLAIN), "buffer-sink")
+            val sink = sink(BufferSink<ByteArray, DataType>(), "buffer-sink")
             val decoder = transformer(Base64Decoder(), "base64-decoder")
             source.connectTo(decoder)
             decoder.connectTo(sink)

src/test/kotlin/dev/silenium/libs/flows/test/BufferSink.kt

@@ -12,7 +12,7 @@ class BufferSink<T, P>(vararg pads: Pair<UInt, P>) : Sink<T, P> {
 
     private val buffer_: MutableMap<UInt, MutableList<FlowItem<T, P>>> = mutableMapOf()
     val buffer: Map<UInt, List<FlowItem<T, P>>> by ::buffer_
-    val flow_ = MutableStateFlow<Map<UInt, List<FlowItem<T, P>>>>(emptyMap())
+    private val flow_ = MutableStateFlow<Map<UInt, List<FlowItem<T, P>>>>(emptyMap())
     val flow: StateFlow<Map<UInt, List<FlowItem<T, P>>>> = flow_.asStateFlow()
 
     override suspend fun submit(item: FlowItem<T, P>): Result<Unit> {