Flow size- and time-based chunked

kotlinx-coroutines-core/common/src/flow/operators/Chunked.kt

@@ -0,0 +1,181 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+@file:JvmMultifileClass
+@file:JvmName("FlowKt")
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.internal.*
+import kotlinx.coroutines.selects.*
+import kotlin.jvm.*
+import kotlin.math.*
+import kotlin.time.*
+
+private const val NO_MAXIMUM = -1
+
+public fun <T> Flow<T>.chunked(maxSize: Int, minSize: Int = 1): Flow<List<T>> {
+    require(minSize in 0 until maxSize)
+    return flow {
+        val accumulator = mutableListOf<T>()
+        collect { value ->
+            accumulator.add(value)
+            if (accumulator.size == maxSize) emit(accumulator.drain())
+        }
+        if (accumulator.size >= minSize) emit(accumulator)
+    }
+}
+
+@ExperimentalTime
+public fun <T> Flow<T>.chunked(
+    chunkDuration: Duration,
+    minSize: Int = 1,
+    maxSize: Int = NO_MAXIMUM
+): Flow<List<T>> = chunked(chunkDuration.toDelayMillis(), minSize, maxSize)
+
+public fun <T> Flow<T>.chunked(
+    chunkDurationMs: Long,
+    minSize: Int = 1,
+    maxSize: Int = NO_MAXIMUM
+): Flow<List<T>> {
+    require(chunkDurationMs > 0)
+    require(minSize >= 0)
+    require(maxSize == NO_MAXIMUM || maxSize >= minSize)
+
+    return if (minSize == 0 && maxSize == NO_MAXIMUM) chunkFixedTimeWindows(chunkDurationMs)
+    else if (minSize == 0) chunkContinousWindows(chunkDurationMs, maxSize)
+    else chunkFloatingWindows(chunkDurationMs, minSize, maxSize)
+}
+
+private fun <T> Flow<T>.chunkFixedTimeWindows(durationMs: Long): Flow<List<T>> = scopedFlow { downstream ->
+    val upstream = produce(capacity = Channel.CHANNEL_DEFAULT_CAPACITY) {
+        val ticker = Ticker(durationMs, this).apply { send(Ticker.Message.Start) }
+        launch {
+            for (tick in ticker) send(Signal.TimeIsUp)
+        }
+        collect { value -> send(Signal.NewElement(value)) }
+        ticker.close()
+    }
+    val accumulator = mutableListOf<T>()
+
+    for (signal in upstream) {
+        when (signal) {
+            is Signal.NewElement -> accumulator.add(signal.value)
+            is Signal.TimeIsUp -> downstream.emit(accumulator.drain())
+        }
+    }
+    if (accumulator.isNotEmpty()) downstream.emit(accumulator)
+}
+
+private fun <T> Flow<T>.chunkContinousWindows(durationMs: Long, maxSize: Int): Flow<List<T>> =
+    scopedFlow { downstream ->
+        val inbox: ReceiveChannel<T> = this@chunkContinousWindows.produceIn(this)
+        val ticker = Ticker(durationMs, this).apply { send(Ticker.Message.Start) }
+        val accumulator = mutableListOf<T>()
+
+        whileSelect {
+            inbox.onReceiveOrClosed.invoke { valueOrClosed ->
+                val isOpen = !valueOrClosed.isClosed
+                if (isOpen) {
+                    accumulator.add(valueOrClosed.value)
+                    if(accumulator.size == maxSize){
+                        ticker.send(Ticker.Message.Reset)
+                        downstream.emit(accumulator.drain())
+                        ticker.send(Ticker.Message.Start)
+                    }
+                }
+                isOpen
+            }
+            ticker.onReceive.invoke {
+                downstream.emit(accumulator.drain())
+                true
+            }
+        }
+
+        ticker.close()
+        if (accumulator.isNotEmpty()) downstream.emit(accumulator)
+    }
+
+private fun <T> Flow<T>.chunkFloatingWindows(
+    durationMs: Long,
+    minSize: Int,
+    maxSize: Int,
+): Flow<List<T>> {
+
+    return scopedFlow { downstream ->
+        val upstream: ReceiveChannel<T> = this@chunkFloatingWindows.produceIn(this)
+        val ticker = Ticker(durationMs, this)
+        val accumulator = mutableListOf<T>()
+
+        whileSelect {
+            upstream.onReceiveOrClosed.invoke { valueOrClosed ->
+                val isOpen = valueOrClosed.isClosed.not()
+                if (isOpen) {
+                    if (accumulator.isEmpty()) ticker.send(Ticker.Message.Start)
+                    accumulator.add(valueOrClosed.value)
+                    if (accumulator.size == maxSize) {
+                        ticker.send(Ticker.Message.Reset)
+                        downstream.emit(accumulator.drain())
+                    }
+                }
+                isOpen
+            }
+            ticker.onReceive.invoke {
+                if (accumulator.size >= minSize) downstream.emit(accumulator.drain())
+                true
+            }
+        }
+
+        ticker.close()
+        if (accumulator.size >= minSize) downstream.emit(accumulator)
+    }
+}
+
+private class Ticker(
+    private val intervalMs: Long,
+    scope: CoroutineScope,
+    private val inbox: Channel<Message> = Channel(),
+    private val ticks: Channel<Unit> = Channel()
+) : SendChannel<Ticker.Message> by inbox, ReceiveChannel<Unit> by ticks {
+
+    init {
+        scope.processMessages()
+    }
+
+    private fun CoroutineScope.processMessages() = launch {
+        var ticker = setupTicks()
+        for (message in inbox) {
+            when (message) {
+                Message.Start -> ticker.start()
+                Message.Reset -> {
+                    ticker.cancel()
+                    ticker = setupTicks()
+                }
+            }
+        }
+        ticker.cancel()
+        ticks.cancel()
+    }
+
+    private fun CoroutineScope.setupTicks() = launch(start = CoroutineStart.LAZY) {
+        while (true) {
+            delay(intervalMs)
+            ticks.send(Unit)
+        }
+    }
+
+    sealed class Message {
+        object Start : Message()
+        object Reset : Message()
+    }
+}
+
+private sealed class Signal<out T> {
+    class NewElement<out T>(val value: T) : Signal<T>()
+    object TimeIsUp : Signal<Nothing>()
+}
+
+private fun <T> MutableList<T>.drain() = toList().also { this.clear() }

