From f6547f102588c6dffae7622a59712c7f61d01957 Mon Sep 17 00:00:00 2001 From: Gianluca Graziadei Date: Sun, 20 Sep 2026 12:00:04 +0200 Subject: [PATCH] Add opt-in virtual threads for blocking I/O thread pools Introduce storm.virtual.threads.enabled (default false) and StormThreadFactory, which returns a virtual-thread ThreadFactory when the flag is on and a non-daemon, normal-priority platform-thread factory otherwise. Thread names are - in both modes. The factory is injected into the existing pools, so core sizes, queues, rejection policy and error handling are unchanged and the *.threads settings keep their meaning as a concurrency bound: - Thrift server handlers for the SASL, TLS and simple transports (Nimbus, Supervisor and DRPC). With the simple transport an executor is now always supplied when the flag is on; without a configured queue size it uses the same unbounded queue THsHaServer builds by default. - AsyncLocalizer download and task executors. - Nimbus AssignmentDistributionService and Supervisor heartbeat pools. - DRPCSpout background executor. - Worker shared executor exposed through TopologyContext, read from the merged topology conf so a topology can override the flag. Spout/bolt executor threads, worker transfer, JCQueue, Netty and timer threads are deliberately left on platform threads. Round-trip tests through NimbusClient assert that Thrift handlers run on virtual threads when the flag is on and on platform threads by default; unit tests cover the factory, the localizer executor and the worker shared executor. ThriftHandlerVirtualThreadsBench in storm-perf measures the flag against an in-process Nimbus Thrift server whose handler simulates blocking I/O. At equal *.threads the throughput is unchanged, while server-side OS threads drop from *.threads to roughly the number of cores. --- conf/defaults.yaml | 1 + .../ThriftHandlerVirtualThreadsBench.java | 286 ++++++++++++++++++ .../src/jvm/org/apache/storm/Config.java | 27 ++ .../storm/daemon/worker/WorkerState.java | 12 +- .../jvm/org/apache/storm/drpc/DRPCSpout.java | 5 +- .../security/auth/SimpleTransportPlugin.java | 17 +- .../auth/sasl/SaslTransportPlugin.java | 6 +- .../security/auth/tls/TlsTransportPlugin.java | 5 +- .../storm/utils/StormThreadFactory.java | 85 ++++++ .../worker/TestUtilsForWorkerState.java | 1 + .../storm/daemon/worker/WorkerStateTest.java | 27 ++ .../storm/utils/StormThreadFactoryTest.java | 97 ++++++ .../storm/daemon/supervisor/Supervisor.java | 3 +- .../storm/localizer/AsyncLocalizer.java | 11 +- .../nimbus/AssignmentDistributionService.java | 6 +- .../storm/localizer/AsyncLocalizerTest.java | 35 +++ .../apache/storm/security/auth/AuthTest.java | 90 +++++- 17 files changed, 702 insertions(+), 12 deletions(-) create mode 100644 examples/storm-perf/src/main/java/org/apache/storm/perf/ThriftHandlerVirtualThreadsBench.java create mode 100644 storm-client/src/jvm/org/apache/storm/utils/StormThreadFactory.java create mode 100644 storm-client/test/jvm/org/apache/storm/utils/StormThreadFactoryTest.java diff --git a/conf/defaults.yaml b/conf/defaults.yaml index a159f6bc133..7ba3ef7e458 100644 --- a/conf/defaults.yaml +++ b/conf/defaults.yaml @@ -38,6 +38,7 @@ storm.zookeeper.auth.password: null storm.zookeeper.ssl.enable: false storm.zookeeper.ssl.hostnameVerification: false storm.cluster.mode: "distributed" # can be distributed or local +storm.virtual.threads.enabled: false storm.local.mode.zmq: false storm.thrift.transport: "org.apache.storm.security.auth.SimpleTransportPlugin" storm.thrift.socket.timeout.ms: 600000 diff --git a/examples/storm-perf/src/main/java/org/apache/storm/perf/ThriftHandlerVirtualThreadsBench.java b/examples/storm-perf/src/main/java/org/apache/storm/perf/ThriftHandlerVirtualThreadsBench.java new file mode 100644 index 00000000000..3e467153caa --- /dev/null +++ b/examples/storm-perf/src/main/java/org/apache/storm/perf/ThriftHandlerVirtualThreadsBench.java @@ -0,0 +1,286 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License + */ + +package org.apache.storm.perf; + +import java.io.IOException; +import java.lang.management.ManagementFactory; +import java.lang.management.ThreadMXBean; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.storm.Config; +import org.apache.storm.generated.Nimbus; +import org.apache.storm.security.auth.SimpleTransportPlugin; +import org.apache.storm.security.auth.ThriftConnectionType; +import org.apache.storm.security.auth.ThriftServer; +import org.apache.storm.utils.NimbusClient; +import org.apache.storm.utils.Utils; + +/** + * Benchmark for {@code storm.virtual.threads.enabled} on the Nimbus Thrift handler pool. + * + *

Starts an in-process {@link ThriftServer} using {@link SimpleTransportPlugin} whose {@code getNimbusConf} + * handler sleeps for a configurable time to simulate blocking I/O (ZooKeeper, blob store), then hammers it with + * N client threads and reports throughput, latency percentiles, peak platform-thread count and RSS. Run it with + * the flag off and on to compare. + * + *

Build and run (JDK 25, Maven 3.9): + * 

+ * mvn -pl storm-client,examples/storm-perf -am install -DskipTests -Dcheckstyle.skip=true -q
+ * mvn -pl examples/storm-perf dependency:build-classpath -Dmdep.outputFile=target/cp.txt -q
+ * java -cp examples/storm-perf/target/classes:$(cat examples/storm-perf/target/cp.txt) \
+ *     org.apache.storm.perf.ThriftHandlerVirtualThreadsBench --mode both --clients 200 --calls 50 --io-ms 10
+ *
+ * + *

