Skip to content

Commit 9bb007d

Browse files
He-PinHe-Pin
committed
feat: Add support for switching scheduler
1 parent b20ec82 commit 9bb007d

15 files changed

+417
-25
lines changed

‎actor-tests/src/test/scala-jdk21-only/org/apache/pekko/dispatch/ForkJoinPoolVirtualThreadSpec.scala

+63
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* license agreements; and to You under the Apache License, version 2.0:
4+
*
5+
* https://www.apache.org/licenses/LICENSE-2.0
6+
*
7+
* This file is part of the Apache Pekko project, which was derived from Akka.
8+
*/
9+
10+
/*
11+
* Copyright (C) 2018-2022 Lightbend Inc. <https://www.lightbend.com>
12+
*/
13+
14+
package org.apache.pekko.dispatch
15+
16+
import com.typesafe.config.ConfigFactory
17+
18+
import org.apache.pekko
19+
import pekko.actor.{ Actor, Props }
20+
import pekko.testkit.{ ImplicitSender, PekkoSpec }
21+
import pekko.util.JavaVersion
22+
23+
object ForkJoinPoolVirtualThreadSpec {
24+
val config = ConfigFactory.parseString("""
25+
|virtual {
26+
| task-dispatcher {
27+
| mailbox-type = "org.apache.pekko.dispatch.SingleConsumerOnlyUnboundedMailbox"
28+
| throughput = 5
29+
| fork-join-executor {
30+
| parallelism-factor = 2
31+
| parallelism-max = 2
32+
| parallelism-min = 2
33+
| virtualize = on
34+
| }
35+
| }
36+
|}
37+
""".stripMargin)
38+
39+
class ThreadNameActor extends Actor {
40+
41+
override def receive = {
42+
case "ping" =>
43+
sender() ! Thread.currentThread().getName
44+
}
45+
}
46+
47+
}
48+
49+
class ForkJoinPoolVirtualThreadSpec extends PekkoSpec(ForkJoinPoolVirtualThreadSpec.config) with ImplicitSender {
50+
import ForkJoinPoolVirtualThreadSpec._
51+
52+
"PekkoForkJoinPool" must {
53+
54+
"support virtualization with Virtual Thread" in {
55+
val actor = system.actorOf(Props(new ThreadNameActor).withDispatcher("virtual.task-dispatcher"))
56+
for (_ <- 1 to 1000) {
57+
actor ! "ping"
58+
expectMsgPF() { case name: String => name should include("virtual.task-dispatcher-virtual-thread") }
59+
}
60+
}
61+
62+
}
63+
}

‎actor-tests/src/test/scala-jdk21-only/org/apache/pekko/dispatch/ThreadPoolVirtualThreadSpec.scala

+64
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.pekko.dispatch
19+
20+
import com.typesafe.config.ConfigFactory
21+
import org.apache.pekko
22+
import pekko.actor.{ Actor, Props }
23+
import pekko.testkit.{ ImplicitSender, PekkoSpec }
24+
25+
object ThreadPoolVirtualThreadSpec {
26+
val config = ConfigFactory.parseString("""
27+
|pekko.actor.default-dispatcher.executor = "thread-pool-executor"
28+
|virtual {
29+
| task-dispatcher {
30+
| mailbox-type = "org.apache.pekko.dispatch.SingleConsumerOnlyUnboundedMailbox"
31+
| throughput = 1
32+
| thread-pool-executor {
33+
| fixed-pool-size = 2
34+
| virtualize = on
35+
| }
36+
| }
37+
|}
38+
""".stripMargin)
39+
40+
class ThreadNameActor extends Actor {
41+
42+
override def receive = {
43+
case "ping" =>
44+
sender() ! Thread.currentThread().getName
45+
}
46+
}
47+
48+
}
49+
50+
class ThreadPoolVirtualThreadSpec extends PekkoSpec(ThreadPoolVirtualThreadSpec.config) with ImplicitSender {
51+
import ThreadPoolVirtualThreadSpec._
52+
53+
"PekkoThreadPoolExecutor" must {
54+
55+
"support virtualization with Virtual Thread" in {
56+
val actor = system.actorOf(Props(new ThreadNameActor).withDispatcher("virtual.task-dispatcher"))
57+
for (_ <- 1 to 1000) {
58+
actor ! "ping"
59+
expectMsgPF() { case name: String => name should include("virtual.task-dispatcher-virtual-thread") }
60+
}
61+
}
62+
63+
}
64+
}

