batchSequence.kts

import java.time.ZonedDateTime

data class MessageWithExpiry(
    val messageSerialized: ByteArray,
    val expiryDate: ZonedDateTime
)
data class Batch(
    val messages: Array<ByteArray>,
    val latestMessageExpiryDate: ZonedDateTime
)

private val MAX_BATCH_SIZE = 8 * 1024 * 1024
fun Sequence<MessageWithExpiry>.batchMessages() = sequence {
    val batchMessages = mutableListOf<MessageWithExpiry>()
    var batchSize = 0
    this@batchMessages.forEach { message ->
        val messageSize = message.messageSerialized.size
        if (batchSize + messageSize > MAX_BATCH_SIZE && batchMessages.isNotEmpty()) {
            yield(batchMessages.toBatch())
            batchMessages.clear()
            batchSize = 0
        }
        batchMessages += message
        batchSize += messageSize
    }
    if (batchMessages.isNotEmpty()) {
        yield(batchMessages.toBatch())
    }
}

fun MutableList<MessageWithExpiry>.toBatch(): Batch {
    return Batch(
        map { it.messageSerialized }.toTypedArray(),
        maxBy { it.expiryDate }!!.expiryDate
    )
}