Options: {@code --mode off|on|both}, {@code --clients N}, {@code --calls M} (per client), + * {@code --io-ms L} (simulated handler I/O), {@code --threads T} ({@code nimbus.thrift.threads}), + * {@code --queue-size Q|none} ({@code nimbus.queue.size}; defaults to the shipped 100000, so both modes use a pool of + * {@code --threads} handlers; {@code none} removes the key so the platform-thread run falls back to THsHaServer's own + * pool, whose effective concurrency is 5), {@code --warmup W} (calls per client excluded from statistics). + * + *

For comparable RSS numbers run each mode in its own JVM ({@code --mode off} then {@code --mode on}). + */ +public final class ThriftHandlerVirtualThreadsBench { + + private static final String HOST = "localhost"; + + private ThriftHandlerVirtualThreadsBench() { + } + + public static void main(String[] args) throws Exception { + Map opts = parseArgs(args); + String mode = opts.getOrDefault("mode", "both").toLowerCase(Locale.ROOT); + int clients = Integer.parseInt(opts.getOrDefault("clients", "200")); + int calls = Integer.parseInt(opts.getOrDefault("calls", "50")); + long ioMs = Long.parseLong(opts.getOrDefault("io-ms", "10")); + int threads = Integer.parseInt(opts.getOrDefault("threads", "64")); + String queueOpt = opts.get("queue-size"); + boolean noQueue = "none".equalsIgnoreCase(queueOpt); + Integer queueSize = queueOpt == null || noQueue ? null : Integer.valueOf(queueOpt); + int warmup = Integer.parseInt(opts.getOrDefault("warmup", "5")); + + System.out.printf(Locale.ROOT, "clients=%d calls=%d io-ms=%d threads=%d queue-size=%s warmup=%d%n", + clients, calls, ioMs, threads, noQueue ? "none" : queueSize, warmup); + if (mode.equals("off") || mode.equals("both")) { + runPhase(false, clients, calls, ioMs, threads, queueSize, noQueue, warmup); + } + if (mode.equals("on") || mode.equals("both")) { + runPhase(true, clients, calls, ioMs, threads, queueSize, noQueue, warmup); + } + } + + private static void runPhase(boolean virtual, int clients, int calls, long ioMs, int threads, + Integer queueSize, boolean noQueue, int warmup) throws Exception { + Map conf = new HashMap<>(Utils.readDefaultConfig()); + conf.put(Config.STORM_THRIFT_TRANSPORT_PLUGIN, SimpleTransportPlugin.class.getName()); + conf.put(Config.NIMBUS_THRIFT_PORT, 0); + conf.put(Config.NIMBUS_THRIFT_THREADS, threads); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, virtual); + conf.put(Config.STORM_NIMBUS_RETRY_TIMES, 0); + if (noQueue) { + conf.remove(Config.NIMBUS_QUEUE_SIZE); + } else if (queueSize != null) { + conf.put(Config.NIMBUS_QUEUE_SIZE, queueSize); + } + + AtomicInteger inFlight = new AtomicInteger(); + AtomicInteger peakInFlight = new AtomicInteger(); + Nimbus.Iface handler = sleepingHandler(ioMs, inFlight, peakInFlight); + ThriftServer server = new ThriftServer(conf, new Nimbus.Processor<>(handler), ThriftConnectionType.NIMBUS); + Thread serveThread = new Thread(server::serve, "bench-thrift-serve"); + serveThread.setDaemon(true); + serveThread.start(); + while (!server.isServing()) { + Thread.sleep(10); + } + int port = server.getPort(); + + ThreadMXBean threadMx = ManagementFactory.getThreadMXBean(); + long rssBeforeKb = rssKb(); + int threadsBefore = threadMx.getThreadCount(); + AtomicInteger peakThreads = new AtomicInteger(threadsBefore); + Thread sampler = new Thread(() -> { + while (!Thread.currentThread().isInterrupted()) { + peakThreads.accumulateAndGet(threadMx.getThreadCount(), Math::max); + try { + Thread.sleep(20); + } catch (InterruptedException e) { + return; + } + } + }, "bench-thread-sampler"); + sampler.setDaemon(true); + sampler.start(); + + long[][] latenciesNs = new long[clients][calls]; + List failures = new ArrayList<>(); + CountDownLatch ready = new CountDownLatch(clients); + CountDownLatch go = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(clients); + for (int c = 0; c < clients; c++) { + final int clientIdx = c; + Thread t = new Thread(() -> { + try (NimbusClient client = NimbusClient.Builder.withConf(conf).withTimeout(60_000) + .buildWithNimbusHostPort(HOST, port)) { + Nimbus.Iface iface = client.getClient(); + for (int i = 0; i < warmup; i++) { + iface.getNimbusConf(); + } + ready.countDown(); + go.await(); + for (int i = 0; i < calls; i++) { + long start = System.nanoTime(); + iface.getNimbusConf(); + latenciesNs[clientIdx][i] = System.nanoTime() - start; + } + } catch (Throwable e) { + synchronized (failures) { + failures.add(e); + } + ready.countDown(); + } finally { + done.countDown(); + } + }, "bench-client-" + c); + t.start(); + } + ready.await(); + long startNs = System.nanoTime(); + go.countDown(); + done.await(); + long elapsedNs = System.nanoTime() - startNs; + + sampler.interrupt(); + sampler.join(1000); + long rssAfterKb = rssKb(); + server.stop(); + serveThread.join(5000); + + report(virtual, clients, calls, ioMs, threads, queueSize, latenciesNs, elapsedNs, failures.size(), + threadsBefore, peakThreads.get(), peakInFlight.get(), rssBeforeKb, rssAfterKb); + } + + private static void report(boolean virtual, int clients, int calls, long ioMs, int threads, Integer queueSize, + long[][] latenciesNs, long elapsedNs, int failures, int threadsBefore, int peakThreads, + int peakInFlight, long rssBeforeKb, long rssAfterKb) { + long[] all = Arrays.stream(latenciesNs).flatMapToLong(Arrays::stream).filter(v -> v > 0).sorted().toArray(); + int total = clients * calls; + double elapsedS = elapsedNs / 1e9; + double throughput = total / elapsedS; + double p50 = pct(all, 0.50); + double p90 = pct(all, 0.90); + double p99 = pct(all, 0.99); + double max = all.length == 0 ? 0 : all[all.length - 1] / 1e6; + double idealThroughput = (double) Math.min(clients, threads) * 1000.0 / ioMs; + + String label = virtual ? "virtual threads ON" : "virtual threads OFF"; + System.out.println(); + System.out.println("=== " + label + " ==="); + System.out.printf(Locale.ROOT, " calls : %d (%d failed)%n", total, failures); + System.out.printf(Locale.ROOT, " elapsed : %.2f s%n", elapsedS); + System.out.printf(Locale.ROOT, " throughput : %.0f calls/s (ideal at %d concurrent handlers: %.0f)%n", + throughput, Math.min(clients, threads), idealThroughput); + System.out.printf(Locale.ROOT, " latency p50/p90/p99: %.1f / %.1f / %.1f ms (max %.1f, floor %d)%n", + p50, p90, p99, max, ioMs); + System.out.printf(Locale.ROOT, " peak handlers busy : %d%n", peakInFlight); + System.out.printf(Locale.ROOT, " platform threads : %d before, %d peak (+%d)%n", + threadsBefore, peakThreads, peakThreads - threadsBefore); + System.out.printf(Locale.ROOT, " RSS : %d MB before, %d MB after%n", + rssBeforeKb / 1024, rssAfterKb / 1024); + System.out.printf(Locale.ROOT, + "JSON {\"virtual\":%b,\"clients\":%d,\"calls\":%d,\"ioMs\":%d,\"threads\":%d,\"queueSize\":%s," + + "\"throughput\":%.1f,\"p50Ms\":%.2f,\"p90Ms\":%.2f,\"p99Ms\":%.2f,\"maxMs\":%.2f," + + "\"peakHandlersBusy\":%d,\"platformThreadsPeak\":%d,\"platformThreadsBefore\":%d," + + "\"rssBeforeMb\":%d,\"rssAfterMb\":%d,\"failures\":%d}%n", + virtual, clients, calls, ioMs, threads, queueSize, throughput, p50, p90, p99, max, + peakInFlight, peakThreads, threadsBefore, rssBeforeKb / 1024, rssAfterKb / 1024, failures); + } + + private static double pct(long[] sorted, double q) { + if (sorted.length == 0) { + return 0; + } + int idx = (int) Math.min(sorted.length - 1, Math.round(q * (sorted.length - 1))); + return sorted[idx] / 1e6; + } + + /** + * A Nimbus handler whose {@code getNimbusConf} blocks for {@code ioMs}; every other method is unsupported. + */ + private static Nimbus.Iface sleepingHandler(long ioMs, AtomicInteger inFlight, AtomicInteger peakInFlight) { + InvocationHandler h = (Object proxy, Method method, Object[] margs) -> { + if (method.getName().equals("getNimbusConf")) { + int now = inFlight.incrementAndGet(); + peakInFlight.accumulateAndGet(now, Math::max); + try { + Thread.sleep(ioMs); + } finally { + inFlight.decrementAndGet(); + } + return "{}"; + } + if (method.getDeclaringClass() == Object.class) { + return method.invoke(inFlight, margs); + } + throw new UnsupportedOperationException(method.getName()); + }; + return (Nimbus.Iface) Proxy.newProxyInstance(Nimbus.Iface.class.getClassLoader(), + new Class[]{ Nimbus.Iface.class }, h); + } + + /** Resident set size from /proc/self/status, or -1 when unavailable. */ + private static long rssKb() { + try { + for (String line : Files.readAllLines(Paths.get("/proc/self/status"))) { + if (line.startsWith("VmRSS:")) { + return Long.parseLong(line.replaceAll("[^0-9]", "")); + } + } + } catch (IOException | RuntimeException e) { + // not Linux, or unreadable + } + return -1; + } + + private static Map parseArgs(String[] args) { + Map opts = new HashMap<>(); + for (int i = 0; i < args.length; i++) { + if (!args[i].startsWith("--")) { + throw new IllegalArgumentException("Unexpected argument: " + args[i]); + } + String key = args[i].substring(2); + if (i + 1 >= args.length || args[i + 1].startsWith("--")) { + throw new IllegalArgumentException("Missing value for --" + key); + } + opts.put(key, args[++i]); + } + return opts; + } +} diff --git a/storm-client/src/jvm/org/apache/storm/Config.java b/storm-client/src/jvm/org/apache/storm/Config.java index f9c09b8973a..3115f7f79af 100644 --- a/storm-client/src/jvm/org/apache/storm/Config.java +++ b/storm-client/src/jvm/org/apache/storm/Config.java @@ -851,6 +851,33 @@ public class Config extends HashMap { */ @IsInteger public static final String TOPOLOGY_WORKER_SHARED_THREAD_POOL_SIZE = "topology.worker.shared.thread.pool.size"; + /** + * When true, Storm's blocking I/O thread pools run on Java virtual threads instead of platform threads. + * Affected pools: the Nimbus and Supervisor Thrift server handlers (SASL, TLS and simple transports), the + * supervisor blob localizer download and task executors, the Nimbus assignment distribution service, the + * supervisor heartbeat executor, the DRPC spout background executor and the worker shared thread pool + * exposed through the TopologyContext. Pool sizes configured through the corresponding {@code *.threads} + * settings keep their meaning as a bound on concurrency. Spout/bolt executor threads, worker transfer + * threads and Netty event loops are never affected. + * + *

DRPC server handler pools ({@code drpc.worker.threads}, {@code drpc.invocations.threads}) are also + * affected, since the DRPC server uses the same transport plugins. + * + *

With the simple (non-authenticated) Thrift transport, Storm builds its own handler pool of + * {@code *.threads} threads whenever a {@code *.queue.size} is configured, which the shipped defaults do + * for Nimbus, Supervisor and DRPC. Only when the queue size is explicitly unset does {@code THsHaServer} + * fall back to its own pool, whose effective concurrency is 5; enabling this flag in that case raises it to + * the configured {@code *.threads} value. + * + *

All virtual threads in a JVM share one carrier scheduler sized to the number of available processors; + * blocking file I/O (for example blob downloads in the supervisor localizer) occupies a carrier, so + * operators enabling this on I/O-heavy supervisors should size {@code jdk.virtualThreadScheduler.parallelism} + * / {@code jdk.virtualThreadScheduler.maxPoolSize} accordingly. + * + *

A topology may override this key in its own config to control its worker shared executor. + */ + @IsBoolean + public static final String STORM_VIRTUAL_THREADS_ENABLED = "storm.virtual.threads.enabled"; /** * The interval in seconds to use for determining whether to throttle error reported to Zookeeper. For example, an interval of 10 * seconds with topology.max.error.report.per.interval set to 5 will only allow 5 errors to be reported to Zookeeper per task for every diff --git a/storm-client/src/jvm/org/apache/storm/daemon/worker/WorkerState.java b/storm-client/src/jvm/org/apache/storm/daemon/worker/WorkerState.java index 59aceb8d65e..bd8883112dc 100644 --- a/storm-client/src/jvm/org/apache/storm/daemon/worker/WorkerState.java +++ b/storm-client/src/jvm/org/apache/storm/daemon/worker/WorkerState.java @@ -68,6 +68,7 @@ import org.apache.storm.security.auth.IAutoCredentials; import org.apache.storm.serialization.ITupleSerializer; import org.apache.storm.serialization.KryoTupleSerializer; +import org.apache.storm.shade.com.google.common.annotations.VisibleForTesting; import org.apache.storm.shade.com.google.common.collect.ImmutableMap; import org.apache.storm.shade.com.google.common.collect.Sets; import org.apache.storm.task.WorkerTopologyContext; @@ -77,6 +78,7 @@ import org.apache.storm.utils.ConfigUtils; import org.apache.storm.utils.JCQueue; import org.apache.storm.utils.ObjectReader; +import org.apache.storm.utils.StormThreadFactory; import org.apache.storm.utils.SupervisorIfaceFactory; import org.apache.storm.utils.ThriftTopologyUtils; import org.apache.storm.utils.Utils; @@ -780,8 +782,16 @@ private Map, JCQueue> mkReceiveQueueMap(Map topologyC } private Map makeDefaultResources() { + return ImmutableMap.of(WorkerTopologyContext.SHARED_EXECUTOR, makeSharedExecutor(topologyConf)); + } + + /** + * Builds the worker shared executor, honoring the thread pool size and virtual-threads flag from the given (merged) conf. + */ + @VisibleForTesting + static ExecutorService makeSharedExecutor(Map conf) { int threadPoolSize = ObjectReader.getInt(conf.get(Config.TOPOLOGY_WORKER_SHARED_THREAD_POOL_SIZE)); - return ImmutableMap.of(WorkerTopologyContext.SHARED_EXECUTOR, Executors.newFixedThreadPool(threadPoolSize)); + return Executors.newFixedThreadPool(threadPoolSize, StormThreadFactory.create(conf, "worker-shared-executor")); } private Map makeUserResources() { diff --git a/storm-client/src/jvm/org/apache/storm/drpc/DRPCSpout.java b/storm-client/src/jvm/org/apache/storm/drpc/DRPCSpout.java index ccd52c34937..58cc861a978 100644 --- a/storm-client/src/jvm/org/apache/storm/drpc/DRPCSpout.java +++ b/storm-client/src/jvm/org/apache/storm/drpc/DRPCSpout.java @@ -44,6 +44,7 @@ import org.apache.storm.utils.ExtendedThreadPoolExecutor; import org.apache.storm.utils.ObjectReader; import org.apache.storm.utils.ServiceRegistry; +import org.apache.storm.utils.StormThreadFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -116,7 +117,9 @@ public void open(Map conf, TopologyContext context, SpoutOutputC if (localDrpcId == null) { background = new ExtendedThreadPoolExecutor(0, Integer.MAX_VALUE, 60L, TimeUnit.SECONDS, - new SynchronousQueue()); + new SynchronousQueue(), + StormThreadFactory.create(conf, + "drpc-spout-" + context.getThisTaskId() + "-background")); futuresMap = new HashMap<>(); int numTasks = context.getComponentTasks(context.getThisComponentId()).size(); int index = context.getThisTaskIndex(); diff --git a/storm-client/src/jvm/org/apache/storm/security/auth/SimpleTransportPlugin.java b/storm-client/src/jvm/org/apache/storm/security/auth/SimpleTransportPlugin.java index f52aacf26c8..5711138cabe 100644 --- a/storm-client/src/jvm/org/apache/storm/security/auth/SimpleTransportPlugin.java +++ b/storm-client/src/jvm/org/apache/storm/security/auth/SimpleTransportPlugin.java @@ -18,8 +18,11 @@ import java.net.UnknownHostException; import java.security.Principal; import java.util.HashSet; +import java.util.Locale; import java.util.Map; import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import javax.security.auth.Subject; @@ -35,6 +38,7 @@ import org.apache.storm.thrift.transport.TTransport; import org.apache.storm.thrift.transport.TTransportException; import org.apache.storm.thrift.transport.layered.TFramedTransport; +import org.apache.storm.utils.StormThreadFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -74,9 +78,18 @@ public TServer getServer(TProcessor processor) throws IOException, TTransportExc serverArgs.maxReadBufferBytes = maxBufferSize; - if (queueSize != null) { + if (queueSize != null || StormThreadFactory.isVirtualEnabled(topoConf)) { + // THsHaServer builds its own platform-thread pool (core 5, unbounded LinkedBlockingQueue) when no + // executor is supplied, so when virtual threads are requested we always supply an executor; + // without a configured queue size we use the same unbounded queue so no request is rejected + // that would not be rejected today. The shipped defaults always configure a queue size. + BlockingQueue workQueue = queueSize != null + ? new ArrayBlockingQueue<>(queueSize) + : new LinkedBlockingQueue<>(); serverArgs.executorService(new ThreadPoolExecutor(numWorkerThreads, numWorkerThreads, - 60, TimeUnit.SECONDS, new ArrayBlockingQueue(queueSize))); + 60, TimeUnit.SECONDS, workQueue, + StormThreadFactory.create(topoConf, + type.name().toLowerCase(Locale.ROOT) + "-handler"))); } //construct THsHaServer diff --git a/storm-client/src/jvm/org/apache/storm/security/auth/sasl/SaslTransportPlugin.java b/storm-client/src/jvm/org/apache/storm/security/auth/sasl/SaslTransportPlugin.java index 745d360a072..ea99a613fb5 100644 --- a/storm-client/src/jvm/org/apache/storm/security/auth/sasl/SaslTransportPlugin.java +++ b/storm-client/src/jvm/org/apache/storm/security/auth/sasl/SaslTransportPlugin.java @@ -16,6 +16,7 @@ import java.io.IOException; import java.net.Socket; import java.security.Principal; +import java.util.Locale; import java.util.Map; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.BlockingQueue; @@ -42,6 +43,7 @@ import org.apache.storm.thrift.transport.TTransportException; import org.apache.storm.thrift.transport.TTransportFactory; import org.apache.storm.utils.ExtendedThreadPoolExecutor; +import org.apache.storm.utils.StormThreadFactory; /** * Base class for SASL authentication plugin. @@ -86,7 +88,9 @@ public TServer getServer(TProcessor processor) throws IOException, TTransportExc workQueue = new ArrayBlockingQueue<>(queueSize); } ThreadPoolExecutor executorService = new ExtendedThreadPoolExecutor(numWorkerThreads, numWorkerThreads, - 60, TimeUnit.SECONDS, workQueue); + 60, TimeUnit.SECONDS, workQueue, + StormThreadFactory.create(conf, + type.name().toLowerCase(Locale.ROOT) + "-handler")); serverArgs.executorService(executorService); return new TThreadPoolServer(serverArgs); } diff --git a/storm-client/src/jvm/org/apache/storm/security/auth/tls/TlsTransportPlugin.java b/storm-client/src/jvm/org/apache/storm/security/auth/tls/TlsTransportPlugin.java index 211cd7bdcd6..cad7d90ad5a 100644 --- a/storm-client/src/jvm/org/apache/storm/security/auth/tls/TlsTransportPlugin.java +++ b/storm-client/src/jvm/org/apache/storm/security/auth/tls/TlsTransportPlugin.java @@ -18,6 +18,7 @@ import java.net.ServerSocket; import java.security.cert.Certificate; import java.security.cert.X509Certificate; +import java.util.Locale; import java.util.Map; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.BlockingQueue; @@ -42,6 +43,7 @@ import org.apache.storm.thrift.transport.TTransport; import org.apache.storm.thrift.transport.TTransportException; import org.apache.storm.utils.ExtendedThreadPoolExecutor; +import org.apache.storm.utils.StormThreadFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -107,7 +109,8 @@ public TServer getServer(TProcessor processor) throws IOException, TTransportExc workQueue = new ArrayBlockingQueue<>(queueSize); } ThreadPoolExecutor executorService = new ExtendedThreadPoolExecutor(numWorkerThreads, numWorkerThreads, - 60, TimeUnit.SECONDS, workQueue); + 60, TimeUnit.SECONDS, workQueue, + StormThreadFactory.create(conf, type.name().toLowerCase(Locale.ROOT) + "-handler")); serverArgs.executorService(executorService); tThreadPoolServer = new TThreadPoolServer(serverArgs); return tThreadPoolServer; diff --git a/storm-client/src/jvm/org/apache/storm/utils/StormThreadFactory.java b/storm-client/src/jvm/org/apache/storm/utils/StormThreadFactory.java new file mode 100644 index 00000000000..b9f03b3a433 --- /dev/null +++ b/storm-client/src/jvm/org/apache/storm/utils/StormThreadFactory.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.storm.utils; + +import java.util.Map; +import java.util.concurrent.ThreadFactory; +import org.apache.storm.Config; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Builds the {@link ThreadFactory} used by Storm's blocking I/O thread pools. + * + *

When {@link Config#STORM_VIRTUAL_THREADS_ENABLED} is true the factory creates virtual threads. + * Otherwise it creates non-daemon platform threads with normal priority, matching + * {@code Executors.defaultThreadFactory()}. In both modes threads are named + * {@code -} with {@code n} starting at 0. + */ +public final class StormThreadFactory { + + private static final Logger LOG = LoggerFactory.getLogger(StormThreadFactory.class); + + private StormThreadFactory() { + } + + /** + * Create a thread factory for a pool. + * + * @param conf cluster or topology configuration + * @param namePrefix prefix for the thread names, e.g. {@code "nimbus-thrift-handler"} + * @return a factory producing virtual or platform threads depending on the configuration + */ + public static ThreadFactory create(Map conf, String namePrefix) { + String prefix = namePrefix + "-"; + if (isVirtualEnabled(conf)) { + return Thread.ofVirtual().name(prefix, 0).factory(); + } + return Thread.ofPlatform() + .name(prefix, 0) + .daemon(false) + .priority(Thread.NORM_PRIORITY) + .factory(); + } + + /** + * Whether virtual threads are enabled for blocking I/O pools. + * + *

Boolean and String values are honoured (strings are parsed as boolean). Any other type + * is treated as disabled. This method never throws. + */ + public static boolean isVirtualEnabled(Map conf) { + if (conf == null) { + return false; + } + Object value = conf.get(Config.STORM_VIRTUAL_THREADS_ENABLED); + if (value == null) { + return false; + } + if (value instanceof Boolean) { + return (Boolean) value; + } + if (value instanceof String) { + return Boolean.parseBoolean((String) value); + } + LOG.warn("Ignoring {} value of type {} for {}; expected a boolean", value, value.getClass().getName(), + Config.STORM_VIRTUAL_THREADS_ENABLED); + return false; + } +} diff --git a/storm-client/test/jvm/org/apache/storm/daemon/worker/TestUtilsForWorkerState.java b/storm-client/test/jvm/org/apache/storm/daemon/worker/TestUtilsForWorkerState.java index 3941f1778e2..0987510250d 100644 --- a/storm-client/test/jvm/org/apache/storm/daemon/worker/TestUtilsForWorkerState.java +++ b/storm-client/test/jvm/org/apache/storm/daemon/worker/TestUtilsForWorkerState.java @@ -77,6 +77,7 @@ public static WorkerState getWorkerState(Map conf, String topolo topologyConf.put(Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS, 30); topologyConf.put(Config.TOPOLOGY_TRANSFER_BUFFER_SIZE, 1000); topologyConf.put(Config.TOPOLOGY_TRANSFER_BATCH_SIZE, 1); + topologyConf.put(Config.TOPOLOGY_WORKER_SHARED_THREAD_POOL_SIZE, 1); IStateStorage stateStorage = null; IStormClusterState stormClusterState = null; diff --git a/storm-client/test/jvm/org/apache/storm/daemon/worker/WorkerStateTest.java b/storm-client/test/jvm/org/apache/storm/daemon/worker/WorkerStateTest.java index d738f61ef77..21669670475 100644 --- a/storm-client/test/jvm/org/apache/storm/daemon/worker/WorkerStateTest.java +++ b/storm-client/test/jvm/org/apache/storm/daemon/worker/WorkerStateTest.java @@ -33,6 +33,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -138,4 +140,29 @@ public void testTransferLocalBatchDropsTuplesForUnknownTasks() throws TException ConfigUtils.setInstance(previousConfigUtils); } } + + @Test + public void sharedExecutorUsesVirtualThreadsWhenEnabled() throws Exception { + Map conf = new HashMap<>(); + conf.put(Config.TOPOLOGY_WORKER_SHARED_THREAD_POOL_SIZE, 2); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + ExecutorService pool = WorkerState.makeSharedExecutor(conf); + try { + assertTrue(pool.submit(() -> Thread.currentThread().isVirtual()).get(10, TimeUnit.SECONDS)); + } finally { + pool.shutdownNow(); + } + } + + @Test + public void sharedExecutorUsesPlatformThreadsByDefault() throws Exception { + Map conf = new HashMap<>(); + conf.put(Config.TOPOLOGY_WORKER_SHARED_THREAD_POOL_SIZE, 2); + ExecutorService pool = WorkerState.makeSharedExecutor(conf); + try { + assertFalse(pool.submit(() -> Thread.currentThread().isVirtual()).get(10, TimeUnit.SECONDS)); + } finally { + pool.shutdownNow(); + } + } } diff --git a/storm-client/test/jvm/org/apache/storm/utils/StormThreadFactoryTest.java b/storm-client/test/jvm/org/apache/storm/utils/StormThreadFactoryTest.java new file mode 100644 index 00000000000..78b24bca904 --- /dev/null +++ b/storm-client/test/jvm/org/apache/storm/utils/StormThreadFactoryTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.storm.utils; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ThreadFactory; + +import org.apache.storm.Config; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class StormThreadFactoryTest { + + @Test + public void flagOffCreatesPlatformThreadsWithPrefix() throws Exception { + Map conf = new HashMap<>(); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, false); + ThreadFactory factory = StormThreadFactory.create(conf, "test-pool"); + + Thread t = factory.newThread(() -> { }); + assertFalse(t.isVirtual()); + assertFalse(t.isDaemon()); + assertEquals(Thread.NORM_PRIORITY, t.getPriority()); + assertEquals("test-pool-0", t.getName()); + assertEquals("test-pool-1", factory.newThread(() -> { }).getName()); + } + + @Test + public void missingFlagBehavesAsOff() { + ThreadFactory factory = StormThreadFactory.create(new HashMap<>(), "test-pool"); + assertFalse(factory.newThread(() -> { }).isVirtual()); + } + + @Test + public void flagOnCreatesVirtualThreadsWithPrefix() throws Exception { + Map conf = new HashMap<>(); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + ThreadFactory factory = StormThreadFactory.create(conf, "test-pool"); + + Thread t = factory.newThread(() -> { }); + assertTrue(t.isVirtual()); + assertEquals("test-pool-0", t.getName()); + assertEquals("test-pool-1", factory.newThread(() -> { }).getName()); + } + + @Test + public void flagOnAcceptsStringValue() { + Map conf = new HashMap<>(); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, "true"); + assertTrue(StormThreadFactory.create(conf, "test-pool").newThread(() -> { }).isVirtual()); + } + + @Test + public void platformThreadForcesNormPriorityRegardlessOfCallerPriority() throws Exception { + Map conf = new HashMap<>(); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, false); + ThreadFactory factory = StormThreadFactory.create(conf, "test-pool"); + + Thread callerThread = Thread.currentThread(); + int originalPriority = callerThread.getPriority(); + try { + callerThread.setPriority(Thread.MAX_PRIORITY); + Thread t = factory.newThread(() -> { }); + assertEquals(Thread.NORM_PRIORITY, t.getPriority()); + } finally { + callerThread.setPriority(originalPriority); + } + } + + @Test + public void unexpectedValueTypeBehavesAsOff() { + Map conf = new HashMap<>(); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, Integer.valueOf(1)); + ThreadFactory factory = StormThreadFactory.create(conf, "test-pool"); + assertFalse(factory.newThread(() -> { }).isVirtual()); + } +} diff --git a/storm-server/src/main/java/org/apache/storm/daemon/supervisor/Supervisor.java b/storm-server/src/main/java/org/apache/storm/daemon/supervisor/Supervisor.java index 50b1bba9f44..3677b7ec84b 100644 --- a/storm-server/src/main/java/org/apache/storm/daemon/supervisor/Supervisor.java +++ b/storm-server/src/main/java/org/apache/storm/daemon/supervisor/Supervisor.java @@ -77,6 +77,7 @@ import org.apache.storm.utils.ObjectReader; import org.apache.storm.utils.ServerConfigUtils; import org.apache.storm.utils.ShellUtils; +import org.apache.storm.utils.StormThreadFactory; import org.apache.storm.utils.Time; import org.apache.storm.utils.Utils; import org.apache.storm.utils.VersionInfo; @@ -150,7 +151,7 @@ public Supervisor(Map conf, IContext sharedContext, ISupervisor this.upTime = Utils.makeUptimeComputer(); this.stormVersion = VersionInfo.getVersion(); this.sharedContext = sharedContext; - this.heartbeatExecutor = Executors.newFixedThreadPool(1); + this.heartbeatExecutor = Executors.newFixedThreadPool(1, StormThreadFactory.create(conf, "supervisor-heartbeat")); this.authorizationHandler = StormCommon.mkAuthorizationHandler( (String) conf.get(DaemonConfig.SUPERVISOR_AUTHORIZER), conf); if (authorizationHandler == null && conf.get(DaemonConfig.NIMBUS_AUTHORIZER) != null) { diff --git a/storm-server/src/main/java/org/apache/storm/localizer/AsyncLocalizer.java b/storm-server/src/main/java/org/apache/storm/localizer/AsyncLocalizer.java index c9ad9e1c3fc..bc434df2080 100644 --- a/storm-server/src/main/java/org/apache/storm/localizer/AsyncLocalizer.java +++ b/storm-server/src/main/java/org/apache/storm/localizer/AsyncLocalizer.java @@ -54,13 +54,13 @@ import org.apache.storm.generated.StormTopology; import org.apache.storm.metric.StormMetricsRegistry; import org.apache.storm.shade.com.google.common.annotations.VisibleForTesting; -import org.apache.storm.shade.com.google.common.util.concurrent.ThreadFactoryBuilder; import org.apache.storm.thrift.transport.TTransportException; import org.apache.storm.utils.ConfigUtils; import org.apache.storm.utils.NimbusLeaderNotFoundException; import org.apache.storm.utils.ObjectReader; import org.apache.storm.utils.ServerConfigUtils; import org.apache.storm.utils.ServerUtils; +import org.apache.storm.utils.StormThreadFactory; import org.apache.storm.utils.Utils; import org.apache.storm.utils.WrappedKeyNotFoundException; import org.slf4j.Logger; @@ -131,9 +131,9 @@ public class AsyncLocalizer implements AutoCloseable { int downloadThreadPoolSize = ObjectReader.getInt(conf.get(DaemonConfig.SUPERVISOR_BLOBSTORE_DOWNLOAD_THREAD_COUNT), 5); downloadExecService = Executors.newScheduledThreadPool(downloadThreadPoolSize, - new ThreadFactoryBuilder().setNameFormat("AsyncLocalizer Download Executor - %d").build()); + StormThreadFactory.create(conf, "AsyncLocalizer-Download-Executor")); taskExecService = Executors.newScheduledThreadPool(3, - new ThreadFactoryBuilder().setNameFormat("AsyncLocalizer Task Executor - %d").build()); + StormThreadFactory.create(conf, "AsyncLocalizer-Task-Executor")); reconstructLocalizedResources(); symlinksDisabled = (boolean) conf.getOrDefault(Config.DISABLE_SYMLINKS, false); @@ -144,6 +144,11 @@ public AsyncLocalizer(Map conf, StormMetricsRegistry metricsRegi this(conf, AdvancedFSOps.make(conf), ConfigUtils.supervisorLocalDir(conf), metricsRegistry); } + @VisibleForTesting + ScheduledExecutorService getDownloadExecService() { + return downloadExecService; + } + @VisibleForTesting LocallyCachedBlob getTopoJar(final String topologyId, String owner) { return topologyBlobs.computeIfAbsent(ConfigUtils.masterStormJarKey(topologyId), diff --git a/storm-server/src/main/java/org/apache/storm/nimbus/AssignmentDistributionService.java b/storm-server/src/main/java/org/apache/storm/nimbus/AssignmentDistributionService.java index ce11722f950..012d0ad339e 100644 --- a/storm-server/src/main/java/org/apache/storm/nimbus/AssignmentDistributionService.java +++ b/storm-server/src/main/java/org/apache/storm/nimbus/AssignmentDistributionService.java @@ -30,6 +30,7 @@ import org.apache.storm.scheduler.INodeAssignmentSentCallBack; import org.apache.storm.utils.ConfigUtils; import org.apache.storm.utils.ObjectReader; +import org.apache.storm.utils.StormThreadFactory; import org.apache.storm.utils.SupervisorClient; import org.apache.storm.utils.Time; import org.slf4j.Logger; @@ -110,6 +111,8 @@ public static AssignmentDistributionService getInstance(Map conf, INodeAssignmen */ public void prepare(Map conf, INodeAssignmentSentCallBack callBack) { this.conf = conf; + @SuppressWarnings("unchecked") + Map typedConf = (Map) conf; this.sendAssignmentCallback = callBack; this.random = new Random(47); @@ -121,7 +124,8 @@ public void prepare(Map conf, INodeAssignmentSentCallBack callBack) { this.assignmentsQueue.put(i, new LinkedBlockingQueue(queueSize)); } //start the thread pool - this.service = Executors.newFixedThreadPool(threadsNum); + this.service = Executors.newFixedThreadPool(threadsNum, + StormThreadFactory.create(typedConf, "nimbus-assignment-distribution")); this.active = true; //start the threads for (int i = 0; i < threadsNum; i++) { diff --git a/storm-server/src/test/java/org/apache/storm/localizer/AsyncLocalizerTest.java b/storm-server/src/test/java/org/apache/storm/localizer/AsyncLocalizerTest.java index be67e0b34e3..212b5660de7 100644 --- a/storm-server/src/test/java/org/apache/storm/localizer/AsyncLocalizerTest.java +++ b/storm-server/src/test/java/org/apache/storm/localizer/AsyncLocalizerTest.java @@ -1040,6 +1040,41 @@ LocalizedResource getBlob(LocalResource localResource, PortAndAssignment pna, Bl } } + @Test + public void downloadExecutorUsesVirtualThreadsWhenEnabled() throws Exception { + try (TmpPath stormLocal = new TmpPath(); TmpPath localizerRoot = new TmpPath()) { + Map conf = new HashMap<>(); + conf.put(Config.STORM_LOCAL_DIR, stormLocal.getPath()); + conf.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + AdvancedFSOps ops = AdvancedFSOps.make(conf); + AsyncLocalizer localizer = new AsyncLocalizer(conf, ops, localizerRoot.getPath(), new StormMetricsRegistry()); + try { + boolean onVirtual = localizer.getDownloadExecService() + .submit(() -> Thread.currentThread().isVirtual()).get(10, TimeUnit.SECONDS); + assertTrue(onVirtual); + } finally { + localizer.close(); + } + } + } + + @Test + public void downloadExecutorUsesPlatformThreadsByDefault() throws Exception { + try (TmpPath stormLocal = new TmpPath(); TmpPath localizerRoot = new TmpPath()) { + Map conf = new HashMap<>(); + conf.put(Config.STORM_LOCAL_DIR, stormLocal.getPath()); + AdvancedFSOps ops = AdvancedFSOps.make(conf); + AsyncLocalizer localizer = new AsyncLocalizer(conf, ops, localizerRoot.getPath(), new StormMetricsRegistry()); + try { + boolean onVirtual = localizer.getDownloadExecService() + .submit(() -> Thread.currentThread().isVirtual()).get(10, TimeUnit.SECONDS); + assertFalse(onVirtual); + } finally { + localizer.close(); + } + } + } + static class TestInputStreamWithMeta extends InputStreamWithMeta { private final long version; private final long fileLength; diff --git a/storm-server/src/test/java/org/apache/storm/security/auth/AuthTest.java b/storm-server/src/test/java/org/apache/storm/security/auth/AuthTest.java index cc2e4acaee7..24ab173a0b5 100644 --- a/storm-server/src/test/java/org/apache/storm/security/auth/AuthTest.java +++ b/storm-server/src/test/java/org/apache/storm/security/auth/AuthTest.java @@ -647,4 +647,92 @@ public void impersonationAuthorizerTest() throws Exception { public interface MyBiConsumer { void accept(T t, U u) throws Exception; } -} \ No newline at end of file + + @Test + public void simpleTransportRunsHandlerOnVirtualThreadWhenEnabled() throws Exception { + Nimbus.Iface impl = mock(Nimbus.Iface.class); + final AtomicReference handlerOnVirtualThread = new AtomicReference<>(); + doAnswer((invocation) -> { + handlerOnVirtualThread.set(Thread.currentThread().isVirtual()); + return null; + }).when(impl).activate(anyString()); + + Map extra = new HashMap<>(); + extra.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + withServer(null, SimpleTransportPlugin.class, impl, null, extra, + (ThriftServer server, Map conf) -> { + try (NimbusClient client = NimbusClient.Builder.withConf(conf).withTimeout(NIMBUS_TIMEOUT) + .buildWithNimbusHostPort("localhost", server.getPort())) { + client.getClient().activate("virtual_thread_test_topology"); + } + assertEquals(Boolean.TRUE, handlerOnVirtualThread.get()); + }); + } + + @Test + public void simpleTransportWithQueueSizeRunsHandlerOnVirtualThreadWhenEnabled() throws Exception { + Nimbus.Iface impl = mock(Nimbus.Iface.class); + final AtomicReference handlerOnVirtualThread = new AtomicReference<>(); + doAnswer((invocation) -> { + handlerOnVirtualThread.set(Thread.currentThread().isVirtual()); + return null; + }).when(impl).activate(anyString()); + + Map extra = new HashMap<>(); + extra.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + extra.put(Config.NIMBUS_QUEUE_SIZE, 8); + withServer(null, SimpleTransportPlugin.class, impl, null, extra, + (ThriftServer server, Map conf) -> { + try (NimbusClient client = NimbusClient.Builder.withConf(conf).withTimeout(NIMBUS_TIMEOUT) + .buildWithNimbusHostPort("localhost", server.getPort())) { + client.getClient().activate("virtual_thread_test_topology"); + } + assertEquals(Boolean.TRUE, handlerOnVirtualThread.get()); + }); + } + + @Test + public void simpleTransportRunsHandlerOnPlatformThreadByDefault() throws Exception { + Nimbus.Iface impl = mock(Nimbus.Iface.class); + final AtomicReference handlerOnVirtualThread = new AtomicReference<>(); + doAnswer((invocation) -> { + handlerOnVirtualThread.set(Thread.currentThread().isVirtual()); + return null; + }).when(impl).activate(anyString()); + + withServer(SimpleTransportPlugin.class, impl, + (ThriftServer server, Map conf) -> { + try (NimbusClient client = NimbusClient.Builder.withConf(conf).withTimeout(NIMBUS_TIMEOUT) + .buildWithNimbusHostPort("localhost", server.getPort())) { + client.getClient().activate("platform_thread_test_topology"); + } + assertEquals(Boolean.FALSE, handlerOnVirtualThread.get()); + }); + } + + @Test + public void digestTransportRunsHandlerOnVirtualThreadWhenEnabled() throws Exception { + Nimbus.Iface impl = mock(Nimbus.Iface.class); + final AtomicReference handlerOnVirtualThread = new AtomicReference<>(); + final AtomicReference user = new AtomicReference<>(); + doAnswer((invocation) -> { + handlerOnVirtualThread.set(Thread.currentThread().isVirtual()); + user.set(new ReqContext(ReqContext.context())); + return null; + }).when(impl).activate(anyString()); + + Map extra = new HashMap<>(); + extra.put(Config.STORM_VIRTUAL_THREADS_ENABLED, true); + withServer(DIGEST_JAAS_CONF, DigestSaslTransportPlugin.class, impl, null, extra, + (ThriftServer server, Map conf) -> { + try (NimbusClient client = NimbusClient.Builder.withConf(conf).withTimeout(NIMBUS_TIMEOUT) + .buildWithNimbusHostPort("localhost", server.getPort())) { + client.getClient().activate("virtual_thread_digest_test_topology"); + } + assertEquals(Boolean.TRUE, handlerOnVirtualThread.get()); + // ReqContext is still populated per request on a virtual handler thread + assertNotNull(user.get()); + assertEquals("bob", user.get().principal().getName()); + }); + } +}