Support profiling in an virtual thread.

Instead of using `Thread.currentThread().threadId()` for the `threadId` in the `Profiler`, we maintain a map of `threadId` to `laneId` for virtual threads so a same lane in the UI can be re-used for different virtual threads.

This change requires using `Profiler#profileableVirtualThreadFactory()` to create VirtualThread which can use `Profiler`.

PiperOrigin-RevId: 622814148
Change-Id: I5a6bad039c95b287456597dceb158b487f61fc86

Author: coeuvre

src/main/java/com/google/devtools/build/lib/bazel/BUILD

@@ -44,6 +44,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/pkgcache",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/local_repository_rule",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/new_local_repository_function",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/new_local_repository_rule",

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -77,6 +77,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.pkgcache.PackageOptions;
+import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.rules.repository.LocalRepositoryFunction;
 import com.google.devtools.build.lib.rules.repository.LocalRepositoryRule;
 import com.google.devtools.build.lib.rules.repository.NewLocalRepositoryFunction;
@@ -121,6 +122,7 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -325,8 +327,11 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
             starlarkRepositoryFunction.setWorkerExecutorService(
                 (ExecutorService)
                     Executors.class
-                        .getDeclaredMethod("newVirtualThreadPerTaskExecutor")
-                        .invoke(null));
+                        .getDeclaredMethod("newThreadPerTaskExecutor", ThreadFactory.class)
+                        .invoke(
+                            null,
+                            Profiler.instance()
+                                .profileableVirtualThreadFactory("starlark-repository-")));
           } catch (ReflectiveOperationException e) {
             if (repoOptions.workerForRepoFetching == RepositoryOptions.WorkerForRepoFetching.AUTO) {
               starlarkRepositoryFunction.setWorkerExecutorService(null);

src/main/java/com/google/devtools/build/lib/profiler/JsonTraceFileWriter.java

@@ -85,7 +85,9 @@ public void start() {
   }
 
   public void enqueue(TraceData data) {
-    if (!metadataPosted.get()) {
+    // We assign a virtual lane for virtual thread and the metadata for the virtual lane is posted
+    // at creation time.
+    if (!Thread.currentThread().isVirtual() && !metadataPosted.get()) {
       metadataPosted.set(Boolean.TRUE);
       queue.add(new ThreadMetadata());
     }

src/main/java/com/google/devtools/build/lib/profiler/Profiler.java

@@ -46,12 +46,14 @@
 import java.util.PriorityQueue;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.zip.GZIPOutputStream;
 import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
 
 /**
  * Blaze internal profiler. Provides facility to report various Blaze tasks and store them
@@ -274,8 +276,7 @@ private static final class SlowestTaskAggregator {
 
     // @ThreadSafe
     void add(TaskData taskData) {
-      Extrema<SlowTask> extrema =
-          extremaAggregators[(int) (Thread.currentThread().getId() % SHARDS)];
+      Extrema<SlowTask> extrema = extremaAggregators[(int) (taskData.threadId % SHARDS)];
       synchronized (extrema) {
         extrema.aggregate(new SlowTask(taskData));
       }
@@ -659,12 +660,7 @@ private void logTask(
    */
   public void logSimpleTask(long startTimeNanos, ProfilerTask type, String description) {
     if (clock != null) {
-      logTask(
-          Thread.currentThread().getId(),
-          startTimeNanos,
-          clock.nanoTime() - startTimeNanos,
-          type,
-          description);
+      logTask(getLaneId(), startTimeNanos, clock.nanoTime() - startTimeNanos, type, description);
     }
   }
 
@@ -682,12 +678,7 @@ public void logSimpleTask(long startTimeNanos, ProfilerTask type, String descrip
    */
   public void logSimpleTask(
       long startTimeNanos, long stopTimeNanos, ProfilerTask type, String description) {
-    logTask(
-        Thread.currentThread().getId(),
-        startTimeNanos,
-        stopTimeNanos - startTimeNanos,
-        type,
-        description);
+    logTask(getLaneId(), startTimeNanos, stopTimeNanos - startTimeNanos, type, description);
   }
 
   /**
@@ -702,12 +693,12 @@ public void logSimpleTask(
    */
   public void logSimpleTaskDuration(
       long startTimeNanos, Duration duration, ProfilerTask type, String description) {
-    logTask(Thread.currentThread().getId(), startTimeNanos, duration.toNanos(), type, description);
+    logTask(getLaneId(), startTimeNanos, duration.toNanos(), type, description);
   }
 
   /** Used to log "events" happening at a specific time - tasks with zero duration. */
   public void logEventAtTime(long atTimeNanos, ProfilerTask type, String description) {
-    logTask(Thread.currentThread().getId(), atTimeNanos, 0, type, description);
+    logTask(getLaneId(), atTimeNanos, 0, type, description);
   }
 
   /** Used to log "events" - tasks with zero duration. */
@@ -741,7 +732,7 @@ private SilentCloseable reallyProfile(long laneId, ProfilerTask type, String des
    * @param description task description. May be stored until the end of the build.
    */
   public SilentCloseable profile(ProfilerTask type, String description) {
-    return profile(Thread.currentThread().getId(), type, description);
+    return profile(getLaneId(), type, description);
   }
 
   private SilentCloseable profile(long laneId, ProfilerTask type, String description) {
@@ -753,7 +744,7 @@ private SilentCloseable profile(long laneId, ProfilerTask type, String descripti
    * profiling.
    */
   public SilentCloseable profile(ProfilerTask type, Supplier<String> description) {
-    return profile(Thread.currentThread().getId(), type, description);
+    return profile(getLaneId(), type, description);
   }
 
   private SilentCloseable profile(long laneId, ProfilerTask type, Supplier<String> description) {
@@ -799,7 +790,7 @@ public SilentCloseable profileAction(
       final long startTimeNanos = clock.nanoTime();
       return () ->
           completeAction(
-              Thread.currentThread().getId(),
+              getLaneId(),
               startTimeNanos,
               type,
               description,
@@ -824,7 +815,7 @@ private boolean countAction(ProfilerTask type, TaskData taskData) {
   }
 
   public void completeTask(long startTimeNanos, ProfilerTask type, String description) {
-    completeTask(Thread.currentThread().getId(), startTimeNanos, type, description);
+    completeTask(getLaneId(), startTimeNanos, type, description);
   }
 
   /** Records the end of the task. */
@@ -969,72 +960,145 @@ public interface FutureSupplier<T> {
   }
 
   public <T> ListenableFuture<T> profileAsync(
-      ProfilerTaskType type, String description, FutureSupplier<T> futureSupplier) {
+      String prefix, String description, FutureSupplier<T> futureSupplier) {
     if (!(isActive() && isProfiling(ProfilerTask.INFO))) {
       return futureSupplier.get(new ScopedProfiler(/* active= */ false, 0));
     }
 
-    long laneId = laneIdGenerator.acquire(type);
+    var lane = multiLaneGenerator.acquire(prefix);
     final long startTimeNanos = clock.nanoTime();
-    var scopedProfiler = new ScopedProfiler(/* active= */ true, laneId);
+    var scopedProfiler = new ScopedProfiler(/* active= */ true, lane.id());
     var future = futureSupplier.get(scopedProfiler);
     future.addListener(
         () -> {
           long endTimeNanos = clock.nanoTime();
           long duration = endTimeNanos - startTimeNanos;
           recordTask(
-              new TaskData(laneId, startTimeNanos, duration, ProfilerTask.INFO, description));
-          laneIdGenerator.release(type, laneId);
+              new TaskData(lane.id(), startTimeNanos, duration, ProfilerTask.INFO, description));
+          multiLaneGenerator.release(prefix, lane);
         },
         MoreExecutors.directExecutor());
     return future;
   }
 
-  private static final long LANE_ID_BASE = 1_000_000;
-  private final AtomicLong nextLaneId = new AtomicLong(LANE_ID_BASE);
-  private final TaskTypeLaneIdGenerator laneIdGenerator = new TaskTypeLaneIdGenerator();
+  private final ThreadLocal<String> virtualThreadPrefix = ThreadLocal.withInitial(() -> null);
+  private final ThreadLocal<Lane> borrowedLane = ThreadLocal.withInitial(() -> null);
 
-  private class TaskTypeLaneIdGenerator {
-    private final Map<ProfilerTaskType, LaneIdGenerator> typeToLaneIdGenerator =
-        Maps.newConcurrentMap();
+  private void registerVirtualThread(String prefix) {
+    var thread = Thread.currentThread();
+    var threadId = thread.threadId();
+    virtualThreadPrefix.set(prefix);
+    thread.setName(prefix + threadId);
+  }
+
+  private void deregisterVirtualThread() {
+    var prefix = checkNotNull(virtualThreadPrefix.get());
+    virtualThreadPrefix.remove();
+    var lane = borrowedLane.get();
+    if (lane != null) {
+      borrowedLane.remove();
+      multiLaneGenerator.release(prefix, lane);
+    }
+  }
+
+  private final AtomicLong nextLaneId = new AtomicLong(1_000_000);
+  private final MultiLaneGenerator multiLaneGenerator = new MultiLaneGenerator();
+
+  private class MultiLaneGenerator {
+    private final Map<String, LaneGenerator> laneGenerators = Maps.newConcurrentMap();
 
-    public long acquire(ProfilerTaskType type) {
-      var laneIdGenerator =
-          typeToLaneIdGenerator.computeIfAbsent(type, unused -> new LaneIdGenerator(type));
-      return laneIdGenerator.acquire();
+    private Lane acquire(String prefix) {
+      checkState(isActive());
+      var laneGenerator =
+          laneGenerators.computeIfAbsent(prefix, unused -> new LaneGenerator(prefix));
+      return laneGenerator.acquire();
     }
 
-    public void release(ProfilerTaskType type, long laneId) {
-      var laneIdGenerator = checkNotNull(typeToLaneIdGenerator.get(type));
-      laneIdGenerator.release(laneId);
+    private void release(String prefix, Lane lane) {
+      checkState(isActive());
+      var laneGenerator = checkNotNull(laneGenerators.get(prefix));
+      laneGenerator.release(lane);
     }
   }
 
-  private class LaneIdGenerator {
-    private final ProfilerTaskType type;
-    private final PriorityQueue<Long> availableLaneIds = new PriorityQueue<>();
+  private record Lane(long id) implements Comparable<Lane> {
+    @Override
+    public int compareTo(Lane o) {
+      return Long.compare(id, o.id);
+    }
+  }
+
+  private class LaneGenerator {
+    private final String prefix;
+
+    @GuardedBy("this")
+    private final PriorityQueue<Lane> availableLanes = new PriorityQueue<>();
 
+    @GuardedBy("this")
     private int count = 0;
 
-    private LaneIdGenerator(ProfilerTaskType type) {
-      this.type = type;
+    private LaneGenerator(String prefix) {
+      this.prefix = prefix;
     }
 
-    public synchronized long acquire() {
-      if (!availableLaneIds.isEmpty()) {
-        return availableLaneIds.poll();
+    public Lane acquire() {
+      long newLaneId;
+      String newLaneName;
+      synchronized (this) {
+        if (!availableLanes.isEmpty()) {
+          return availableLanes.poll();
+        }
+
+        newLaneId = nextLaneId.getAndIncrement();
+        int newLaneIndex = count++;
+        newLaneName = prefix + newLaneIndex + " (Virtual)";
       }
-      var newLaneId = Profiler.this.nextLaneId.getAndIncrement();
-      var threadMetadata = new ThreadMetadata(type.getName(count++), newLaneId, LANE_ID_BASE);
+
+      var threadMetadata = new ThreadMetadata(newLaneName, newLaneId);
       var writer = Profiler.this.writerRef.get();
       if (writer != null) {
         writer.enqueue(threadMetadata);
       }
-      return newLaneId;
+      return new Lane(newLaneId);
     }
 
-    public synchronized void release(long laneId) {
-      availableLaneIds.add(laneId);
+    public synchronized void release(Lane lane) {
+      availableLanes.add(lane);
     }
   }
+
+  private long getLaneId() {
+    var currentThread = Thread.currentThread();
+    var threadId = currentThread.threadId();
+    if (!currentThread.isVirtual()) {
+      return threadId;
+    }
+
+    var lane = borrowedLane.get();
+    if (lane == null) {
+      var prefix = virtualThreadPrefix.get();
+      checkNotNull(
+          prefix,
+          "Current virtual thread is not registered. Did you use"
+              + " Profiler#profileableVirtualThreadFactor to create a VirtualThread?");
+      lane = multiLaneGenerator.acquire(prefix);
+      borrowedLane.set(lane);
+    }
+    return lane.id();
+  }
+
+  public ThreadFactory profileableVirtualThreadFactory(String prefix) {
+    return r ->
+        Thread.ofVirtual()
+            .unstarted(
+                () -> {
+                  var profiler = Profiler.instance();
+                  profiler.registerVirtualThread(prefix);
+                  try {
+                    r.run();
+                  } finally {
+                    profiler.deregisterVirtualThread();
+                  }
+                });
+  }
 }

src/main/java/com/google/devtools/build/lib/profiler/ThreadMetadata.java

@@ -24,6 +24,12 @@ class ThreadMetadata implements TraceData {
   private final long threadId;
   private final long sortIndex;
 
+  public ThreadMetadata(String readableName, long threadId) {
+    this.readableName = readableName;
+    this.threadId = threadId;
+    this.sortIndex = getSortIndex(readableName);
+  }
+
   public ThreadMetadata(String readableName, long threadId, long sortIndex) {
     this.readableName = readableName;
     this.threadId = threadId;

src/main/java/com/google/devtools/build/lib/runtime/QuiescingExecutorsImpl.java

@@ -25,6 +25,7 @@
 import com.google.devtools.build.lib.concurrent.QuiescingExecutor;
 import com.google.devtools.build.lib.concurrent.QuiescingExecutors;
 import com.google.devtools.build.lib.pkgcache.PackageOptions;
+import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.skyframe.ParallelEvaluatorErrorClassifier;
 import com.google.devtools.common.options.OptionsProvider;
 import java.util.concurrent.Executors;
@@ -145,7 +146,8 @@ public QuiescingExecutor getMergedAnalysisAndExecutionExecutor() {
             /* parallelism= */ cpuHeavySkyKeysThreadPoolSize, SKYFRAME_EVALUATOR_CPU_HEAVY),
         useAsyncExecution
             ? Executors.newThreadPerTaskExecutor(
-                Thread.ofVirtual().name(SKYFRAME_EVALUATOR_EXECUTION + "-%s", 0).factory())
+                Profiler.instance()
+                    .profileableVirtualThreadFactory(SKYFRAME_EVALUATOR_EXECUTION + "-"))
             : AbstractQueueVisitor.createExecutorService(
                 /* parallelism= */ executionParallelism, SKYFRAME_EVALUATOR_EXECUTION),
         ExceptionHandlingMode.FAIL_FAST,