ImmediatelyScalingThreadPoolTaskExecutor.kt

/**
 * A custom ThreadPoolTaskExecutor based on this SO answer https://stackoverflow.com/a/24493856.
 *
 * The main distinction of this executor is that it will not wait for the queue to be full before scaling up the number
 * of threads in the pool. Instead, it will scale up the number of threads in the pool (up to the maximum) as soon as
 * a new task is submitted and no idle threads are available to run it. Then, once the number of threads reaches the
 * maximum, it will start to queue tasks.
 */
class ImmediatelyScalingThreadPoolTaskExecutor : ThreadPoolTaskExecutor() {

    override fun createQueue(queueCapacity: Int): BlockingQueue<Runnable> {
        return object : LinkedTransferQueue<Runnable>() {
            override fun offer(e: Runnable): Boolean {
                // tryTransfer will return false if no thread is available to run the task which will cause the thread
                // pool to either spawn a new thread or reject the task. If the task is rejected, the custom
                // RejectedExecutionHandler will put the task into the queue
                return tryTransfer(e)
            }
        }
    }

    override fun initializeExecutor(threadFactory: ThreadFactory, ignore: RejectedExecutionHandler): ExecutorService {
        val rejectedExecutionHandler = RejectedExecutionHandler { r, executor ->
            try {
                executor.queue.put(r)
            } catch (e: InterruptedException) {
                Thread.currentThread().interrupt()
            }
        }
        return super.initializeExecutor(threadFactory, rejectedExecutionHandler)
    }
}