Fix ResourceTracker in benchmark to work fine with VIRTUAL_THREAD (#234)

benchmark/src/main/java/com/linecorp/decaton/benchmark/DecatonRunner.java

@@ -16,17 +16,24 @@
 
 package com.linecorp.decaton.benchmark;
 
+import java.lang.management.ThreadInfo;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 
+import com.sun.management.ThreadMXBean;
+
 import com.linecorp.decaton.processor.TaskMetadata;
 import com.linecorp.decaton.processor.metrics.Metrics;
 import com.linecorp.decaton.processor.runtime.DecatonTask;
@@ -43,7 +50,9 @@
 import io.micrometer.core.instrument.Clock;
 import io.micrometer.core.instrument.logging.LoggingMeterRegistry;
 import io.micrometer.core.instrument.logging.LoggingRegistryConfig;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class DecatonRunner implements Runner {
     private static final Map<String, Function<String, Object>> propertyConstructors =
             new HashMap<String, Function<String, Object>>() {{
@@ -52,6 +61,7 @@ public class DecatonRunner implements Runner {
                 put(ProcessorProperties.CONFIG_LOGGING_MDC_ENABLED.name(), Boolean::parseBoolean);
             }};
 
+    private SubPartitionRuntime subPartitionRuntime;
     private ProcessorSubscription subscription;
     private LoggingMeterRegistry registry;
 
@@ -71,7 +81,7 @@ public void init(Config config, Recording recording, ResourceTracker resourceTra
         // value than zero with the default "latest" reset policy.
         props.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
 
-        SubPartitionRuntime subPartitionRuntime = SubPartitionRuntime.THREAD_POOL;
+        subPartitionRuntime = SubPartitionRuntime.THREAD_POOL;
         List<Property<?>> properties = new ArrayList<>();
         for (Map.Entry<String, String> entry : config.parameters().entrySet()) {
             String name = entry.getKey();
@@ -97,6 +107,8 @@ public String get(String key) {
         }, Clock.SYSTEM);
         Metrics.register(registry);
 
+        maybeSetupForkJoinPoolDrip();
+
         CountDownLatch startLatch = new CountDownLatch(1);
 
         subscription = SubscriptionBuilder
@@ -113,22 +125,69 @@ public String get(String key) {
                                                                 TaskMetadata.builder().build(), task, bytes);
                                                     })
                                          .thenProcess(
-                                                 (ctx, task) -> {
-                                                     resourceTracker.track(Thread.currentThread().getId());
-                                                     recording.process(task);
-                                                 }))
+                                                 (ctx, task) -> recording.process(task)))
                 .stateListener(state -> {
                     if (state == SubscriptionStateListener.State.RUNNING) {
                         startLatch.countDown();
                     }
                 })
                 .build();
-        resourceTracker.track(subscription.getId());
         subscription.start();
 
         startLatch.await();
     }
 
