Support IO-bound work with its own thread (#69)

Author: Ben Asher

build.gradle.kts

@@ -4,7 +4,7 @@ val coroutinesVersion = "1.3.9"
 val atomicfuVersion = "0.14.4"
 
 plugins {
-    kotlin("multiplatform") version "1.4.0"
+    kotlin("multiplatform") version "1.4.10"
     id("org.jetbrains.dokka") version "0.10.0"
     id("maven-publish")
     id("signing")

gradle.properties

@@ -1 +1 @@
-VERSION=0.6.1
+VERSION=0.6.2

src/nativeMain/kotlin/com/autodesk/coroutineworker/BackgroundCoroutineWorkQueueExecutor.kt

@@ -45,6 +45,11 @@ internal class BackgroundCoroutineWorkQueueExecutor<WorkItem : CoroutineWorkItem
      */
     private val queueThread = Worker.start()
 
+    /**
+     * Special worker for IO work
+     */
+    private val ioWorker = Worker.start(name = ioWorkerName)
+
     /**
      * The wrapped (allow freezing and mutable access on single thread) queue of WorkItems
      */
@@ -84,32 +89,52 @@ internal class BackgroundCoroutineWorkQueueExecutor<WorkItem : CoroutineWorkItem
     }.result
 
     /**
-     * Queues an item to be executed
+     * Queues an item to be executed in the general worker pool
      */
-    fun enqueueWork(item: WorkItem) = queueThread.executeAfter(
-        operation = {
-            queue.enqueue(item)
-            // start a worker if we have more workers to start
-            val activeWorkerCount = _numActiveWorkers.value
-            if (activeWorkerCount < numWorkers) {
-                pool.performWork {
+    fun enqueueWork(item: WorkItem, isIoWork: Boolean) {
+        if (isIoWork) {
+            ioWorker.executeAfter(
+                operation = {
                     runBlocking {
-                        // error if we accidentally freeze coroutine internals
                         this.ensureNeverFrozen()
-                        processWorkItems()
+                        performWorkHandlingExceptions(item, this)
                     }
-                }
-                _numActiveWorkers.increment()
-            }
-        }.freeze()
-    )
+                }.freeze()
+            )
+        } else {
+            queueThread.executeAfter(
+                operation = {
+                    queue.enqueue(item)
+                    // start a worker if we have more workers to start
+                    val activeWorkerCount = _numActiveWorkers.value
+                    if (activeWorkerCount < numWorkers) {
+                        pool.performWork {
+                            runBlocking {
+                                // error if we accidentally freeze coroutine internals
+                                this.ensureNeverFrozen()
+                                processWorkItems(this)
+                            }
+                        }
+                        _numActiveWorkers.increment()
+                    }
+                }.freeze()
+            )
+        }
+    }
 
-    suspend fun CoroutineScope.processWorkItems() {
+    private suspend fun processWorkItems(scope: CoroutineScope) {
         val workItem = dequeueWork() ?: return
 
+        performWorkHandlingExceptions(workItem, scope)
+
+        // execute a coroutine to attempt to process the next work item, if possible
+        scope.launch { processWorkItems(scope) }
+    }
+
+    private suspend fun performWorkHandlingExceptions(workItem: WorkItem, scope: CoroutineScope) {
         // Execute the work in a job that can be cancelled
         try {
-            async {
+            scope.async {
                 workItem.work(this)
             }.await()
         } catch (_: CancellationException) {
@@ -122,9 +147,6 @@ internal class BackgroundCoroutineWorkQueueExecutor<WorkItem : CoroutineWorkItem
                 throw e
             }
         }
-
-        // execute a coroutine to attempt to process the next work item, if possible
-        launch { processWorkItems() }
     }
 
     init { freeze() }
@@ -136,6 +158,16 @@ internal class BackgroundCoroutineWorkQueueExecutor<WorkItem : CoroutineWorkItem
         internal fun setUnhandledExceptionHook(handler: (Throwable) -> Unit) {
             UNHANDLED_EXCEPTION_HOOK.value = handler.freeze()
         }
+
+        /**
+         * The name of the IO worker
+         */
+        private const val ioWorkerName = "com.autodesk.coroutineworker.ioworker"
+
+        /**
+         * Returns whether we're already running on the IO thread
+         */
+        internal fun shouldPerformIoWorkInline() = Worker.current.name == ioWorkerName
     }
 }
 

