Repository: Kotlin/kotlinx.coroutines
File: kotlinx-coroutines-core/jvm/src/scheduling/WorkQueue.kt (Kotlin)

Code before:
            if (!lastScheduled.isBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        return pollWithMode(onlyBlocking = true /* only blocking */)
    }

    fun pollCpu(): Task? {
        while (true) { // Poll the slot
            val lastScheduled = lastScheduledTask.value ?: break
            if (lastScheduled.isBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        return pollWithMode(onlyBlocking = false /* only cpu */)
    }

    private fun pollWithMode(/* Only blocking OR only CPU */ onlyBlocking: Boolean): Task? {
        val start = consumerIndex.value
        var end = producerIndex.value
        // CPU or (BLOCKING & hasBlocking)
        val shouldProceed = !onlyBlocking || blockingTasksInBuffer.value > 0
        while (start != end && shouldProceed) {
            val task = tryExtractFromTheMiddle(--end, onlyBlocking)
            if (task != null) {
                return task
            }
        }
        return null
    }

    private fun tryExtractFromTheMiddle(index: Int, onlyBlocking: Boolean): Task? {
        if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
        val arrayIndex = index and MASK
        val value = buffer[arrayIndex]
        if (value != null && value.isBlocking == onlyBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
            if (onlyBlocking) blockingTasksInBuffer.decrementAndGet()
            return value
        }
        return null
    }

    fun offloadAllWorkTo(globalQueue: GlobalQueue) {
        lastScheduledTask.getAndSet(null)?.let { globalQueue.addLast(it) }
        while (pollTo(globalQueue)) {
            // Steal everything
        }

Review comment (suggestion):
`shouldProceed` isn't updated deliberately, I find making it as part of `while` a bit more compact 

Code after:
    // NB: ONLY for runSingleTask method
    fun pollBlocking(): Task? = pollWithExclusiveMode(onlyBlocking = true /* only blocking */)

    // Polls for CPU task, invoked only by the owner
    // NB: ONLY for runSingleTask method
    fun pollCpu(): Task? = pollWithExclusiveMode(onlyBlocking = false /* only cpu */)

    private fun pollWithExclusiveMode(/* Only blocking OR only CPU */ onlyBlocking: Boolean): Task? {
        while (true) { // Poll the slot
            val lastScheduled = lastScheduledTask.value ?: break
            if (lastScheduled.isBlocking != onlyBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        // Failed to poll the slot, scan the queue
        val start = consumerIndex.value
        var end = producerIndex.value
        // Bail out if there is no blocking work for us
        while (start != end) {
            if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
            val task = tryExtractFromTheMiddle(--end, onlyBlocking)
            if (task != null) {
                return task
            }
        }
        return null
    }

    private fun tryExtractFromTheMiddle(index: Int, onlyBlocking: Boolean): Task? {
        val arrayIndex = index and MASK
        val value = buffer[arrayIndex]
        if (value != null && value.isBlocking == onlyBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
            if (onlyBlocking) blockingTasksInBuffer.decrementAndGet()
            return value
        }
        return null
    }

    fun offloadAllWorkTo(globalQueue: GlobalQueue) {
        lastScheduledTask.getAndSet(null)?.let { globalQueue.addLast(it) }
        while (pollTo(globalQueue)) {
            // Steal everything
        }
    }

    /**
     * Contract on return value is the same as for [trySteal]
     */
    private fun tryStealLastScheduled(stealingMode: StealingMode, stolenTaskRef: ObjectRef<Task?>): Long {