+    /**
+     * This value comes from VirtualThread#createDefaultScheduler(), which defaults keep-alive of ForkJoinPool
+     * to 30 seconds (as of Java21).
+     */
+    private static final long DRIP_INTERVAL = 30;
+    /**
+     * The intention of this method is to setup a scheduled work submitting some VirtualThread instances
+     * periodically, in order to make sure the {@link ForkJoinPool} to keep-alive the created carrier thread
+     * during warmup phase.
+     * It is necessary because {@link #onWarmupComplete(ResourceTracker)} depends on the alive list of
+     * ForkJoinPool threads to observe cpu time and memory allocation profile, which cannot be obtained from
+     * VirtualThread instances directly.
+     */
+    private void maybeSetupForkJoinPoolDrip() {
+        if (subPartitionRuntime != SubPartitionRuntime.VIRTUAL_THREAD) {
+            return;
+        }
+        ScheduledExecutorService executor = Executors.newScheduledThreadPool(
+                1, Thread.ofPlatform().daemon().factory());
+        executor.scheduleAtFixedRate(() -> {
+            for (int i = 0; i < Runtime.getRuntime().availableProcessors(); i++) {
+                Thread.startVirtualThread(() -> {
+                    try {
+                        TimeUnit.SECONDS.sleep(DRIP_INTERVAL);
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                });
+            }
+        }, 0, DRIP_INTERVAL, TimeUnit.SECONDS);
+    }
+
+    @Override
+    public void onWarmupComplete(ResourceTracker resourceTracker) {
+        // As we support VIRTUAL_THREAD runtime now, we can't use the technique to obtain the target Thread ID
+        // in the execution context, as in VirtualThread gets created and discarded for all tasks separately.
+        // Instead we need to observe resource usage based on the carrier thread (platform thread) side, while
+        // stdlib doesn't provide a way to resolve the carrier thread's ID.
+        ThreadMXBean threadMxBean = ResourceTracker.getSunThreadMxBean();
+        ThreadInfo[] threadInfos = threadMxBean.dumpAllThreads(true, true);
+        for (ThreadInfo threadInfo : threadInfos) {
+            String name = threadInfo.getThreadName();
+            log.debug("Tracking target check for thread name: {}", name);
+            if (name.startsWith("DecatonSubscriptionThread-") || name.startsWith("PartitionProcessorThread-") ||
+                subPartitionRuntime == SubPartitionRuntime.VIRTUAL_THREAD && name.startsWith("ForkJoinPool-")) {
+                log.info("Tracking resource of thread {} - {}", threadInfo.getThreadId(), name);
+                resourceTracker.track(threadInfo.getThreadId());
+            }
+        }
+    }
+
     @Override
     public void close() throws Exception {
         if (registry != null) {

benchmark/src/main/java/com/linecorp/decaton/benchmark/InProcessExecution.java

@@ -78,6 +78,7 @@ public BenchmarkResult execute(Config config, Consumer<Stage> stageCallback) thr
                 awaitJITGetsSettled();
             }
             profiling.stop();
+            runner.onWarmupComplete(resourceTracker);
 
             profiling.start();
             JvmTracker jvmTracker = JvmTracker.create();

benchmark/src/main/java/com/linecorp/decaton/benchmark/ResourceTracker.java

@@ -38,7 +38,6 @@ public static class TrackingValues {
         long allocatedBytes;
     }
 
-    private static final ThreadLocal<Boolean> tracking = ThreadLocal.withInitial(() -> false);
     private final ConcurrentMap<Long, TrackingValues> targetThreadIds;
     private final ThreadMXBean threadMxBean;
 
@@ -47,7 +46,7 @@ public ResourceTracker() {
         threadMxBean = getSunThreadMxBean();
     }
 
-    private static ThreadMXBean getSunThreadMxBean() {
+    static ThreadMXBean getSunThreadMxBean() {
         java.lang.management.ThreadMXBean mxBean = ManagementFactory.getThreadMXBean();
         if (mxBean instanceof ThreadMXBean) {
             return (ThreadMXBean) mxBean;
@@ -65,10 +64,6 @@ private static ThreadMXBean getSunThreadMxBean() {
      * @param threadId thread ID to start tracking.
      */
     public void track(long threadId) {
-        if (tracking.get()) {
-            return;
-        }
-        tracking.set(true);
         targetThreadIds.computeIfAbsent(
                 threadId,
                 key -> new TrackingValues(threadMxBean.getThreadCpuTime(threadId),

benchmark/src/main/java/com/linecorp/decaton/benchmark/Runner.java

@@ -46,6 +46,12 @@ class Config {
      */
     void init(Config config, Recording recording, ResourceTracker resourceTracker) throws InterruptedException;
 
+    /**
+     * Called after all warmup tasks are processed and before the stage to proceed for actual measurement.
+     * @param resourceTracker resource tracking interface (use is optional).
+     */
+    void onWarmupComplete(ResourceTracker resourceTracker);
+
     @Override
     default void close() throws Exception {
         // noop by default