src/nativeMain/kotlin/com/autodesk/coroutineworker/CoroutineUtils.kt

@@ -1,9 +1,12 @@
 package com.autodesk.coroutineworker
 
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Delay
+import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.delay
 import kotlin.coroutines.ContinuationInterceptor
+import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.coroutineContext
 import kotlin.native.concurrent.AtomicInt
 import kotlin.native.concurrent.AtomicReference
@@ -64,3 +67,14 @@ public actual suspend fun <T> threadSafeSuspendCallback(startAsync: (CompletionL
         throw e
     }
 }
+
+/**
+ * Symbolic dispatcher used to trigger IO-bound-thread-like behavior on native
+ */
+public object IODispatcher : CoroutineDispatcher() {
+    override fun dispatch(context: CoroutineContext, block: Runnable) {
+        throw UnsupportedOperationException(
+            "This dispatcher is symbolic and should not be used to run anything."
+        )
+    }
+}

src/nativeMain/kotlin/com/autodesk/coroutineworker/CoroutineWorker.kt

@@ -50,21 +50,16 @@ public actual class CoroutineWorker internal actual constructor() {
         private val executor = BackgroundCoroutineWorkQueueExecutor<WorkItem>(4)
 
         public actual fun execute(block: suspend CoroutineScope.() -> Unit): CoroutineWorker {
-            return CoroutineWorker().also {
-                val state = it.state
-                executor.enqueueWork(
-                    WorkItem(
-                        { state.cancelled },
-                        { state.completed = true },
-                        block
-                    )
-                )
-            }
+            return executeInternal(false, block)
         }
 
         public actual suspend fun <T> withContext(jvmContext: CoroutineContext, block: suspend CoroutineScope.() -> T): T {
+            val isIoWork = jvmContext == IODispatcher
+            if (isIoWork && BackgroundCoroutineWorkQueueExecutor.shouldPerformIoWorkInline()) {
+                return coroutineScope(block)
+            }
             return threadSafeSuspendCallback<T> { completion ->
-                val job = execute {
+                val job = executeInternal(isIoWork) {
                     val result = runCatching {
                         block()
                     }
@@ -74,6 +69,20 @@ public actual class CoroutineWorker internal actual constructor() {
             }
         }
 
+        private fun executeInternal(isIoWork: Boolean, block: suspend CoroutineScope.() -> Unit): CoroutineWorker {
+            return CoroutineWorker().also {
+                val state = it.state
+                executor.enqueueWork(
+                    WorkItem(
+                        { state.cancelled },
+                        { state.completed = true },
+                        block
+                    ),
+                    isIoWork
+                )
+            }
+        }
+
         /** CoroutineWorker's CoroutineWorkItem class that listens for cancellation */
         private class WorkItem(
             val cancelled: () -> Boolean,

src/nativeTest/kotlin/com/autodesk/coroutineworker/CoroutineUtilsTest.kt

@@ -1,7 +1,10 @@
 package com.autodesk.coroutineworker
 
+import kotlinx.atomicfu.atomic
+import kotlinx.coroutines.Dispatchers
 import kotlin.native.concurrent.isFrozen
 import kotlin.test.Test
+import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
 class CoroutineUtilsTest {
@@ -19,4 +22,38 @@ class CoroutineUtilsTest {
             assertTrue(called)
         }
     }
+
+    @Test
+    fun `withContext with special dispatcher uses special thread`() {
+        testRunBlocking {
+            CoroutineWorker.withContext(IODispatcher) {
+                // should return true because it's on the special thread
+                assertTrue(BackgroundCoroutineWorkQueueExecutor.shouldPerformIoWorkInline())
+            }
+        }
+    }
+
+    @Test
+    fun `withContext without special dispatcher does not use special thread`() {
+        testRunBlocking {
+            CoroutineWorker.withContext(Dispatchers.Default) {
+                // should return false because it's not on the special thread
+                assertFalse(BackgroundCoroutineWorkQueueExecutor.shouldPerformIoWorkInline())
+            }
+        }
+    }
+
+    @Test
+    fun `multiple withContext on IODispatcher works without timeout`() {
+        testRunBlocking {
+            val ran = atomic(false)
+            CoroutineWorker.withContext(IODispatcher) {
+                // should return false because it's not on the special thread
+                CoroutineWorker.withContext(IODispatcher) {
+                    ran.value = true
+                }
+            }
+            assertTrue(ran.value)
+        }
+    }
 }