Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 1 | package kotlinx.coroutines.experimental |
| 2 | |
| 3 | import java.util.concurrent.Executors |
| 4 | import java.util.concurrent.ScheduledExecutorService |
| 5 | import java.util.concurrent.TimeUnit |
| 6 | import java.util.concurrent.atomic.AtomicInteger |
| 7 | import kotlin.concurrent.thread |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 8 | import kotlin.coroutines.CoroutineContext |
| 9 | |
| 10 | /** |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 11 | * Creates new coroutine execution context with the a single thread and built-in [yield] and [delay] support. |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 12 | * All continuations are dispatched immediately when invoked inside the thread of this context. |
| 13 | * Resources of this pool (its thread) are reclaimed when job of this context is cancelled. |
| 14 | * The specified [name] defines the name of the new thread. |
| 15 | * An optional [parent] job may be specified upon creation. |
| 16 | */ |
| 17 | fun newSingleThreadContext(name: String, parent: Job? = null): CoroutineContext = |
| 18 | newFixedThreadPoolContext(1, name, parent) |
| 19 | |
| 20 | /** |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 21 | * Creates new coroutine execution context with the fixed-size thread-pool and built-in [yield] and [delay] support. |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 22 | * All continuations are dispatched immediately when invoked inside the threads of this context. |
| 23 | * Resources of this pool (its threads) are reclaimed when job of this context is cancelled. |
| 24 | * The specified [name] defines the names of the threads. |
| 25 | * An optional [parent] job may be specified upon creation. |
| 26 | */ |
| 27 | fun newFixedThreadPoolContext(nThreads: Int, name: String, parent: Job? = null): CoroutineContext { |
| 28 | require(nThreads >= 1) { "Expected at least one thread, but $nThreads specified" } |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 29 | val job = Job(parent) |
| 30 | return job + ThreadPoolDispatcher(nThreads, name, job) |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 31 | } |
| 32 | |
| 33 | private val thisThreadContext = ThreadLocal<ThreadPoolDispatcher>() |
| 34 | |
| 35 | private class ThreadPoolDispatcher( |
| 36 | nThreads: Int, |
| 37 | name: String, |
| 38 | val job: Job |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 39 | ) : CoroutineDispatcher(), Yield, Delay { |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 40 | val threadNo = AtomicInteger() |
| 41 | val executor: ScheduledExecutorService = Executors.newScheduledThreadPool(nThreads) { target -> |
| 42 | thread(start = false, isDaemon = true, |
| 43 | name = if (nThreads == 1) name else name + "-" + threadNo.incrementAndGet()) { |
| 44 | thisThreadContext.set(this@ThreadPoolDispatcher) |
| 45 | target.run() |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | init { |
| 50 | job.onCompletion { executor.shutdown() } |
| 51 | } |
| 52 | |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 53 | override fun isDispatchNeeded(context: CoroutineContext): Boolean = thisThreadContext.get() != this |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 54 | |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 55 | override fun dispatch(context: CoroutineContext, block: Runnable) = executor.execute(block) |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 56 | |
Roman Elizarov | d528e3e | 2017-01-23 15:40:05 +0300 | [diff] [blame] | 57 | override fun scheduleResume(continuation: CancellableContinuation<Unit>) { |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 58 | executor.scheduleResume(continuation) |
Roman Elizarov | d528e3e | 2017-01-23 15:40:05 +0300 | [diff] [blame] | 59 | } |
| 60 | |
| 61 | override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) { |
Roman Elizarov | 67891d8 | 2017-01-23 16:47:20 +0300 | [diff] [blame^] | 62 | executor.scheduleResumeAfterDelay(time, unit, continuation) |
Roman Elizarov | 3754f95 | 2017-01-18 20:47:54 +0300 | [diff] [blame] | 63 | } |
| 64 | } |