如何无延迟测试Kotlin协程多通道任务调度逻辑？

我实现了一个Kotlin协程通道调度器，能按设置的worker数量分发消息，功能正常，但单元测试必须加delay才能避免阻塞并得到正确结果，这种方式太粗糙，求不依赖delay的可靠测试方案。

调度器代码：

fun <INPUT_TYPE, OUTPUT_TYPE> CoroutineScope.dispatchWorkConcurrently(
    inputChannel: ReceiveChannel<INPUT_TYPE>,
    processWork: suspend (INPUT_TYPE, SendChannel<OUTPUT_TYPE>) -> Unit,
    logger: Logger,
    nameOfWorkers: String,
    maxWorkers: Int,
    workerCounter: AtomicInteger
): ReceiveChannel<OUTPUT_TYPE> {
    val coroutineScope = this + SupervisorJob()

    suspend fun launchWorker(
        input: Channel<INPUT_TYPE>,
        output: Channel<OUTPUT_TYPE>
    ) {
        val coroutineExceptionHandler = CoroutineExceptionHandler { _, throwable ->
            workerCounter.decrementAndGet()
            logger.error("Exception in $nameOfWorkers: ${throwable.localizedMessage}")
        }
        val newCounterValue = workerCounter.incrementAndGet()
        coroutineScope.launch(coroutineExceptionHandler + CoroutineName("$nameOfWorkers $newCounterValue")) {
            for (work in input) {
                processWork(work, output)
            }
        }
    }

    val internalChannel = Channel<INPUT_TYPE>()
    val outputChannel = Channel<OUTPUT_TYPE>()

    coroutineScope.launch {
        for (message in inputChannel) {
            if (!internalChannel.trySend(message).isSuccess) {
                val allowedToLaunchNewWorker = workerCounter.get() < maxWorkers
                if (allowedToLaunchNewWorker) {
                    launchWorker(internalChannel, outputChannel)
                }

                internalChannel.send(message)
            }
        }
    }

    return outputChannel
}

原测试代码：

@Test
fun `correctly dispatch jobs with multiple workers`() = runBlocking {
    val charPool : List<Char> = ('a'..'z').toList()
    val count = AtomicInteger(0);
    suspend fun processWork(message: Char, output: SendChannel<String>) {
        delay(1)
        println("Does some work with: $message")
        count.getAndIncrement()
    }

    this.dispatchWorkConcurrently(
        inputChannel = inputMessageChannel,
        processWork = ::processWork,
        logger = LoggerFactory.getLogger(ChannelUtilTest::class.java),
        nameOfWorkers = javaClass.simpleName,
        maxWorkers = 10,
        workerCounter = workerCounter
    )

    charPool.forEach {
        inputMessageChannel.send(it)
    }
    delay(20)

    assertEquals(26, count.get())
    assertEquals(10, workerCounter.get())
}

要解决测试依赖delay的问题，核心是让测试能准确感知所有任务的完成状态，而非盲目等待。需要同时调整调度器的生命周期逻辑和测试代码的同步方式。

1. 优化调度器的生命周期管理 #

首先修改调度器代码，确保输入通道关闭后，内部通道也会关闭，worker处理完任务后能正常退出并更新计数器：