‎actor-tests/src/test/scala-jdk21-only/org/apache/pekko/dispatch/VirtualThreadPoolDispatcherSpec.scala

-2
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,6 @@ object VirtualThreadPoolDispatcherSpec {
4343
class VirtualThreadPoolDispatcherSpec extends PekkoSpec(VirtualThreadPoolDispatcherSpec.config) with ImplicitSender {
4444
import VirtualThreadPoolDispatcherSpec._
4545

46-
val Iterations = 1000
47-
4846
"VirtualThreadPool support" must {
4947

5048
"handle simple dispatch" in {

‎actor/src/main/resources/reference.conf

+12
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,12 @@ pekko {
487487
# This config is new in Pekko v1.1.0 and only has an effect if you are running with JDK 9 and above.
488488
# Read the documentation on `java.util.concurrent.ForkJoinPool` to find out more. Default in hex is 0x7fff.
489489
maximum-pool-size = 32767
490+
491+
# This config is new in Pekko v1.2.0 and only has an effect if you are running with JDK 21 and above,
492+
# When set to `on` but underlying runtime does not support virtual threads, an Exception will throw.
493+
# Virtualize this dispatcher as a virtual-thread-executor
494+
# Valid values are: `on`, `off`
495+
virtualize = off
490496
}
491497

492498
# This will be used if you have set "executor = "thread-pool-executor""
@@ -538,6 +544,12 @@ pekko {
538544

539545
# Allow core threads to time out
540546
allow-core-timeout = on
547+
548+
# This config is new in Pekko v1.2.0 and only has an effect if you are running with JDK 21 and above,
549+
# When set to `on` but underlying runtime does not support virtual threads, an Exception will throw.
550+
# Virtualize this dispatcher as a virtual-thread-executor
551+
# Valid values are: `on`, `off`
552+
virtualize = off
541553
}
542554

543555
# This will be used if you have set "executor = "virtual-thread-executor"

‎actor/src/main/scala/org/apache/pekko/dispatch/AbstractDispatcher.scala

+5-1
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ class ThreadPoolExecutorConfigurator(config: Config, prerequisites: DispatcherPr
444444
@unused prerequisites: DispatcherPrerequisites): ThreadPoolConfigBuilder = {
445445
import org.apache.pekko.util.Helpers.ConfigOps
446446
val builder =
447-
ThreadPoolConfigBuilder(ThreadPoolConfig())
447+
ThreadPoolConfigBuilder(ThreadPoolConfig(virtualize = false))
448448
.setKeepAliveTime(config.getMillisDuration("keep-alive-time"))
449449
.setAllowCoreThreadTimeout(config.getBoolean("allow-core-timeout"))
450450
.configure(Some(config.getInt("task-queue-size")).flatMap {
@@ -474,6 +474,10 @@ class ThreadPoolExecutorConfigurator(config: Config, prerequisites: DispatcherPr
474474
config.getInt("max-pool-size-max"))
475475
else
476476
builder.setFixedPoolSize(config.getInt("fixed-pool-size"))
477+
478+
if (config.getBoolean("virtualize")) {
479+
builder.setVirtualize(true)
480+
} else builder
477481
}
478482

479483
def createExecutorServiceFactory(id: String, threadFactory: ThreadFactory): ExecutorServiceFactory =

‎actor/src/main/scala/org/apache/pekko/dispatch/Dispatchers.scala

+1-1
Original file line numberDiff line numberDiff line change
@@ -390,7 +390,7 @@ class PinnedDispatcherConfigurator(config: Config, prerequisites: DispatcherPrer
390390
this.getClass,
391391
"PinnedDispatcher [%s] not configured to use ThreadPoolExecutor, falling back to default config.".format(
392392
config.getString("id"))))
393-
ThreadPoolConfig()
393+
ThreadPoolConfig(virtualize = false)
394394
}
395395

396396
/**

‎actor/src/main/scala/org/apache/pekko/dispatch/ForkJoinExecutorConfigurator.scala

+31-9
Original file line numberDiff line numberDiff line change
@@ -86,15 +86,28 @@ class ForkJoinExecutorConfigurator(config: Config, prerequisites: DispatcherPrer
8686
}
8787

8888
class ForkJoinExecutorServiceFactory(
89+
val id: String,
8990
val threadFactory: ForkJoinPool.ForkJoinWorkerThreadFactory,
9091
val parallelism: Int,
9192
val asyncMode: Boolean,
92-
val maxPoolSize: Int)
93+
val maxPoolSize: Int,
94+
val virtualize: Boolean)
9395
extends ExecutorServiceFactory {
96+
def this(threadFactory: ForkJoinPool.ForkJoinWorkerThreadFactory,
97+
parallelism: Int,
98+
asyncMode: Boolean,
99+
maxPoolSize: Int,
100+
virtualize: Boolean) =
101+
this(null, threadFactory, parallelism, asyncMode, maxPoolSize, virtualize)
102+
103+
def this(threadFactory: ForkJoinPool.ForkJoinWorkerThreadFactory,
104+
parallelism: Int,
105+
asyncMode: Boolean) = this(threadFactory, parallelism, asyncMode, ForkJoinPoolConstants.MaxCap, false)
94106

95107
def this(threadFactory: ForkJoinPool.ForkJoinWorkerThreadFactory,
96108
parallelism: Int,
97-
asyncMode: Boolean) = this(threadFactory, parallelism, asyncMode, ForkJoinPoolConstants.MaxCap)
109+
asyncMode: Boolean,
110+
maxPoolSize: Int) = this(threadFactory, parallelism, asyncMode, maxPoolSize, false)
98111

99112
private def pekkoJdk9ForkJoinPoolClassOpt: Option[Class[_]] =
100113
Try(Class.forName("org.apache.pekko.dispatch.PekkoJdk9ForkJoinPool")).toOption
@@ -116,12 +129,19 @@ class ForkJoinExecutorConfigurator(config: Config, prerequisites: DispatcherPrer
116129
def this(threadFactory: ForkJoinPool.ForkJoinWorkerThreadFactory, parallelism: Int) =
117130
this(threadFactory, parallelism, asyncMode = true)
118131

119-
def createExecutorService: ExecutorService = pekkoJdk9ForkJoinPoolHandleOpt match {
120-
case Some(handle) =>
121-
handle.invoke(parallelism, threadFactory, maxPoolSize,
122-
MonitorableThreadFactory.doNothing, asyncMode).asInstanceOf[ExecutorService]
123-
case _ =>
124-
new PekkoForkJoinPool(parallelism, threadFactory, MonitorableThreadFactory.doNothing, asyncMode)
132+
def createExecutorService: ExecutorService = {
133+
val forkJoinPool = pekkoJdk9ForkJoinPoolHandleOpt match {
134+
case Some(handle) =>
135+
handle.invoke(parallelism, threadFactory, maxPoolSize,
136+
MonitorableThreadFactory.doNothing, asyncMode).asInstanceOf[ExecutorService]
137+
case _ =>
138+
new PekkoForkJoinPool(parallelism, threadFactory, MonitorableThreadFactory.doNothing, asyncMode)
139+
}
140+
if (virtualize) {
141+
new VirtualizedExecutorService(id, forkJoinPool)
142+
} else {
143+
forkJoinPool
144+
}
125145
}
126146
}
127147

@@ -143,12 +163,14 @@ class ForkJoinExecutorConfigurator(config: Config, prerequisites: DispatcherPrer
143163
}
144164

145165
new ForkJoinExecutorServiceFactory(
166+
id,
146167
validate(tf),
147168
ThreadPoolConfig.scaledPoolSize(
148169
config.getInt("parallelism-min"),
149170
config.getDouble("parallelism-factor"),
150171
config.getInt("parallelism-max")),
151172
asyncMode,
152-
config.getInt("maximum-pool-size"))
173+
config.getInt("maximum-pool-size"),
174+
config.getBoolean("virtualize"))
153175
}
154176
}

‎actor/src/main/scala/org/apache/pekko/dispatch/ThreadPoolBuilder.scala

+62-8
Original file line numberDiff line numberDiff line change
@@ -58,12 +58,32 @@ object ThreadPoolConfig {
5858
def reusableQueue(queue: BlockingQueue[Runnable]): QueueFactory = () => queue
5959

6060
def reusableQueue(queueFactory: QueueFactory): QueueFactory = reusableQueue(queueFactory())
61+
62+
// TODO remove in Pekko 2.0 after change to normal class
63+
def unapply(config: ThreadPoolConfig)
64+
: Option[(Boolean, Int, Int, Duration, ThreadPoolConfig.QueueFactory, RejectedExecutionHandler)] =
65+
Some((config.allowCorePoolTimeout, config.corePoolSize, config.maxPoolSize, config.threadTimeout,
66+
config.queueFactory, config.rejectionPolicy))
67+
68+
// TODO remove in Pekko 2.0 after change to normal class
69+
def apply(allowCorePoolTimeout: Boolean,
70+
corePoolSize: Int,
71+
maxPoolSize: Int,
72+
threadTimeout: Duration,
73+
queueFactory: ThreadPoolConfig.QueueFactory,
74+
rejectionPolicy: RejectedExecutionHandler): ThreadPoolConfig =
75+
ThreadPoolConfig(allowCorePoolTimeout, corePoolSize, maxPoolSize, threadTimeout, queueFactory, rejectionPolicy,
76+
virtualize = false)
6177
}
6278

6379
/**
6480
* Function0 without the fun stuff (mostly for the sake of the Java API side of things)
6581
*/
6682
trait ExecutorServiceFactory {
83+
84+
/**
85+
* Create a new ExecutorService
86+
*/
6787
def createExecutorService: ExecutorService
6888
}
6989

@@ -77,14 +97,26 @@ trait ExecutorServiceFactoryProvider {
7797
/**
7898
* A small configuration DSL to create ThreadPoolExecutors that can be provided as an ExecutorServiceFactoryProvider to Dispatcher
7999
*/
100+
//TODO don't use case class for this in 2.0, it's not a good fit
80101
final case class ThreadPoolConfig(
81102
allowCorePoolTimeout: Boolean = ThreadPoolConfig.defaultAllowCoreThreadTimeout,
82103
corePoolSize: Int = ThreadPoolConfig.defaultCorePoolSize,
83104
maxPoolSize: Int = ThreadPoolConfig.defaultMaxPoolSize,
84105
threadTimeout: Duration = ThreadPoolConfig.defaultTimeout,
85106
queueFactory: ThreadPoolConfig.QueueFactory = ThreadPoolConfig.linkedBlockingQueue(),
86-
rejectionPolicy: RejectedExecutionHandler = ThreadPoolConfig.defaultRejectionPolicy)
107+
rejectionPolicy: RejectedExecutionHandler = ThreadPoolConfig.defaultRejectionPolicy,
108+
virtualize: Boolean)
87109
extends ExecutorServiceFactoryProvider {
110+
// TODO remove in Pekko 2.0 after change to normal class
111+
def this(allowCorePoolTimeout: Boolean,
112+
corePoolSize: Int,
113+
maxPoolSize: Int,
114+
threadTimeout: Duration,
115+
queueFactory: ThreadPoolConfig.QueueFactory,
116+
rejectionPolicy: RejectedExecutionHandler) =
117+
this(allowCorePoolTimeout, corePoolSize, maxPoolSize, threadTimeout, queueFactory, rejectionPolicy,
118+
virtualize = false)
119+
88120
// Written explicitly to permit non-inlined defn; this is necessary for downstream instrumentation that stores extra
89121
// context information on the config
90122
@noinline
@@ -94,13 +126,31 @@ final case class ThreadPoolConfig(
94126
maxPoolSize: Int = maxPoolSize,
95127
threadTimeout: Duration = threadTimeout,
96128
queueFactory: ThreadPoolConfig.QueueFactory = queueFactory,
97-
rejectionPolicy: RejectedExecutionHandler = rejectionPolicy
129+
rejectionPolicy: RejectedExecutionHandler = rejectionPolicy,
130+
virtualize: Boolean = virtualize
131+
): ThreadPoolConfig =
132+
ThreadPoolConfig(allowCorePoolTimeout, corePoolSize, maxPoolSize, threadTimeout, queueFactory, rejectionPolicy,
133+
virtualize)
134+
135+
// TODO remove in Pekko 2.0 after change to normal class
136+
@noinline
137+
def copy(allowCorePoolTimeout: Boolean,
138+
corePoolSize: Int,
139+
maxPoolSize: Int,
140+
threadTimeout: Duration,
141+
queueFactory: ThreadPoolConfig.QueueFactory,
142+
rejectionPolicy: RejectedExecutionHandler
98143
): ThreadPoolConfig =
99-
ThreadPoolConfig(allowCorePoolTimeout, corePoolSize, maxPoolSize, threadTimeout, queueFactory, rejectionPolicy)
144+
ThreadPoolConfig(allowCorePoolTimeout, corePoolSize, maxPoolSize, threadTimeout, queueFactory, rejectionPolicy,
145+
virtualize)
146+
147+
class ThreadPoolExecutorServiceFactory(val id: String,
148+
val threadFactory: ThreadFactory) extends ExecutorServiceFactory {
149+
150+
def this(threadFactory: ThreadFactory) = this(null, threadFactory)
100151

101-
class ThreadPoolExecutorServiceFactory(val threadFactory: ThreadFactory) extends ExecutorServiceFactory {
102152
def createExecutorService: ExecutorService = {
103-
val service: ThreadPoolExecutor = new ThreadPoolExecutor(
153+
val executor: ThreadPoolExecutor = new ThreadPoolExecutor(
104154
corePoolSize,
105155
maxPoolSize,
106156
threadTimeout.length,
@@ -110,18 +160,19 @@ final case class ThreadPoolConfig(
110160
rejectionPolicy) with LoadMetrics {
111161
def atFullThrottle(): Boolean = this.getActiveCount >= this.getPoolSize
112162
}
113-
service.allowCoreThreadTimeOut(allowCorePoolTimeout)
114-
service
163+
executor.allowCoreThreadTimeOut(allowCorePoolTimeout)
164+
if (virtualize) new VirtualizedExecutorService(id, executor) else executor
115165
}
116166
}
167+
117168
def createExecutorServiceFactory(id: String, threadFactory: ThreadFactory): ExecutorServiceFactory = {
118169
val tf = threadFactory match {
119170
case m: MonitorableThreadFactory =>
120171
// add the dispatcher id to the thread names
121172
m.withName(m.name + "-" + id)
122173
case other => other
123174
}
124-
new ThreadPoolExecutorServiceFactory(tf)
175+
new ThreadPoolExecutorServiceFactory(id, tf)
125176
}
126177
}
127178

@@ -178,6 +229,9 @@ final case class ThreadPoolConfigBuilder(config: ThreadPoolConfig) {
178229
def setQueueFactory(newQueueFactory: QueueFactory): ThreadPoolConfigBuilder =
179230
this.copy(config = config.copy(queueFactory = newQueueFactory))
180231

232+
def setVirtualize(virtualize: Boolean): ThreadPoolConfigBuilder =
233+
this.copy(config = config.copy(virtualize = virtualize))
234+
181235
def configure(fs: Option[Function[ThreadPoolConfigBuilder, ThreadPoolConfigBuilder]]*): ThreadPoolConfigBuilder =
182236
fs.foldLeft(this)((c, f) => f.map(_(c)).getOrElse(c))
183237
}

0 commit comments

Comments
 (0)