kotlinx-coroutines-core/common/test/flow/operators/ChunkedTest.kt

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.flow
+
+import kotlinx.coroutines.*
+import kotlin.test.*
+import kotlin.time.*
+
+@ExperimentalTime
+class ChunkedTest : TestBase() {
+
+    @Test
+    fun testEmptyFlowChunking() = runTest {
+        val emptyFlow = emptyFlow<Int>()
+        val result = measureTimedValue {
+            emptyFlow.chunked(10.seconds).toList()
+        }
+
+        assertTrue { result.value.isEmpty() }
+        assertTrue { result.duration.inSeconds < 1 }
+    }
+
+    @ExperimentalTime
+    @Test
+    fun testSingleFastElementChunking() = runTest {
+        val fastFlow = flow { emit(1) }
+
+        val result = measureTimedValue {
+            fastFlow.chunked(10.seconds).toList()
+        }
+
+        assertTrue { result.value.size == 1 && result.value.first().contains(1) }
+        assertTrue { result.duration.inSeconds < 1 }
+    }
+
+    @ExperimentalTime
+    @Test
+    fun testMultipleFastElementsChunking() = runTest {
+        val fastFlow = flow {
+            for(i in 1..1000) emit(1)
+        }
+
+        val result = measureTimedValue {
+            fastFlow.chunked(10.seconds).toList()
+        }
+
+        assertTrue { result.value.size == 1 && result.value.first().size == 1000 }
+        assertTrue { result.duration.inSeconds < 1 }
+    }
+
+    @Test
+    fun testFixedTimeWindowChunkingWithZeroMinimumSize() = withVirtualTime {
+        val intervalFlow = flow {
+            delay(1500)
+            emit(1)
+            delay(1500)
+            emit(2)
+            delay(1500)
+            emit(3)
+        }
+        val chunks = intervalFlow.chunked(2.seconds, minSize = 0).toList()
+
+        assertEquals (3, chunks.size)
+        assertTrue { chunks.all { it.size == 1 } }
+
+        finish(1)
+    }
+
+    @Test
+    fun testDefaultChunkingWithFloatingWindows() = withVirtualTime {
+        val intervalFlow = flow {
+            delay(1500)
+            emit(1)
+            delay(1500)
+            emit(2)
+            delay(1500)
+            emit(3)
+        }
+        val chunks = intervalFlow.chunked(2.seconds).toList()
+
+        assertEquals (2, chunks.size)
+        assertTrue { chunks.first().size == 2 && chunks[1].size == 1 }
+
+        finish(1)
+    }
+
+    @Test
+    fun testRespectingMinValue() = withVirtualTime {
+        val intervalFlow = flow {
+            delay(1500)
+            emit(1)
+            delay(1500)
+            emit(2)
+            delay(1500)
+            emit(3)
+        }
+        val chunks = intervalFlow.chunked(2.seconds, minSize = 3).toList()
+
+        assertTrue { chunks.size == 1 }
+        assertTrue { chunks.first().size == 3 }
+
+        finish(1)
+    }
+
+    @Test
+    fun testRespectingMaxValueWithContinousWindows() = withVirtualTime {
+        val intervalFlow = flow {
+            delay(1500)
+            emit(1)
+            emit(2)
+            emit(3)
+            emit(4)
+            delay(1500)
+            emit(5)
+            delay(1500)
+            emit(6)
+        }
+        val chunks = intervalFlow.chunked(2.seconds, minSize = 0, maxSize = 3).toList()
+
+        assertEquals(3, chunks.size)
+        assertEquals(3, chunks.first().size)
+        assertEquals(2, chunks[1].size)
+        assertTrue { chunks[1].containsAll(listOf(4, 5)) }
+
+        finish(1)
+    }
+
+    @Test
+    fun testRespectingMaxValueAndResetingTickerWithNonContinousWindows() = withVirtualTime {
+        val intervalFlow = flow {
+            delay(1500)
+            emit(1)
+            emit(2)
+            emit(3)
+            delay(1500)
+            emit(4)
+            emit(5)
+            delay(1500)
+            emit(6)
+        }
+        val chunks = intervalFlow.chunked(2.seconds, maxSize = 3).toList()
+
+        assertEquals(2, chunks.size)
+        assertEquals(3, chunks.first().size)
+        assertEquals(3, chunks[1].size)
+        assertTrue { chunks[1].containsAll(listOf(4, 5, 6)) }
+
+        finish(1)
+    }
+
+    @Test
+    fun testSizeBasedChunking() = runTest {
+        val flow = flow {
+            for (i in 1..10) emit(i)
+        }
+
+        val chunks = flow.chunked(maxSize = 3).toList()
+
+        assertEquals(4, chunks.size)
+    }
+
+    @Test
+    fun testSizeBasedChunkingWithMinSize() = runTest {
+        val flow = flow {
+            for (i in 1..10) emit(i)
+        }
+
+        val chunks = flow.chunked(maxSize = 3, minSize = 2).toList()
+
+        assertEquals(3, chunks.size)
+    }
+
+}