fun <INPUT_TYPE, OUTPUT_TYPE> CoroutineScope.dispatchWorkConcurrently(
    inputChannel: ReceiveChannel<INPUT_TYPE>,
    processWork: suspend (INPUT_TYPE, SendChannel<OUTPUT_TYPE>) -> Unit,
    logger: Logger,
    nameOfWorkers: String,
    maxWorkers: Int,
    workerCounter: AtomicInteger,
    maxWorkerCount: AtomicInteger? = null // 可选：用于跟踪最大活跃worker数
): ReceiveChannel<OUTPUT_TYPE> {
    val coroutineScope = this + SupervisorJob()

    suspend fun launchWorker(
        input: Channel<INPUT_TYPE>,
        output: Channel<OUTPUT_TYPE>
    ) {
        val coroutineExceptionHandler = CoroutineExceptionHandler { _, throwable ->
            workerCounter.decrementAndGet()
            logger.error("Exception in $nameOfWorkers: ${throwable.localizedMessage}")
        }
        val newCounterValue = workerCounter.incrementAndGet()
        // 记录最大活跃worker数（可选）
        maxWorkerCount?.updateAndGet { currentMax -> maxOf(currentMax, newCounterValue) }
        
        coroutineScope.launch(coroutineExceptionHandler + CoroutineName("$nameOfWorkers $newCounterValue")) {
            try {
                for (work in input) {
                    processWork(work, output)
                }
            } finally {
                // worker处理完所有任务后递减计数器
                workerCounter.decrementAndGet()
            }
        }
    }

    val internalChannel = Channel<INPUT_TYPE>()
    val outputChannel = Channel<OUTPUT_TYPE>()

    coroutineScope.launch {
        // 用consumeEach自动遍历输入通道，通道关闭后自动结束循环
        inputChannel.consumeEach { message ->
            if (!internalChannel.trySend(message).isSuccess) {
                val allowedToLaunchNewWorker = workerCounter.get() < maxWorkers
                if (allowedToLaunchNewWorker) {
                    launchWorker(internalChannel, outputChannel)
                }
                internalChannel.send(message)
            }
        }
        // 输入通道关闭后，关闭内部通道，让worker处理完剩余任务后退出
        internalChannel.close()
    }

    return outputChannel
}

2. 编写可靠的单元测试 #

使用runTest（kotlinx-coroutines-test库提供）替代runBlocking，并通过关闭通道+等待worker退出的方式同步任务状态：

@Test
fun `correctly dispatch jobs with multiple workers`() = runTest {
    val charPool: List<Char> = ('a'..'z').toList()
    val count = AtomicInteger(0)
    val inputMessageChannel = Channel<Char>(Channel.UNLIMITED)
    val workerCounter = AtomicInteger(0)
    val maxWorkerCount = AtomicInteger(0) // 用于验证最大启动的worker数

    // 模拟业务处理逻辑（去掉测试用的delay）
    suspend fun processWork(message: Char, output: SendChannel<String>) {
        println("Does some work with: $message")
        count.getAndIncrement()
    }

    // 启动调度器
    this.dispatchWorkConcurrently(
        inputChannel = inputMessageChannel,
        processWork = ::processWork,
        logger = LoggerFactory.getLogger(ChannelUtilTest::class.java),
        nameOfWorkers = javaClass.simpleName,
        maxWorkers = 10,
        workerCounter = workerCounter,
        maxWorkerCount = maxWorkerCount
    )

    // 发送所有测试任务
    charPool.forEach {
        inputMessageChannel.send(it)
    }
    // 关闭输入通道，通知调度器没有更多任务
    inputMessageChannel.close()

    // 等待所有worker处理完成（加超时避免死锁）
    withTimeout(5000) {
        while (workerCounter.get() > 0) {
            yield() // 让出协程调度权，让worker有机会执行完任务
        }
    }

    // 验证结果
    assertEquals(26, count.get())
    assertEquals(10, maxWorkerCount.get()) // 验证最多启动了10个worker
}

关键改进说明 #

• 通道生命周期同步：输入通道关闭后，调度器会关闭内部通道，worker遍历完内部通道的剩余任务后会退出，避免永久阻塞。
• worker状态跟踪：在worker协程的finally块中递减计数器，确保无论正常完成还是异常退出都能更新状态。
• 测试同步方式：用withTimeout+循环检查计数器的方式替代delay，既保证能等待任务完成，又不会无限阻塞；yield()让测试协程主动让出调度权，让worker协程有机会执行。
• 使用runTest：专门为协程测试设计的函数，能自动处理协程调度，避免测试中出现不必要的延迟或调度问题。

内容的提问来源于stack exchange，提问作者MrAndre