Chunk with interval and size only

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

@@ -12,170 +12,83 @@ 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)
- }
+public object ChunkConstraint {
+ public const val NO_MAXIMUM: Int = Int.MAX_VALUE
+ public const val NO_INTERVAL: Long = Long.MAX_VALUE
 }
 
 @ExperimentalTime
 public fun <T> Flow<T>.chunked(
- chunkDuration: Duration,
- minSize: Int = 1,
- maxSize: Int = NO_MAXIMUM
-): Flow<List<T>> = chunked(chunkDuration.toDelayMillis(), minSize, maxSize)
+ interval: Duration,
+ size: Int
+): Flow<List<T>> = chunked(interval.toLongMilliseconds(), size)
 
 public fun <T> Flow<T>.chunked(
- chunkDurationMs: Long,
- minSize: Int = 1,
- maxSize: Int = NO_MAXIMUM
+ intervalMs: Long,
+ size: Int
 ): 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>()
+ require(intervalMs >= 0)
+ require(size > 0)
 
- 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)
+ return if(intervalMs != ChunkConstraint.NO_INTERVAL) chunkedTimeBased(intervalMs, size)
+ else chunkedSizeBased(size)
 }
 
-private fun <T> Flow<T>.chunkContinousWindows(durationMs: Long, maxSize: Int): Flow<List<T>> =
- scopedFlow { downstream ->
- val inbox: ReceiveChannel<T> = [email protected](this)
- val ticker = Ticker(durationMs, this).apply { send(Ticker.Message.Start) }
- val accumulator = mutableListOf<T>()
+private fun <T> Flow<T>.chunkedTimeBased(intervalMs: Long, size: Int): Flow<List<T>> = scopedFlow { downstream ->
+ val buffer = Channel<T>(size)
+ val emitSemaphore = Channel<Unit>()
+ val collectSemaphore = Channel<Unit>()
 
- 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
+ launch {
+ collect { value ->
+ val hasCapacity = buffer.offer(value)
+ if (!hasCapacity) {
+ emitSemaphore.send(Unit)
+ collectSemaphore.receive()
+ buffer.send(value)
  }
  }
-
- ticker.close()
- if (accumulator.isNotEmpty()) downstream.emit(accumulator)
+ emitSemaphore.close()
+ buffer.close()
  }
 
-private fun <T> Flow<T>.chunkFloatingWindows(
- durationMs: Long,
- minSize: Int,
- maxSize: Int,
-): Flow<List<T>> {
-
- return scopedFlow { downstream ->
- val upstream: ReceiveChannel<T> = [email protected](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
- }
+ whileSelect {
+ emitSemaphore.onReceiveOrClosed { valueOrClosed ->
+ buffer.drain().takeIf { it.isNotEmpty() }?.let { downstream.emit(it) }
+ val shouldCollectNextChunk = valueOrClosed.isClosed.not()
+ if (shouldCollectNextChunk) collectSemaphore.send(Unit) else collectSemaphore.close()
+ shouldCollectNextChunk
  }
-
- 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()
- }
- }
+ onTimeout(intervalMs) {
+ downstream.emit(buffer.awaitFirstAndDrain())
+ true
  }
- ticker.cancel()
- ticks.cancel()
  }
+}
 
- private fun CoroutineScope.setupTicks() = launch(start = CoroutineStart.LAZY) {
- while (true) {
- delay(intervalMs)
- ticks.send(Unit)
- }
- }
+private suspend fun <T> ReceiveChannel<T>.awaitFirstAndDrain(): List<T> {
+ val first = receiveOrClosed().takeIf { it.isClosed.not() }?.value ?: return emptyList()
+ return drain(mutableListOf(first))
+}
 
- sealed class Message {
- object Start : Message()
- object Reset : Message()
+private tailrec fun <T> ReceiveChannel<T>.drain(acc: MutableList<T> = mutableListOf()): List<T> {
+ val item = poll()
+ return if (item == null) acc
+ else {
+ acc.add(item)
+ drain(acc)
  }
 }
 
-private sealed class Signal<out T> {
- class NewElement<out T>(val value: T) : Signal<T>()
- object TimeIsUp : Signal<Nothing>()
+private fun <T> Flow<T>.chunkedSizeBased(size: Int): Flow<List<T>> = flow {
+ val buffer = mutableListOf<T>()
+ collect { value ->
+ buffer.add(value)
+ if(buffer.size == size) emit(buffer.drain())
+ }
+ if(buffer.isNotEmpty()) emit(buffer)
 }
 
 private fun <T> MutableList<T>.drain() = toList().also { this.clear() }