src/main/java/com/google/devtools/build/lib/concurrent/PriorityWorkerPool.java - bazel - Git at Google

 // Copyright 2022 The Bazel Authors. All rights reserved.
 //
 // Licensed 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 com.google.devtools.build.lib.concurrent;

 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.devtools.build.lib.concurrent.PaddedAddresses.createPaddedBaseAddress;
 import static com.google.devtools.build.lib.concurrent.PaddedAddresses.getAlignedAddress;
 import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.DO_CPU_HEAVY_TASK;
 import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.DO_TASK;
 import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.IDLE;
 import static java.lang.Thread.currentThread;

 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.unsafe.UnsafeProvider;
 import java.lang.ref.Cleaner;
 import java.lang.ref.PhantomReference;
 import java.lang.ref.ReferenceQueue;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.ForkJoinWorkerThread;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 import sun.misc.Unsafe;

 /**
  * Inner implementation of {@link TieredPriorityExecutor}.
  *
  * <p>The main motivation for this additional layering is to facilitate garbage collection. The
  * {@link PriorityWorkerPool#WorkerThread}s have references to their enclosing {@link
  * PriorityWorkerPool}. Since threads are garbage collection roots, this makes the entire {@link
  * PriorityWorkerPool} ineligible for garbage collection.
  *
  * <p>The {@link PriorityWorkerPool} has no backreferences to the enclosing {@link
  * TieredPriorityExecutor}, so the {@link TieredPriorityExecutor} is eligible for garbage collection
  * and is able to perform cleanup tasks.
  */
 final class PriorityWorkerPool {
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

   /** The configured size of the thread pool. */
   private final int poolSize;

   /** The number of CPU permits configured. */
   private final int cpuPermits;

   private ForkJoinPool pool;

   /**
    * Queue for non-CPU heavy tasks.
    *
    * <p>An interesting alternative to consider is to place unprioritized tasks directly into {@link
    * #pool}, which could reduce the work performed by the system. Doing this results in about a
    * {@code 4%} end-to-end regression in our benchmark. The likely cause for this is that FIFO
    * behavior is very important for performance because it reflects the ordering of prioritized
    * tasks.
    */
   private final TaskFifo queue;

   private final ConcurrentSkipListSet<ComparableRunnable> cpuHeavyQueue =
       new ConcurrentSkipListSet<>();

   private final String name;

   /**
    * Cache of workers for interrupt handling.
    *
    * <p>A {@link Cache} allows us to use weak keys, which are the only relevant objects here. The
    * values are an unintentional side effect of the library and populated arbitrarily.
    */
   private final Cache<WorkerThread, Object> workers = Caffeine.newBuilder().weakKeys().build();

   /**
    * A synchronization mechanism used when waiting for quiescence.
    *
    * <p>Quiescence can be reached if either of the following conditions is satisfied.
    *
    * <ul>
    *   <li>No tasks are enqueued and all workers are idle.
    *   <li>There was a catastrophe (see {@link TieredPriorityExecutor#awaitQuiescence}).
    * </ul>
    */
   private final Object quiescenceMonitor = new Object();

   private final ErrorClassifier errorClassifier;

   /**
    * The most severe unhandled exception thrown by a worker thread, according to {@link
    * #errorClassifier}. This exception gets propagated to the calling thread of {@link
    * TieredPriorityExecutor#awaitQuiescence}. We use the most severe error for the sake of not
    * masking e.g. crashes in worker threads after the first critical error that can occur due to
    * race conditions in client code.
    *
    * <p>If {@link AbstractQueueVisitor} clients don't like the semantics of storing and propagating
    * the most severe error, then they should be provide an {@link ErrorClassifier} that does the
    * right thing (e.g. to cause the _first_ error to be propagated, you'd want to provide an {@link
    * ErrorClassifier} that gives all errors the exact same {@link
    * ErrorClassifier.ErrorClassification}).
    *
    * <p>Note that this is not a performance-critical path.
    */
   private final AtomicReference<Throwable> unhandled = new AtomicReference<>();

   PriorityWorkerPool(
       Cleaner cleaner, String name, int poolSize, int cpuPermits, ErrorClassifier errorClassifier) {
     checkArgument(poolSize <= (THREADS_MASK >> THREADS_BIT_OFFSET), poolSize);
     checkArgument(cpuPermits <= (CPU_PERMITS_MASK >> CPU_PERMITS_BIT_OFFSET), cpuPermits);

     this.name = name;
     this.poolSize = poolSize;
     this.cpuPermits = cpuPermits;

     this.pool = newForkJoinPool();
     this.errorClassifier = errorClassifier;

     long baseAddress = createPaddedBaseAddress(5);
     cleaner.register(this, new AddressFreer(baseAddress));
     this.countersAddress = getAlignedAddress(baseAddress, /* offset= */ 0);

     this.queue =
         new TaskFifo(
             /* sizeAddress= */ getAlignedAddress(baseAddress, /* offset= */ 1),
             /* appendIndexAddress= */ getAlignedAddress(baseAddress, /* offset= */ 2),
             /* takeIndexAddress= */ getAlignedAddress(baseAddress, /* offset= */ 3));

     this.activeWorkerCountAddress = getAlignedAddress(baseAddress, /* offset= */ 4);

     resetExecutionCounters();
   }

   int poolSize() {
     return poolSize;
   }

   int cpuPermits() {
     return cpuPermits;
   }

   void execute(Runnable rawTask) {
     if (rawTask instanceof ComparableRunnable) {
       var task = (ComparableRunnable) rawTask;
       if (task.isCpuHeavy()) {
         cpuHeavyQueue.add(task);
         if (acquireThreadAndCpuPermitElseReleaseCpuHeavyTask()) {
           UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
           pool.execute(RUN_CPU_HEAVY_TASK);
         }
         return;
       }
     }

     if (!queue.tryAppend(rawTask)) {
       if (!isCancelled()) {
         // The task queue is full (and the pool is not cancelled). Enqueues the task directly in the
         // ForkJoinPool. This should be rare in practice.
         UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
         pool.execute(
             () -> {
               try {
                 rawTask.run();
               } catch (Throwable uncaught) {
                 handleUncaughtError(uncaught);
               } finally {
                 workerBecomingIdle();
               }
             });
       }
       return;
     }

     if (acquireThreadElseReleaseTask()) {
       UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
       pool.execute(RUN_TASK);
     }
   }

   /**
    * An object to {@link Object#wait} or {@link Object#notifyAll} on for quiescence.
    *
    * <p>The pool will {@link Object#notifyAll} this object when {@link #isQuiescent} becomes true.
    */
   Object quiescenceMonitor() {
     return quiescenceMonitor;
   }

   @Nullable
   Throwable unhandled() {
     return unhandled.get();
   }

   /**
    * Sets the pool to stop processing tasks and interrupts all workers.
    *
    * <p>Calling cancel is on a cancelled pool is a noop. A cancelled pool can be {@link #reset}.
    */
   void cancel() {
     markCancelled();

     workers.asMap().keySet().forEach(Thread::interrupt);
   }

   /** Shuts down a pool and frees all resources. */
   private void cleanup() {
     checkState(isQuiescent(), "cleanup called on pool that was not quiescent: %s", this);

     // There's no particular significance to the teardown order here, given that the pool is
     // quiescent. It has the appearance of a logical ordering for cosmetic reasons only.
     queue.clear();
     cpuHeavyQueue.clear();
     workers.invalidateAll();
     pool.shutdown();
   }

   /** Cleans up then resets this pool. */
   void reset() {
     cleanup();
     unhandled.set(null);

     resetExecutionCounters();
     pool = newForkJoinPool();
   }

   /**
    * Makes this pool eligible for garbage collection.
    *
    * <p>Intended for registration with {@link java.lang.ref.Cleaner}.
    */
   void dispose() {
     cancel();
     synchronized (quiescenceMonitor) {
       while (!isQuiescent()) {
         // This should only be reachable if there was a catastrophe because otherwise
         // `TieredPriorityExecutor.awaitQuiescence` would have already ensured quiescence. It's
         // not clear that Bazel can recover from this state. The appropriate action, nonetheless, is
         // to wait for quiescence, then shutdown the pool.
         try {
           quiescenceMonitor.wait();
         } catch (InterruptedException e) {
           // We don't expect this to ever happen, given this is running on a cleaner thread. Logs a
           // warning in case it somehow happens.
           logger.atWarning().withCause(e).log("%s interrupted while cleaning up.", this);
         }
       }
     }
     cleanup();
   }

   @VisibleForTesting
   PhantomReference<ForkJoinPool> registerPoolDisposalMonitorForTesting(
       ReferenceQueue<ForkJoinPool> referenceQueue) {
     return new PhantomReference<>(pool, referenceQueue);
   }

   @Override
   public String toString() {
     var threadStates = new TreeMap<Thread.State, Integer>();
     for (var w : workers.asMap().keySet()) {
       threadStates.compute(w.getState(), (k, v) -> v == null ? 1 : (v + 1));
     }
     return toStringHelper(this)
         .add("available", formatSnapshot(getExecutionCounters()))
         .add("|queue|", queue.size())
         .add("|cpu queue|", cpuHeavyQueue.size())
         .add("threads", threadStates)
         .add("unhandled", unhandled.get())
         .add("pool", pool)
         .toString();
   }

   /**
    * Handles errors created by submitted tasks.
    *
    * <p>Behavior adheres to documentation of {@link TieredPriorityExecutor#awaitQuiescence}.
    */
   private void handleUncaughtError(Throwable error) {
     boolean critical = false;
     var classification = errorClassifier.classify(error);
     switch (classification) {
       case AS_CRITICAL_AS_POSSIBLE:
       case CRITICAL_AND_LOG:
         logger.atWarning().withCause(error).log("Found critical error in queue visitor");
         // fall through
       case CRITICAL:
         critical = true;
         break;
       case NOT_CRITICAL:
         break;
     }

     Throwable unhandledSnapshot;
     do {
       unhandledSnapshot = unhandled.get();
       if (unhandledSnapshot != null
           && errorClassifier.classify(unhandledSnapshot).compareTo(classification) >= 0) {
         break; // Skips saving anything less severe.
       }
     } while (!unhandled.compareAndSet(unhandledSnapshot, error));

     if (critical) {
       cancel();
     }
   }

   private ForkJoinPool newForkJoinPool() {
     return new ForkJoinPool(
         poolSize,
         pool -> {
           var worker = new WorkerThread(pool, name);
           workers.put(worker, "A non-null value, as required by Caffeine.");
           return worker;
         },
         /* handler= */ null,
         /* asyncMode= */ false);
   }

   /**
    * {@link WorkerThread#runLoop} implements a small state machine.
    *
    * <p>After completing a task, the worker checks if there are any available tasks that it may
    * execute, subject to CPU permit constraints. On finding and reserving an appropriate task, the
    * worker returns its next planned activity, {@link #IDLE} if it finds nothing to do.
    */
   enum NextWorkerActivity {
     /** The worker will stop. */
     IDLE,
     /** The worker will perform a non-CPU heavy task. */
     DO_TASK,
     /** The worker will perform a CPU heavy task. */
     DO_CPU_HEAVY_TASK
   }

   /**
    * Performs a task in a {@link WorkerThread} then loops.
    *
    * <p>Passed to {@link ForkJoinPool#execute} when a non-CPU-heavy task execution is needed.
    */
   // This could be a static method reference, but this makes it absolutely clear that no per-task
   // garbage is generated and results in nicer stack traces.
   private static final Runnable RUN_TASK = new LoopStarter(DO_TASK);

   /**
    * Performs a CPU heavy task in a {@link WorkerThread} then loops.
    *
    * <p>Passed to {@link ForkJoinPool#execute} when a CPU-heavy task execution is needed.
    */
   private static final Runnable RUN_CPU_HEAVY_TASK = new LoopStarter(DO_CPU_HEAVY_TASK);

   private static class LoopStarter implements Runnable {
     private final NextWorkerActivity activity;

     private LoopStarter(NextWorkerActivity activity) {
       this.activity = activity;
     }

     @Override
     public void run() {
       ((WorkerThread) currentThread()).runLoop(activity);
     }
   }

   class WorkerThread extends ForkJoinWorkerThread {

     private WorkerThread(ForkJoinPool pool, String name) {
       super(pool);
       setName(name + "-" + getPoolIndex());
     }

     /**
      * The worker runs a loop that scans for and runs available tasks.
      *
      * <p>This reduces costs associated with stopping and restarting threads.
      */
     private void runLoop(NextWorkerActivity nextActivity) {
       while (true) {
         switch (nextActivity) {
           case IDLE:
             workerBecomingIdle();
             return;
           case DO_TASK:
             dequeueTaskAndRun();
             nextActivity = getActivityFollowingTask();
             break;
           case DO_CPU_HEAVY_TASK:
             dequeueCpuHeavyTaskAndRun();
             nextActivity = getActivityFollowingCpuHeavyTask();
             break;
         }
       }
     }
   }

   private void dequeueTaskAndRun() {
     try {
       var task = queue.take();
       task.run();
     } catch (Throwable uncaught) {
       handleUncaughtError(uncaught);
     }
   }

   private void dequeueCpuHeavyTaskAndRun() {
     try {
       cpuHeavyQueue.pollFirst().run();
     } catch (Throwable uncaught) {
       handleUncaughtError(uncaught);
     }
   }

   private void workerBecomingIdle() {
     if (UNSAFE.getAndAddInt(null, activeWorkerCountAddress, -1) == 1) {
       synchronized (quiescenceMonitor) {
         quiescenceMonitor.notifyAll();
       }
     }
   }

   // The constants below apply to the 64-bit execution counters value.

   private static final long CANCEL_BIT = 0x8000_0000_0000_0000L;

   private static final long CPU_PERMITS_MASK = 0x7FF0_0000_0000_0000L;
   private static final int CPU_PERMITS_BIT_OFFSET = 52;
   private static final long ONE_CPU_PERMIT = 1L << CPU_PERMITS_BIT_OFFSET;

   private static final long THREADS_MASK = 0x000F_FFE0_0000_0000L;
   private static final int THREADS_BIT_OFFSET = 37;
   private static final long ONE_THREAD = 1L << THREADS_BIT_OFFSET;

   private static final long TASKS_MASK = 0x0000_001F_FF80_0000L;
   private static final int TASKS_BIT_OFFSET = 23;
   private static final long ONE_TASK = 1L << TASKS_BIT_OFFSET;
   static final int TASKS_MAX_VALUE = (int) (TASKS_MASK >> TASKS_BIT_OFFSET);

   private static final long CPU_HEAVY_TASKS_MASK = 0x0000_0000_07F_FFFFL;
   private static final int CPU_HEAVY_TASKS_BIT_OFFSET = 0;
   private static final long ONE_CPU_HEAVY_TASK = 1L << CPU_HEAVY_TASKS_BIT_OFFSET;

   private static final long CPU_HEAVY_RESOURCES = ONE_CPU_PERMIT + ONE_THREAD;

   static {
     checkState(
         ONE_CPU_PERMIT == (CPU_PERMITS_MASK & -CPU_PERMITS_MASK),
         "Inconsistent CPU Permits Constants");
     checkState(ONE_THREAD == (THREADS_MASK & -THREADS_MASK), "Inconistent Threads Constants");
     checkState(
         ONE_CPU_HEAVY_TASK == (CPU_HEAVY_TASKS_MASK & -CPU_HEAVY_TASKS_MASK),
         "Inconsistent CPU Heavy Task Constants");
   }

   /**
    * Address of the execution counters value, consisting of 5 fields packed into a 64-bit long.
    *
    * <ol>
    *   <li>Canceled - (1 bit) true for cancelled.
    *   <li>CPU Permits - (11 bits) how many CPU heavy permits are available.
    *   <li>Threads - (15 bits) how many threads are available.
    *   <li>Tasks - (14 bits) how many non-CPU heavy tasks are inflight.
    *   <li>CPU Heavy Tasks - (23 bits) how many CPU heavy tasks are inflight.
    * </ol>
    *
    * <p>Convenience constants for field access and manipulation are above.
    */
   private final long countersAddress;

   private final long activeWorkerCountAddress;

   boolean isQuiescent() {
     return UNSAFE.getInt(null, activeWorkerCountAddress) == 0;
   }

   boolean isCancelled() {
     return getExecutionCounters() < 0;
   }

   private void markCancelled() {
     long snapshot;
     do {
       snapshot = getExecutionCounters();
       if (snapshot < 0) {
         return; // Already cancelled.
       }
     } while (!tryUpdateExecutionCounters(snapshot, snapshot | CANCEL_BIT));
   }

   private void resetExecutionCounters() {
     UNSAFE.putLong(
         null,
         countersAddress,
         (((long) poolSize) << THREADS_BIT_OFFSET)
             | (((long) cpuPermits) << CPU_PERMITS_BIT_OFFSET));
     UNSAFE.putInt(null, activeWorkerCountAddress, 0);
   }

   private boolean acquireThreadElseReleaseTask() {
     long snapshot;
     do {
       snapshot = UNSAFE.getLongVolatile(null, countersAddress);
       boolean acquired = (snapshot & THREADS_MASK) > 0 && snapshot >= 0;
       long target = snapshot + (acquired ? -ONE_THREAD : ONE_TASK);
       if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
         return acquired;
       }
     } while (true);
   }

   private boolean acquireThreadAndCpuPermitElseReleaseCpuHeavyTask() {
     long snapshot;
     do {
       snapshot = UNSAFE.getLongVolatile(null, countersAddress);
       boolean acquired =
           (snapshot & (CANCEL_BIT | CPU_PERMITS_MASK)) > 0 && (snapshot & THREADS_MASK) > 0;
       long target = snapshot + (acquired ? -(ONE_THREAD + ONE_CPU_PERMIT) : ONE_CPU_HEAVY_TASK);
       if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
         return acquired;
       }
     } while (true);
   }

   /**
    * Worker threads determine their next action after completing a task using this method.
    *
    * <p>This acquires a CPU permit when returning {@link NextWorkerActivity#DO_CPU_HEAVY_TASK}.
    */
   private NextWorkerActivity getActivityFollowingTask() {
     long snapshot = UNSAFE.getLongVolatile(null, countersAddress);
     do {
       if ((snapshot & (CANCEL_BIT | TASKS_MASK)) > 0) {
         if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, snapshot - ONE_TASK)) {
           return DO_TASK;
         }
       } else if ((snapshot & (CANCEL_BIT | CPU_HEAVY_TASKS_MASK)) > 0
           && (snapshot & CPU_PERMITS_MASK) != 0) {
         if (UNSAFE.compareAndSwapLong(
             null, countersAddress, snapshot, snapshot - (ONE_CPU_HEAVY_TASK + ONE_CPU_PERMIT))) {
           return DO_CPU_HEAVY_TASK;
         }
       } else {
         long target = snapshot + ONE_THREAD;
         if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
           return IDLE;
         }
       }
       snapshot = UNSAFE.getLong(null, countersAddress);
     } while (true);
   }

   /**
    * Worker threads call this to determine their next action after completing a CPU heavy task.
    *
    * <p>This releases a CPU permit when returning {@link NextWorkerActivity#IDLE} or {@link
    * NextWorkerActivity#DO_TASK}.
    */
   private NextWorkerActivity getActivityFollowingCpuHeavyTask() {
     long snapshot = UNSAFE.getLongVolatile(null, countersAddress);
     do {
       if ((snapshot & (CANCEL_BIT | TASKS_MASK)) > 0) {
         if (UNSAFE.compareAndSwapLong(
             null, countersAddress, snapshot, snapshot + (ONE_CPU_PERMIT - ONE_TASK))) {
           return DO_TASK;
         }
       } else if ((snapshot & (CANCEL_BIT | CPU_HEAVY_TASKS_MASK)) > 0) {
         if (UNSAFE.compareAndSwapLong(
             null, countersAddress, snapshot, snapshot - ONE_CPU_HEAVY_TASK)) {
           return DO_CPU_HEAVY_TASK;
         }
       } else {
         long target = snapshot + CPU_HEAVY_RESOURCES;
         if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
           return IDLE;
         }
       }
       snapshot = UNSAFE.getLong(null, countersAddress);
     } while (true);
   }

   // Throughout this class, the following wrappers are used where possible, but they are often not
   // inlined by the JVM even though they show up on profiles, so they are inlined explicitly in
   // numerous cases.

   private long getExecutionCounters() {
     return UNSAFE.getLongVolatile(null, countersAddress);
   }

   private boolean tryUpdateExecutionCounters(long snapshot, long target) {
     return UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target);
   }

   private static String formatSnapshot(long snapshot) {
     return String.format(
         "{cancelled=%b, threads=%d, cpuPermits=%d, tasks=%d, cpuHeavyTasks=%d}",
         snapshot < 0,
         (snapshot & THREADS_MASK) >> THREADS_BIT_OFFSET,
         (snapshot & CPU_PERMITS_MASK) >> CPU_PERMITS_BIT_OFFSET,
         (snapshot & TASKS_MASK) >> TASKS_BIT_OFFSET,
         (snapshot & CPU_HEAVY_TASKS_MASK) >> CPU_HEAVY_TASKS_BIT_OFFSET);
   }

   private static final Unsafe UNSAFE = UnsafeProvider.unsafe();
 }
	// Copyright 2022 The Bazel Authors. All rights reserved.
	//
	// Licensed 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 com.google.devtools.build.lib.concurrent;

	import static com.google.common.base.MoreObjects.toStringHelper;
	import static com.google.common.base.Preconditions.checkArgument;
	import static com.google.common.base.Preconditions.checkState;
	import static com.google.devtools.build.lib.concurrent.PaddedAddresses.createPaddedBaseAddress;
	import static com.google.devtools.build.lib.concurrent.PaddedAddresses.getAlignedAddress;
	import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.DO_CPU_HEAVY_TASK;
	import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.DO_TASK;
	import static com.google.devtools.build.lib.concurrent.PriorityWorkerPool.NextWorkerActivity.IDLE;
	import static java.lang.Thread.currentThread;

	import com.github.benmanes.caffeine.cache.Cache;
	import com.github.benmanes.caffeine.cache.Caffeine;
	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.flogger.GoogleLogger;
	import com.google.devtools.build.lib.unsafe.UnsafeProvider;
	import java.lang.ref.Cleaner;
	import java.lang.ref.PhantomReference;
	import java.lang.ref.ReferenceQueue;
	import java.util.TreeMap;
	import java.util.concurrent.ConcurrentSkipListSet;
	import java.util.concurrent.ForkJoinPool;
	import java.util.concurrent.ForkJoinWorkerThread;
	import java.util.concurrent.atomic.AtomicReference;
	import javax.annotation.Nullable;
	import sun.misc.Unsafe;

	/**
	* Inner implementation of {@link TieredPriorityExecutor}.
	*
	* <p>The main motivation for this additional layering is to facilitate garbage collection. The
	* {@link PriorityWorkerPool#WorkerThread}s have references to their enclosing {@link
	* PriorityWorkerPool}. Since threads are garbage collection roots, this makes the entire {@link
	* PriorityWorkerPool} ineligible for garbage collection.
	*
	* <p>The {@link PriorityWorkerPool} has no backreferences to the enclosing {@link
	* TieredPriorityExecutor}, so the {@link TieredPriorityExecutor} is eligible for garbage collection
	* and is able to perform cleanup tasks.
	*/
	final class PriorityWorkerPool {
	private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

	/** The configured size of the thread pool. */
	private final int poolSize;

	/** The number of CPU permits configured. */
	private final int cpuPermits;

	private ForkJoinPool pool;

	/**
	* Queue for non-CPU heavy tasks.
	*
	* <p>An interesting alternative to consider is to place unprioritized tasks directly into {@link
	* #pool}, which could reduce the work performed by the system. Doing this results in about a
	* {@code 4%} end-to-end regression in our benchmark. The likely cause for this is that FIFO
	* behavior is very important for performance because it reflects the ordering of prioritized
	* tasks.
	*/
	private final TaskFifo queue;

	private final ConcurrentSkipListSet<ComparableRunnable> cpuHeavyQueue =
	new ConcurrentSkipListSet<>();

	private final String name;

	/**
	* Cache of workers for interrupt handling.
	*
	* <p>A {@link Cache} allows us to use weak keys, which are the only relevant objects here. The
	* values are an unintentional side effect of the library and populated arbitrarily.
	*/
	private final Cache<WorkerThread, Object> workers = Caffeine.newBuilder().weakKeys().build();

	/**
	* A synchronization mechanism used when waiting for quiescence.
	*
	* <p>Quiescence can be reached if either of the following conditions is satisfied.
	*
	* <ul>
	* <li>No tasks are enqueued and all workers are idle.
	* <li>There was a catastrophe (see {@link TieredPriorityExecutor#awaitQuiescence}).
	* </ul>
	*/
	private final Object quiescenceMonitor = new Object();

	private final ErrorClassifier errorClassifier;

	/**
	* The most severe unhandled exception thrown by a worker thread, according to {@link
	* #errorClassifier}. This exception gets propagated to the calling thread of {@link
	* TieredPriorityExecutor#awaitQuiescence}. We use the most severe error for the sake of not
	* masking e.g. crashes in worker threads after the first critical error that can occur due to
	* race conditions in client code.
	*
	* <p>If {@link AbstractQueueVisitor} clients don't like the semantics of storing and propagating
	* the most severe error, then they should be provide an {@link ErrorClassifier} that does the
	* right thing (e.g. to cause the _first_ error to be propagated, you'd want to provide an {@link
	* ErrorClassifier} that gives all errors the exact same {@link
	* ErrorClassifier.ErrorClassification}).
	*
	* <p>Note that this is not a performance-critical path.
	*/
	private final AtomicReference<Throwable> unhandled = new AtomicReference<>();

	PriorityWorkerPool(
	Cleaner cleaner, String name, int poolSize, int cpuPermits, ErrorClassifier errorClassifier) {
	checkArgument(poolSize <= (THREADS_MASK >> THREADS_BIT_OFFSET), poolSize);
	checkArgument(cpuPermits <= (CPU_PERMITS_MASK >> CPU_PERMITS_BIT_OFFSET), cpuPermits);

	this.name = name;
	this.poolSize = poolSize;
	this.cpuPermits = cpuPermits;

	this.pool = newForkJoinPool();
	this.errorClassifier = errorClassifier;

	long baseAddress = createPaddedBaseAddress(5);
	cleaner.register(this, new AddressFreer(baseAddress));
	this.countersAddress = getAlignedAddress(baseAddress, /* offset= */ 0);

	this.queue =
	new TaskFifo(
	/* sizeAddress= / getAlignedAddress(baseAddress, / offset= */ 1),
	/* appendIndexAddress= / getAlignedAddress(baseAddress, / offset= */ 2),
	/* takeIndexAddress= / getAlignedAddress(baseAddress, / offset= */ 3));

	this.activeWorkerCountAddress = getAlignedAddress(baseAddress, /* offset= */ 4);

	resetExecutionCounters();
	}

	int poolSize() {
	return poolSize;
	}

	int cpuPermits() {
	return cpuPermits;
	}

	void execute(Runnable rawTask) {
	if (rawTask instanceof ComparableRunnable) {
	var task = (ComparableRunnable) rawTask;
	if (task.isCpuHeavy()) {
	cpuHeavyQueue.add(task);
	if (acquireThreadAndCpuPermitElseReleaseCpuHeavyTask()) {
	UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
	pool.execute(RUN_CPU_HEAVY_TASK);
	}
	return;
	}
	}

	if (!queue.tryAppend(rawTask)) {
	if (!isCancelled()) {
	// The task queue is full (and the pool is not cancelled). Enqueues the task directly in the
	// ForkJoinPool. This should be rare in practice.
	UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
	pool.execute(
	() -> {
	try {
	rawTask.run();
	} catch (Throwable uncaught) {
	handleUncaughtError(uncaught);
	} finally {
	workerBecomingIdle();
	}
	});
	}
	return;
	}

	if (acquireThreadElseReleaseTask()) {
	UNSAFE.getAndAddInt(null, activeWorkerCountAddress, 1);
	pool.execute(RUN_TASK);
	}
	}

	/**
	* An object to {@link Object#wait} or {@link Object#notifyAll} on for quiescence.
	*
	* <p>The pool will {@link Object#notifyAll} this object when {@link #isQuiescent} becomes true.
	*/
	Object quiescenceMonitor() {
	return quiescenceMonitor;
	}

	@Nullable
	Throwable unhandled() {
	return unhandled.get();
	}

	/**
	* Sets the pool to stop processing tasks and interrupts all workers.
	*
	* <p>Calling cancel is on a cancelled pool is a noop. A cancelled pool can be {@link #reset}.
	*/
	void cancel() {
	markCancelled();

	workers.asMap().keySet().forEach(Thread::interrupt);
	}

	/** Shuts down a pool and frees all resources. */
	private void cleanup() {
	checkState(isQuiescent(), "cleanup called on pool that was not quiescent: %s", this);

	// There's no particular significance to the teardown order here, given that the pool is
	// quiescent. It has the appearance of a logical ordering for cosmetic reasons only.
	queue.clear();
	cpuHeavyQueue.clear();
	workers.invalidateAll();
	pool.shutdown();
	}

	/** Cleans up then resets this pool. */
	void reset() {
	cleanup();
	unhandled.set(null);

	resetExecutionCounters();
	pool = newForkJoinPool();
	}

	/**
	* Makes this pool eligible for garbage collection.
	*
	* <p>Intended for registration with {@link java.lang.ref.Cleaner}.
	*/
	void dispose() {
	cancel();
	synchronized (quiescenceMonitor) {
	while (!isQuiescent()) {
	// This should only be reachable if there was a catastrophe because otherwise
	// `TieredPriorityExecutor.awaitQuiescence` would have already ensured quiescence. It's
	// not clear that Bazel can recover from this state. The appropriate action, nonetheless, is
	// to wait for quiescence, then shutdown the pool.
	try {
	quiescenceMonitor.wait();
	} catch (InterruptedException e) {
	// We don't expect this to ever happen, given this is running on a cleaner thread. Logs a
	// warning in case it somehow happens.
	logger.atWarning().withCause(e).log("%s interrupted while cleaning up.", this);
	}
	}
	}
	cleanup();
	}

	@VisibleForTesting
	PhantomReference<ForkJoinPool> registerPoolDisposalMonitorForTesting(
	ReferenceQueue<ForkJoinPool> referenceQueue) {
	return new PhantomReference<>(pool, referenceQueue);
	}

	@Override
	public String toString() {
	var threadStates = new TreeMap<Thread.State, Integer>();
	for (var w : workers.asMap().keySet()) {
	threadStates.compute(w.getState(), (k, v) -> v == null ? 1 : (v + 1));
	}
	return toStringHelper(this)
	.add("available", formatSnapshot(getExecutionCounters()))
	.add("\|queue\|", queue.size())
	.add("\|cpu queue\|", cpuHeavyQueue.size())
	.add("threads", threadStates)
	.add("unhandled", unhandled.get())
	.add("pool", pool)
	.toString();
	}

	/**
	* Handles errors created by submitted tasks.
	*
	* <p>Behavior adheres to documentation of {@link TieredPriorityExecutor#awaitQuiescence}.
	*/
	private void handleUncaughtError(Throwable error) {
	boolean critical = false;
	var classification = errorClassifier.classify(error);
	switch (classification) {
	case AS_CRITICAL_AS_POSSIBLE:
	case CRITICAL_AND_LOG:
	logger.atWarning().withCause(error).log("Found critical error in queue visitor");
	// fall through
	case CRITICAL:
	critical = true;
	break;
	case NOT_CRITICAL:
	break;
	}

	Throwable unhandledSnapshot;
	do {
	unhandledSnapshot = unhandled.get();
	if (unhandledSnapshot != null
	&& errorClassifier.classify(unhandledSnapshot).compareTo(classification) >= 0) {
	break; // Skips saving anything less severe.
	}
	} while (!unhandled.compareAndSet(unhandledSnapshot, error));

	if (critical) {
	cancel();
	}
	}

	private ForkJoinPool newForkJoinPool() {
	return new ForkJoinPool(
	poolSize,
	pool -> {
	var worker = new WorkerThread(pool, name);
	workers.put(worker, "A non-null value, as required by Caffeine.");
	return worker;
	},
	/* handler= */ null,
	/* asyncMode= */ false);
	}

	/**
	* {@link WorkerThread#runLoop} implements a small state machine.
	*
	* <p>After completing a task, the worker checks if there are any available tasks that it may
	* execute, subject to CPU permit constraints. On finding and reserving an appropriate task, the
	* worker returns its next planned activity, {@link #IDLE} if it finds nothing to do.
	*/
	enum NextWorkerActivity {
	/** The worker will stop. */
	IDLE,
	/** The worker will perform a non-CPU heavy task. */
	DO_TASK,
	/** The worker will perform a CPU heavy task. */
	DO_CPU_HEAVY_TASK
	}

	/**
	* Performs a task in a {@link WorkerThread} then loops.
	*
	* <p>Passed to {@link ForkJoinPool#execute} when a non-CPU-heavy task execution is needed.
	*/
	// This could be a static method reference, but this makes it absolutely clear that no per-task
	// garbage is generated and results in nicer stack traces.
	private static final Runnable RUN_TASK = new LoopStarter(DO_TASK);

	/**
	* Performs a CPU heavy task in a {@link WorkerThread} then loops.
	*
	* <p>Passed to {@link ForkJoinPool#execute} when a CPU-heavy task execution is needed.
	*/
	private static final Runnable RUN_CPU_HEAVY_TASK = new LoopStarter(DO_CPU_HEAVY_TASK);

	private static class LoopStarter implements Runnable {
	private final NextWorkerActivity activity;

	private LoopStarter(NextWorkerActivity activity) {
	this.activity = activity;
	}

	@Override
	public void run() {
	((WorkerThread) currentThread()).runLoop(activity);
	}
	}

	class WorkerThread extends ForkJoinWorkerThread {

	private WorkerThread(ForkJoinPool pool, String name) {
	super(pool);
	setName(name + "-" + getPoolIndex());
	}

	/**
	* The worker runs a loop that scans for and runs available tasks.
	*
	* <p>This reduces costs associated with stopping and restarting threads.
	*/
	private void runLoop(NextWorkerActivity nextActivity) {
	while (true) {
	switch (nextActivity) {
	case IDLE:
	workerBecomingIdle();
	return;
	case DO_TASK:
	dequeueTaskAndRun();
	nextActivity = getActivityFollowingTask();
	break;
	case DO_CPU_HEAVY_TASK:
	dequeueCpuHeavyTaskAndRun();
	nextActivity = getActivityFollowingCpuHeavyTask();
	break;
	}
	}
	}
	}

	private void dequeueTaskAndRun() {
	try {
	var task = queue.take();
	task.run();
	} catch (Throwable uncaught) {
	handleUncaughtError(uncaught);
	}
	}

	private void dequeueCpuHeavyTaskAndRun() {
	try {
	cpuHeavyQueue.pollFirst().run();
	} catch (Throwable uncaught) {
	handleUncaughtError(uncaught);
	}
	}

	private void workerBecomingIdle() {
	if (UNSAFE.getAndAddInt(null, activeWorkerCountAddress, -1) == 1) {
	synchronized (quiescenceMonitor) {
	quiescenceMonitor.notifyAll();
	}
	}
	}

	// The constants below apply to the 64-bit execution counters value.

	private static final long CANCEL_BIT = 0x8000_0000_0000_0000L;

	private static final long CPU_PERMITS_MASK = 0x7FF0_0000_0000_0000L;
	private static final int CPU_PERMITS_BIT_OFFSET = 52;
	private static final long ONE_CPU_PERMIT = 1L << CPU_PERMITS_BIT_OFFSET;

	private static final long THREADS_MASK = 0x000F_FFE0_0000_0000L;
	private static final int THREADS_BIT_OFFSET = 37;
	private static final long ONE_THREAD = 1L << THREADS_BIT_OFFSET;

	private static final long TASKS_MASK = 0x0000_001F_FF80_0000L;
	private static final int TASKS_BIT_OFFSET = 23;
	private static final long ONE_TASK = 1L << TASKS_BIT_OFFSET;
	static final int TASKS_MAX_VALUE = (int) (TASKS_MASK >> TASKS_BIT_OFFSET);

	private static final long CPU_HEAVY_TASKS_MASK = 0x0000_0000_07F_FFFFL;
	private static final int CPU_HEAVY_TASKS_BIT_OFFSET = 0;
	private static final long ONE_CPU_HEAVY_TASK = 1L << CPU_HEAVY_TASKS_BIT_OFFSET;

	private static final long CPU_HEAVY_RESOURCES = ONE_CPU_PERMIT + ONE_THREAD;

	static {
	checkState(
	ONE_CPU_PERMIT == (CPU_PERMITS_MASK & -CPU_PERMITS_MASK),
	"Inconsistent CPU Permits Constants");
	checkState(ONE_THREAD == (THREADS_MASK & -THREADS_MASK), "Inconistent Threads Constants");
	checkState(
	ONE_CPU_HEAVY_TASK == (CPU_HEAVY_TASKS_MASK & -CPU_HEAVY_TASKS_MASK),
	"Inconsistent CPU Heavy Task Constants");
	}

	/**
	* Address of the execution counters value, consisting of 5 fields packed into a 64-bit long.
	*
	* <ol>
	* <li>Canceled - (1 bit) true for cancelled.
	* <li>CPU Permits - (11 bits) how many CPU heavy permits are available.
	* <li>Threads - (15 bits) how many threads are available.
	* <li>Tasks - (14 bits) how many non-CPU heavy tasks are inflight.
	* <li>CPU Heavy Tasks - (23 bits) how many CPU heavy tasks are inflight.
	* </ol>
	*
	* <p>Convenience constants for field access and manipulation are above.
	*/
	private final long countersAddress;

	private final long activeWorkerCountAddress;

	boolean isQuiescent() {
	return UNSAFE.getInt(null, activeWorkerCountAddress) == 0;
	}

	boolean isCancelled() {
	return getExecutionCounters() < 0;
	}

	private void markCancelled() {
	long snapshot;
	do {
	snapshot = getExecutionCounters();
	if (snapshot < 0) {
	return; // Already cancelled.
	}
	} while (!tryUpdateExecutionCounters(snapshot, snapshot \| CANCEL_BIT));
	}

	private void resetExecutionCounters() {
	UNSAFE.putLong(
	null,
	countersAddress,
	(((long) poolSize) << THREADS_BIT_OFFSET)
	\| (((long) cpuPermits) << CPU_PERMITS_BIT_OFFSET));
	UNSAFE.putInt(null, activeWorkerCountAddress, 0);
	}

	private boolean acquireThreadElseReleaseTask() {
	long snapshot;
	do {
	snapshot = UNSAFE.getLongVolatile(null, countersAddress);
	boolean acquired = (snapshot & THREADS_MASK) > 0 && snapshot >= 0;
	long target = snapshot + (acquired ? -ONE_THREAD : ONE_TASK);
	if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
	return acquired;
	}
	} while (true);
	}

	private boolean acquireThreadAndCpuPermitElseReleaseCpuHeavyTask() {
	long snapshot;
	do {
	snapshot = UNSAFE.getLongVolatile(null, countersAddress);
	boolean acquired =
	(snapshot & (CANCEL_BIT \| CPU_PERMITS_MASK)) > 0 && (snapshot & THREADS_MASK) > 0;
	long target = snapshot + (acquired ? -(ONE_THREAD + ONE_CPU_PERMIT) : ONE_CPU_HEAVY_TASK);
	if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
	return acquired;
	}
	} while (true);
	}

	/**
	* Worker threads determine their next action after completing a task using this method.
	*
	* <p>This acquires a CPU permit when returning {@link NextWorkerActivity#DO_CPU_HEAVY_TASK}.
	*/
	private NextWorkerActivity getActivityFollowingTask() {
	long snapshot = UNSAFE.getLongVolatile(null, countersAddress);
	do {
	if ((snapshot & (CANCEL_BIT \| TASKS_MASK)) > 0) {
	if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, snapshot - ONE_TASK)) {
	return DO_TASK;
	}
	} else if ((snapshot & (CANCEL_BIT \| CPU_HEAVY_TASKS_MASK)) > 0
	&& (snapshot & CPU_PERMITS_MASK) != 0) {
	if (UNSAFE.compareAndSwapLong(
	null, countersAddress, snapshot, snapshot - (ONE_CPU_HEAVY_TASK + ONE_CPU_PERMIT))) {
	return DO_CPU_HEAVY_TASK;
	}
	} else {
	long target = snapshot + ONE_THREAD;
	if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
	return IDLE;
	}
	}
	snapshot = UNSAFE.getLong(null, countersAddress);
	} while (true);
	}

	/**
	* Worker threads call this to determine their next action after completing a CPU heavy task.
	*
	* <p>This releases a CPU permit when returning {@link NextWorkerActivity#IDLE} or {@link
	* NextWorkerActivity#DO_TASK}.
	*/
	private NextWorkerActivity getActivityFollowingCpuHeavyTask() {
	long snapshot = UNSAFE.getLongVolatile(null, countersAddress);
	do {
	if ((snapshot & (CANCEL_BIT \| TASKS_MASK)) > 0) {
	if (UNSAFE.compareAndSwapLong(
	null, countersAddress, snapshot, snapshot + (ONE_CPU_PERMIT - ONE_TASK))) {
	return DO_TASK;
	}
	} else if ((snapshot & (CANCEL_BIT \| CPU_HEAVY_TASKS_MASK)) > 0) {
	if (UNSAFE.compareAndSwapLong(
	null, countersAddress, snapshot, snapshot - ONE_CPU_HEAVY_TASK)) {
	return DO_CPU_HEAVY_TASK;
	}
	} else {
	long target = snapshot + CPU_HEAVY_RESOURCES;
	if (UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target)) {
	return IDLE;
	}
	}
	snapshot = UNSAFE.getLong(null, countersAddress);
	} while (true);
	}

	// Throughout this class, the following wrappers are used where possible, but they are often not
	// inlined by the JVM even though they show up on profiles, so they are inlined explicitly in
	// numerous cases.

	private long getExecutionCounters() {
	return UNSAFE.getLongVolatile(null, countersAddress);
	}

	private boolean tryUpdateExecutionCounters(long snapshot, long target) {
	return UNSAFE.compareAndSwapLong(null, countersAddress, snapshot, target);
	}

	private static String formatSnapshot(long snapshot) {
	return String.format(
	"{cancelled=%b, threads=%d, cpuPermits=%d, tasks=%d, cpuHeavyTasks=%d}",
	snapshot < 0,
	(snapshot & THREADS_MASK) >> THREADS_BIT_OFFSET,
	(snapshot & CPU_PERMITS_MASK) >> CPU_PERMITS_BIT_OFFSET,
	(snapshot & TASKS_MASK) >> TASKS_BIT_OFFSET,
	(snapshot & CPU_HEAVY_TASKS_MASK) >> CPU_HEAVY_TASKS_BIT_OFFSET);
	}

	private static final Unsafe UNSAFE = UnsafeProvider.unsafe();
	}