src/test/java/com/google/devtools/build/lib/concurrent/CancellableTaskTest.java - bazel - Git at Google

 // Copyright 2026 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.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.testutil.TestUtils.WAIT_TIMEOUT_MILLISECONDS;
 import static com.google.devtools.build.lib.testutil.TestUtils.WAIT_TIMEOUT_SECONDS;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.junit.Assert.assertThrows;

 import com.google.devtools.build.lib.testutil.TestThread;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 import org.junit.runner.RunWith;

 /** Tests for {@link CancellableTask}. */
 @RunWith(TestParameterInjector.class)
 public final class CancellableTaskTest {

   /** The cancellation entry point under test, which must not affect the cancellation semantics. */
   private enum CancelVariant {
     INTERRUPTIBLE {
       @Override
       void cancel(CancellableTask<?> task) throws InterruptedException {
         task.cancelAndAwait();
       }
     },
     REQUEST_THEN_AWAIT {
       @Override
       void cancel(CancellableTask<?> task) throws InterruptedException {
         task.requestCancellation();
         task.awaitCompletion();
       }
     };

     abstract void cancel(CancellableTask<?> task) throws InterruptedException;
   }

   @Test
   public void runTwice_throws() throws Exception {
     var task = new CancellableTask<>(() -> {});

     assertThat(task.runIfNotCancelled()).isTrue();

     assertThrows(IllegalStateException.class, task::runIfNotCancelled);
   }

   @Test
   public void cancelBeforeRun_preventsTaskFromRunning(@TestParameter CancelVariant variant)
       throws Exception {
     var ran = new AtomicBoolean();
     var task = new CancellableTask<>(() -> ran.set(true));

     variant.cancel(task);

     assertThat(task.runIfNotCancelled()).isFalse();
     assertThat(ran.get()).isFalse();
   }

   @Test
   public void cancelTwiceBeforeRun_returnsWithoutBlocking(@TestParameter CancelVariant variant)
       throws Exception {
     var task = new CancellableTask<>(() -> {});

     variant.cancel(task);
     variant.cancel(task);
   }

   @Test
   public void cancelDuringRun_interruptsAndAwaitsTask(@TestParameter CancelVariant variant)
       throws Exception {
     var taskStarted = new Semaphore(0);
     var taskInterrupted = new Semaphore(0);
     var taskMayFinish = new Semaphore(0);
     var cancellationReturned = new Semaphore(0);
     var task =
         new CancellableTask<>(
             () -> {
               taskStarted.release();
               try {
                 taskMayFinish.acquire();
               } catch (InterruptedException e) {
                 taskInterrupted.release();
                 taskMayFinish.acquireUninterruptibly();
               }
             });
     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var canceller =
         new TestThread(
             () -> {
               variant.cancel(task);
               cancellationReturned.release();
             });
     canceller.start();
     try {
       assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
       assertThat(taskInterrupted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     } finally {
       taskMayFinish.release();
     }

     assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void cancelAfterRun_returnsWithoutBlocking(@TestParameter CancelVariant variant)
       throws Exception {
     var task = new CancellableTask<>(() -> {});

     assertThat(task.runIfNotCancelled()).isTrue();

     variant.cancel(task);
     // Repeated cancellation is allowed and must not block either.
     variant.cancel(task);
   }

   @Test
   public void completion_runsOnceAfterBody() throws Exception {
     var events = new ArrayList<String>();
     var task = new CancellableTask<>(() -> events.add("body"), () -> events.add("completion"));

     assertThat(task.runIfNotCancelled()).isTrue();

     assertThat(events).containsExactly("body", "completion").inOrder();
     // Cancellation after completion must not run the completion action again.
     task.cancelAndAwait();
     assertThat(events).containsExactly("body", "completion").inOrder();
   }

   @Test
   public void completion_runsWhenBodyThrows() {
     var completionRuns = new AtomicInteger();
     var task =
         new CancellableTask<>(
             () -> {
               throw new IOException("task failed");
             },
             completionRuns::incrementAndGet);

     assertThrows(IOException.class, task::runIfNotCancelled);

     assertThat(completionRuns.get()).isEqualTo(1);
   }

   @Test
   public void cancelBeforeRun_runsCompletionExactlyOnce(@TestParameter CancelVariant variant)
       throws Exception {
     var completionRuns = new AtomicInteger();
     var task = new CancellableTask<>(() -> {}, completionRuns::incrementAndGet);

     variant.cancel(task);
     assertThat(completionRuns.get()).isEqualTo(1);

     // Neither a repeated cancellation nor the prevented run runs the completion action again.
     variant.cancel(task);
     assertThat(task.runIfNotCancelled()).isFalse();
     assertThat(completionRuns.get()).isEqualTo(1);
   }

   @Test
   public void cancelBeforeRun_concurrentCancellerAwaitsCompletion(
       @TestParameter CancelVariant variant) throws Exception {
     var completionStarted = new Semaphore(0);
     var completionMayFinish = new Semaphore(0);
     var task =
         new CancellableTask<>(
             () -> {},
             () -> {
               completionStarted.release();
               completionMayFinish.acquireUninterruptibly();
             });
     var winner = new TestThread(() -> variant.cancel(task));
     winner.start();
     assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var loserReturned = new Semaphore(0);
     var loser =
         new TestThread(
             () -> {
               variant.cancel(task);
               loserReturned.release();
             });
     loser.start();
     try {
       // The losing canceler must not return while the winner's completion action is running.
       assertThat(loserReturned.tryAcquire(100, MILLISECONDS)).isFalse();
     } finally {
       completionMayFinish.release();
     }

     assertThat(loserReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     winner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     loser.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void cancelDuringCompletion_awaitsWithoutInterruptingCompletion(
       @TestParameter CancelVariant variant) throws Exception {
     var completionStarted = new Semaphore(0);
     var completionMayFinish = new Semaphore(0);
     var completionWasInterrupted = new AtomicBoolean();
     var task =
         new CancellableTask<>(
             () -> {},
             () -> {
               completionStarted.release();
               try {
                 completionMayFinish.acquire();
               } catch (InterruptedException e) {
                 completionWasInterrupted.set(true);
               }
             });
     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var cancellationReturned = new Semaphore(0);
     var canceller =
         new TestThread(
             () -> {
               variant.cancel(task);
               cancellationReturned.release();
             });
     canceller.start();
     try {
       // Completion is part of the terminal lifecycle, but cancellation only interrupts the body.
       assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
       assertThat(completionWasInterrupted.get()).isFalse();
     } finally {
       completionMayFinish.release();
     }

     assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(completionWasInterrupted.get()).isFalse();
   }

   @Test
   public void runLosingToCancelBeforeStart_awaitsCompletion(@TestParameter CancelVariant variant)
       throws Exception {
     var completionStarted = new Semaphore(0);
     var completionMayFinish = new Semaphore(0);
     var task =
         new CancellableTask<>(
             () -> {},
             () -> {
               completionStarted.release();
               completionMayFinish.acquireUninterruptibly();
             });
     var canceller = new TestThread(() -> variant.cancel(task));
     canceller.start();
     assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var runReturned = new Semaphore(0);
     var ran = new AtomicBoolean(true);
     var runner =
         new TestThread(
             () -> {
               ran.set(task.runIfNotCancelled());
               runReturned.release();
             });
     runner.start();
     try {
       assertThat(runReturned.tryAcquire(100, MILLISECONDS)).isFalse();
     } finally {
       completionMayFinish.release();
     }

     assertThat(runReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(ran.get()).isFalse();
   }

   @Test
   public void taskThrows_propagatesExceptionAndUnblocksCanceller(
       @TestParameter CancelVariant variant) throws Exception {
     var taskStarted = new Semaphore(0);
     var taskMayThrow = new Semaphore(0);
     var task =
         new CancellableTask<IOException>(
             () -> {
               taskStarted.release();
               taskMayThrow.acquireUninterruptibly();
               throw new IOException("task failed");
             });
     var thrown = new AtomicReference<IOException>();
     var runner =
         new TestThread(() -> thrown.set(assertThrows(IOException.class, task::runIfNotCancelled)));
     runner.start();
     assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var cancellationReturned = new Semaphore(0);
     var canceller =
         new TestThread(
             () -> {
               variant.cancel(task);
               cancellationReturned.release();
             });
     canceller.start();
     try {
       assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
     } finally {
       taskMayThrow.release();
     }

     // A task that leaves its body by throwing must still unblock cancelers.
     assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(thrown.get()).hasMessageThat().isEqualTo("task failed");
   }

   @Test
   public void cancelFromTaskBody_throws(@TestParameter CancelVariant variant) throws Exception {
     var taskRef = new AtomicReference<CancellableTask<Exception>>();
     var task =
         new CancellableTask<Exception>(
             () -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
     taskRef.set(task);

     assertThat(task.runIfNotCancelled()).isTrue();
   }

   @Test
   public void cancelFromCompletionAfterBody_throws(@TestParameter CancelVariant variant)
       throws Exception {
     var taskRef = new AtomicReference<CancellableTask<Exception>>();
     var task =
         new CancellableTask<Exception>(
             () -> {},
             () -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
     taskRef.set(task);

     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void cancelFromCancelBeforeStartCompletion_throws(@TestParameter CancelVariant variant)
       throws Exception {
     var taskRef = new AtomicReference<CancellableTask<Exception>>();
     var task =
         new CancellableTask<Exception>(
             () -> {},
             () -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
     taskRef.set(task);

     var canceller = new TestThread(() -> variant.cancel(task));
     canceller.start();
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(task.runIfNotCancelled()).isFalse();
   }

   @Test
   public void runFromCancelBeforeStartCompletion_throws() throws Exception {
     var taskRef = new AtomicReference<CancellableTask<Exception>>();
     var task =
         new CancellableTask<Exception>(
             () -> {},
             () ->
                 assertThrows(IllegalStateException.class, () -> taskRef.get().runIfNotCancelled()));
     taskRef.set(task);

     var canceller = new TestThread(task::cancelAndAwait);
     canceller.start();
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void completionThrowsAfterBody_stillSignalsTerminalState(
       @TestParameter CancelVariant variant) throws Exception {
     var task =
         new CancellableTask<>(
             () -> {},
             () -> {
               throw new IllegalStateException("completion failed");
             });

     assertThat(assertThrows(IllegalStateException.class, task::runIfNotCancelled))
         .hasMessageThat()
         .isEqualTo("completion failed");

     assertLaterCancellationReturns(task, variant);
   }

   @Test
   public void completionThrowsAfterCancelBeforeStart_stillSignalsTerminalState(
       @TestParameter CancelVariant variant) throws Exception {
     var task =
         new CancellableTask<>(
             () -> {},
             () -> {
               throw new IllegalStateException("completion failed");
             });

     assertThat(assertThrows(IllegalStateException.class, () -> variant.cancel(task)))
         .hasMessageThat()
         .isEqualTo("completion failed");

     assertLaterCancellationReturns(task, variant);
     assertThat(task.runIfNotCancelled()).isFalse();
   }

   private static void assertLaterCancellationReturns(CancellableTask<?> task, CancelVariant variant)
       throws Exception {
     var cancellationReturned = new Semaphore(0);
     var canceller =
         new TestThread(
             () -> {
               variant.cancel(task);
               cancellationReturned.release();
             });
     canceller.start();

     assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void requestCancellation_beforeRun_preventsStartAndRunsCompletion() throws Exception {
     var ran = new AtomicBoolean();
     var completionRuns = new AtomicInteger();
     var task = new CancellableTask<>(() -> ran.set(true), completionRuns::incrementAndGet);

     task.requestCancellation();

     assertThat(completionRuns.get()).isEqualTo(1);
     assertThat(task.runIfNotCancelled()).isFalse();
     assertThat(ran.get()).isFalse();
     assertThat(completionRuns.get()).isEqualTo(1);
   }

   @Test
   public void requestCancellation_duringRun_interruptsWithoutAwaiting() throws Exception {
     var taskStarted = new Semaphore(0);
     var taskInterrupted = new Semaphore(0);
     var taskMayFinish = new Semaphore(0);
     var task =
         new CancellableTask<>(
             () -> {
               taskStarted.release();
               try {
                 taskMayFinish.acquire();
               } catch (InterruptedException e) {
                 taskInterrupted.release();
                 taskMayFinish.acquireUninterruptibly();
               }
             });
     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     // Returns while the body is still running: it only finishes after the release below.
     task.requestCancellation();

     assertThat(taskInterrupted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     taskMayFinish.release();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
   }

   @Test
   public void requestCancellation_thenAwaitCompletion_doesNotInterruptTaskAgain() throws Exception {
     var taskStarted = new Semaphore(0);
     var cleanupStarted = new Semaphore(0);
     var cleanupMayFinish = new Semaphore(0);
     var interruptionCount = new AtomicInteger();
     var task =
         new CancellableTask<>(
             () -> {
               taskStarted.release();
               try {
                 new Semaphore(0).acquire();
               } catch (InterruptedException e) {
                 interruptionCount.incrementAndGet();
               }
               cleanupStarted.release();
               try {
                 cleanupMayFinish.acquire();
               } catch (InterruptedException e) {
                 interruptionCount.incrementAndGet();
               }
             });
     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     task.requestCancellation();
     assertThat(cleanupStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     var waiter = new TestThread(task::awaitCompletion);
     waiter.start();
     try {
       // Awaiting a previously requested cancellation must not send another interrupt that could
       // abort the task's cleanup.
       assertThat(waiter.isAlive()).isTrue();
       assertThat(interruptionCount.get()).isEqualTo(1);
     } finally {
       cleanupMayFinish.release();
     }

     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     waiter.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(interruptionCount.get()).isEqualTo(1);
   }

   @Test
   public void cancelAndAwait_interruptedWhileAwaiting_retryInterruptsAgainAndAwaits()
       throws Exception {
     var taskStarted = new Semaphore(0);
     var taskMayFinish = new Semaphore(0);
     var taskInterruptions = new Semaphore(0);
     var interruptionCount = new AtomicInteger();
     var task =
         new CancellableTask<>(
             () -> {
               taskStarted.release();
               while (true) {
                 try {
                   taskMayFinish.acquire();
                   return;
                 } catch (InterruptedException e) {
                   interruptionCount.incrementAndGet();
                   taskInterruptions.release();
                 }
               }
             });
     var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
     runner.start();
     assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     var cancellerThrew = new Semaphore(0);
     var cancellationReturned = new Semaphore(0);
     var canceller =
         new TestThread(
             () -> {
               assertThrows(InterruptedException.class, task::cancelAndAwait);
               cancellerThrew.release();
               // A retry interrupts the still-running task again and awaits quiescence.
               task.cancelAndAwait();
               cancellationReturned.release();
             });
     canceller.start();
     assertThat(taskInterruptions.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

     canceller.interrupt();
     assertThat(cancellerThrew.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     assertThat(taskInterruptions.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     try {
       assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
       assertThat(interruptionCount.get()).isEqualTo(2);
     } finally {
       taskMayFinish.release();
     }

     assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
     runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
     assertThat(interruptionCount.get()).isEqualTo(2);
   }

   @Test
   public void runAfterCancel_whileInterrupted_restoresInterruptBit() throws Exception {
     var task = new CancellableTask<>(() -> {});
     task.cancelAndAwait();
     Thread.currentThread().interrupt();

     assertThat(task.runIfNotCancelled()).isFalse();

     assertThat(Thread.interrupted()).isTrue();
   }

   @Test
   public void concurrentRunAndCancel_exactlyOneClaimWins() throws Exception {
     for (int i = 0; i < 1000; i++) {
       var bodyRuns = new AtomicInteger();
       var bodyRunning = new AtomicBoolean();
       var completionRuns = new AtomicInteger();
       var task =
           new CancellableTask<>(
               () -> {
                 bodyRunning.set(true);
                 bodyRuns.incrementAndGet();
                 bodyRunning.set(false);
               },
               completionRuns::incrementAndGet);
       var barrier = new CyclicBarrier(3);
       var ran = new AtomicBoolean();
       var runner =
           new TestThread(
               () -> {
                 barrier.await();
                 ran.set(task.runIfNotCancelled());
               });
       var cancellers = new TestThread[2];
       for (int c = 0; c < cancellers.length; c++) {
         cancellers[c] =
             new TestThread(
                 () -> {
                   barrier.await();
                   task.cancelAndAwait();
                   // A normal return from cancellation guarantees that the body is not running
                   // and that the completion action has finished.
                   assertThat(bodyRunning.get()).isFalse();
                   assertThat(completionRuns.get()).isEqualTo(1);
                 });
       }
       // Rotate which thread arrives at the barrier last: that thread trips the barrier and
       // proceeds without parking while the others are still being woken, so it tends to win the
       // claim. Rotation makes both outcomes of the race common.
       var startOrder =
           switch (i % 3) {
             case 0 -> new TestThread[] {runner, cancellers[0], cancellers[1]};
             case 1 -> new TestThread[] {cancellers[0], cancellers[1], runner};
             default -> new TestThread[] {cancellers[0], runner, cancellers[1]};
           };
       for (var thread : startOrder) {
         thread.start();
       }
       runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
       for (var canceller : cancellers) {
         canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
       }

       // The claim is atomic: either the task ran to completion exactly once and no canceler
       // prevented it, or cancellation prevented it from running.
       assertThat(ran.get()).isEqualTo(bodyRuns.get() == 1);
       assertThat(bodyRuns.get()).isAtMost(1);
       assertThat(completionRuns.get()).isEqualTo(1);
     }
   }
 }
	// Copyright 2026 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.truth.Truth.assertThat;
	import static com.google.devtools.build.lib.testutil.TestUtils.WAIT_TIMEOUT_MILLISECONDS;
	import static com.google.devtools.build.lib.testutil.TestUtils.WAIT_TIMEOUT_SECONDS;
	import static java.util.concurrent.TimeUnit.MILLISECONDS;
	import static java.util.concurrent.TimeUnit.SECONDS;
	import static org.junit.Assert.assertThrows;

	import com.google.devtools.build.lib.testutil.TestThread;
	import com.google.testing.junit.testparameterinjector.TestParameter;
	import com.google.testing.junit.testparameterinjector.TestParameterInjector;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.concurrent.CyclicBarrier;
	import java.util.concurrent.Semaphore;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.concurrent.atomic.AtomicReference;
	import org.junit.Test;
	import org.junit.runner.RunWith;

	/** Tests for {@link CancellableTask}. */
	@RunWith(TestParameterInjector.class)
	public final class CancellableTaskTest {

	/** The cancellation entry point under test, which must not affect the cancellation semantics. */
	private enum CancelVariant {
	INTERRUPTIBLE {
	@Override
	void cancel(CancellableTask<?> task) throws InterruptedException {
	task.cancelAndAwait();
	}
	},
	REQUEST_THEN_AWAIT {
	@Override
	void cancel(CancellableTask<?> task) throws InterruptedException {
	task.requestCancellation();
	task.awaitCompletion();
	}
	};

	abstract void cancel(CancellableTask<?> task) throws InterruptedException;
	}

	@Test
	public void runTwice_throws() throws Exception {
	var task = new CancellableTask<>(() -> {});

	assertThat(task.runIfNotCancelled()).isTrue();

	assertThrows(IllegalStateException.class, task::runIfNotCancelled);
	}

	@Test
	public void cancelBeforeRun_preventsTaskFromRunning(@TestParameter CancelVariant variant)
	throws Exception {
	var ran = new AtomicBoolean();
	var task = new CancellableTask<>(() -> ran.set(true));

	variant.cancel(task);

	assertThat(task.runIfNotCancelled()).isFalse();
	assertThat(ran.get()).isFalse();
	}

	@Test
	public void cancelTwiceBeforeRun_returnsWithoutBlocking(@TestParameter CancelVariant variant)
	throws Exception {
	var task = new CancellableTask<>(() -> {});

	variant.cancel(task);
	variant.cancel(task);
	}

	@Test
	public void cancelDuringRun_interruptsAndAwaitsTask(@TestParameter CancelVariant variant)
	throws Exception {
	var taskStarted = new Semaphore(0);
	var taskInterrupted = new Semaphore(0);
	var taskMayFinish = new Semaphore(0);
	var cancellationReturned = new Semaphore(0);
	var task =
	new CancellableTask<>(
	() -> {
	taskStarted.release();
	try {
	taskMayFinish.acquire();
	} catch (InterruptedException e) {
	taskInterrupted.release();
	taskMayFinish.acquireUninterruptibly();
	}
	});
	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var canceller =
	new TestThread(
	() -> {
	variant.cancel(task);
	cancellationReturned.release();
	});
	canceller.start();
	try {
	assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	assertThat(taskInterrupted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	} finally {
	taskMayFinish.release();
	}

	assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void cancelAfterRun_returnsWithoutBlocking(@TestParameter CancelVariant variant)
	throws Exception {
	var task = new CancellableTask<>(() -> {});

	assertThat(task.runIfNotCancelled()).isTrue();

	variant.cancel(task);
	// Repeated cancellation is allowed and must not block either.
	variant.cancel(task);
	}

	@Test
	public void completion_runsOnceAfterBody() throws Exception {
	var events = new ArrayList<String>();
	var task = new CancellableTask<>(() -> events.add("body"), () -> events.add("completion"));

	assertThat(task.runIfNotCancelled()).isTrue();

	assertThat(events).containsExactly("body", "completion").inOrder();
	// Cancellation after completion must not run the completion action again.
	task.cancelAndAwait();
	assertThat(events).containsExactly("body", "completion").inOrder();
	}

	@Test
	public void completion_runsWhenBodyThrows() {
	var completionRuns = new AtomicInteger();
	var task =
	new CancellableTask<>(
	() -> {
	throw new IOException("task failed");
	},
	completionRuns::incrementAndGet);

	assertThrows(IOException.class, task::runIfNotCancelled);

	assertThat(completionRuns.get()).isEqualTo(1);
	}

	@Test
	public void cancelBeforeRun_runsCompletionExactlyOnce(@TestParameter CancelVariant variant)
	throws Exception {
	var completionRuns = new AtomicInteger();
	var task = new CancellableTask<>(() -> {}, completionRuns::incrementAndGet);

	variant.cancel(task);
	assertThat(completionRuns.get()).isEqualTo(1);

	// Neither a repeated cancellation nor the prevented run runs the completion action again.
	variant.cancel(task);
	assertThat(task.runIfNotCancelled()).isFalse();
	assertThat(completionRuns.get()).isEqualTo(1);
	}

	@Test
	public void cancelBeforeRun_concurrentCancellerAwaitsCompletion(
	@TestParameter CancelVariant variant) throws Exception {
	var completionStarted = new Semaphore(0);
	var completionMayFinish = new Semaphore(0);
	var task =
	new CancellableTask<>(
	() -> {},
	() -> {
	completionStarted.release();
	completionMayFinish.acquireUninterruptibly();
	});
	var winner = new TestThread(() -> variant.cancel(task));
	winner.start();
	assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var loserReturned = new Semaphore(0);
	var loser =
	new TestThread(
	() -> {
	variant.cancel(task);
	loserReturned.release();
	});
	loser.start();
	try {
	// The losing canceler must not return while the winner's completion action is running.
	assertThat(loserReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	} finally {
	completionMayFinish.release();
	}

	assertThat(loserReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	winner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	loser.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void cancelDuringCompletion_awaitsWithoutInterruptingCompletion(
	@TestParameter CancelVariant variant) throws Exception {
	var completionStarted = new Semaphore(0);
	var completionMayFinish = new Semaphore(0);
	var completionWasInterrupted = new AtomicBoolean();
	var task =
	new CancellableTask<>(
	() -> {},
	() -> {
	completionStarted.release();
	try {
	completionMayFinish.acquire();
	} catch (InterruptedException e) {
	completionWasInterrupted.set(true);
	}
	});
	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var cancellationReturned = new Semaphore(0);
	var canceller =
	new TestThread(
	() -> {
	variant.cancel(task);
	cancellationReturned.release();
	});
	canceller.start();
	try {
	// Completion is part of the terminal lifecycle, but cancellation only interrupts the body.
	assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	assertThat(completionWasInterrupted.get()).isFalse();
	} finally {
	completionMayFinish.release();
	}

	assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(completionWasInterrupted.get()).isFalse();
	}

	@Test
	public void runLosingToCancelBeforeStart_awaitsCompletion(@TestParameter CancelVariant variant)
	throws Exception {
	var completionStarted = new Semaphore(0);
	var completionMayFinish = new Semaphore(0);
	var task =
	new CancellableTask<>(
	() -> {},
	() -> {
	completionStarted.release();
	completionMayFinish.acquireUninterruptibly();
	});
	var canceller = new TestThread(() -> variant.cancel(task));
	canceller.start();
	assertThat(completionStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var runReturned = new Semaphore(0);
	var ran = new AtomicBoolean(true);
	var runner =
	new TestThread(
	() -> {
	ran.set(task.runIfNotCancelled());
	runReturned.release();
	});
	runner.start();
	try {
	assertThat(runReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	} finally {
	completionMayFinish.release();
	}

	assertThat(runReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(ran.get()).isFalse();
	}

	@Test
	public void taskThrows_propagatesExceptionAndUnblocksCanceller(
	@TestParameter CancelVariant variant) throws Exception {
	var taskStarted = new Semaphore(0);
	var taskMayThrow = new Semaphore(0);
	var task =
	new CancellableTask<IOException>(
	() -> {
	taskStarted.release();
	taskMayThrow.acquireUninterruptibly();
	throw new IOException("task failed");
	});
	var thrown = new AtomicReference<IOException>();
	var runner =
	new TestThread(() -> thrown.set(assertThrows(IOException.class, task::runIfNotCancelled)));
	runner.start();
	assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var cancellationReturned = new Semaphore(0);
	var canceller =
	new TestThread(
	() -> {
	variant.cancel(task);
	cancellationReturned.release();
	});
	canceller.start();
	try {
	assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	} finally {
	taskMayThrow.release();
	}

	// A task that leaves its body by throwing must still unblock cancelers.
	assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(thrown.get()).hasMessageThat().isEqualTo("task failed");
	}

	@Test
	public void cancelFromTaskBody_throws(@TestParameter CancelVariant variant) throws Exception {
	var taskRef = new AtomicReference<CancellableTask<Exception>>();
	var task =
	new CancellableTask<Exception>(
	() -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
	taskRef.set(task);

	assertThat(task.runIfNotCancelled()).isTrue();
	}

	@Test
	public void cancelFromCompletionAfterBody_throws(@TestParameter CancelVariant variant)
	throws Exception {
	var taskRef = new AtomicReference<CancellableTask<Exception>>();
	var task =
	new CancellableTask<Exception>(
	() -> {},
	() -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
	taskRef.set(task);

	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void cancelFromCancelBeforeStartCompletion_throws(@TestParameter CancelVariant variant)
	throws Exception {
	var taskRef = new AtomicReference<CancellableTask<Exception>>();
	var task =
	new CancellableTask<Exception>(
	() -> {},
	() -> assertThrows(IllegalStateException.class, () -> variant.cancel(taskRef.get())));
	taskRef.set(task);

	var canceller = new TestThread(() -> variant.cancel(task));
	canceller.start();
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(task.runIfNotCancelled()).isFalse();
	}

	@Test
	public void runFromCancelBeforeStartCompletion_throws() throws Exception {
	var taskRef = new AtomicReference<CancellableTask<Exception>>();
	var task =
	new CancellableTask<Exception>(
	() -> {},
	() ->
	assertThrows(IllegalStateException.class, () -> taskRef.get().runIfNotCancelled()));
	taskRef.set(task);

	var canceller = new TestThread(task::cancelAndAwait);
	canceller.start();
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void completionThrowsAfterBody_stillSignalsTerminalState(
	@TestParameter CancelVariant variant) throws Exception {
	var task =
	new CancellableTask<>(
	() -> {},
	() -> {
	throw new IllegalStateException("completion failed");
	});

	assertThat(assertThrows(IllegalStateException.class, task::runIfNotCancelled))
	.hasMessageThat()
	.isEqualTo("completion failed");

	assertLaterCancellationReturns(task, variant);
	}

	@Test
	public void completionThrowsAfterCancelBeforeStart_stillSignalsTerminalState(
	@TestParameter CancelVariant variant) throws Exception {
	var task =
	new CancellableTask<>(
	() -> {},
	() -> {
	throw new IllegalStateException("completion failed");
	});

	assertThat(assertThrows(IllegalStateException.class, () -> variant.cancel(task)))
	.hasMessageThat()
	.isEqualTo("completion failed");

	assertLaterCancellationReturns(task, variant);
	assertThat(task.runIfNotCancelled()).isFalse();
	}

	private static void assertLaterCancellationReturns(CancellableTask<?> task, CancelVariant variant)
	throws Exception {
	var cancellationReturned = new Semaphore(0);
	var canceller =
	new TestThread(
	() -> {
	variant.cancel(task);
	cancellationReturned.release();
	});
	canceller.start();

	assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void requestCancellation_beforeRun_preventsStartAndRunsCompletion() throws Exception {
	var ran = new AtomicBoolean();
	var completionRuns = new AtomicInteger();
	var task = new CancellableTask<>(() -> ran.set(true), completionRuns::incrementAndGet);

	task.requestCancellation();

	assertThat(completionRuns.get()).isEqualTo(1);
	assertThat(task.runIfNotCancelled()).isFalse();
	assertThat(ran.get()).isFalse();
	assertThat(completionRuns.get()).isEqualTo(1);
	}

	@Test
	public void requestCancellation_duringRun_interruptsWithoutAwaiting() throws Exception {
	var taskStarted = new Semaphore(0);
	var taskInterrupted = new Semaphore(0);
	var taskMayFinish = new Semaphore(0);
	var task =
	new CancellableTask<>(
	() -> {
	taskStarted.release();
	try {
	taskMayFinish.acquire();
	} catch (InterruptedException e) {
	taskInterrupted.release();
	taskMayFinish.acquireUninterruptibly();
	}
	});
	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	// Returns while the body is still running: it only finishes after the release below.
	task.requestCancellation();

	assertThat(taskInterrupted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	taskMayFinish.release();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	@Test
	public void requestCancellation_thenAwaitCompletion_doesNotInterruptTaskAgain() throws Exception {
	var taskStarted = new Semaphore(0);
	var cleanupStarted = new Semaphore(0);
	var cleanupMayFinish = new Semaphore(0);
	var interruptionCount = new AtomicInteger();
	var task =
	new CancellableTask<>(
	() -> {
	taskStarted.release();
	try {
	new Semaphore(0).acquire();
	} catch (InterruptedException e) {
	interruptionCount.incrementAndGet();
	}
	cleanupStarted.release();
	try {
	cleanupMayFinish.acquire();
	} catch (InterruptedException e) {
	interruptionCount.incrementAndGet();
	}
	});
	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	task.requestCancellation();
	assertThat(cleanupStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	var waiter = new TestThread(task::awaitCompletion);
	waiter.start();
	try {
	// Awaiting a previously requested cancellation must not send another interrupt that could
	// abort the task's cleanup.
	assertThat(waiter.isAlive()).isTrue();
	assertThat(interruptionCount.get()).isEqualTo(1);
	} finally {
	cleanupMayFinish.release();
	}

	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	waiter.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(interruptionCount.get()).isEqualTo(1);
	}

	@Test
	public void cancelAndAwait_interruptedWhileAwaiting_retryInterruptsAgainAndAwaits()
	throws Exception {
	var taskStarted = new Semaphore(0);
	var taskMayFinish = new Semaphore(0);
	var taskInterruptions = new Semaphore(0);
	var interruptionCount = new AtomicInteger();
	var task =
	new CancellableTask<>(
	() -> {
	taskStarted.release();
	while (true) {
	try {
	taskMayFinish.acquire();
	return;
	} catch (InterruptedException e) {
	interruptionCount.incrementAndGet();
	taskInterruptions.release();
	}
	}
	});
	var runner = new TestThread(() -> assertThat(task.runIfNotCancelled()).isTrue());
	runner.start();
	assertThat(taskStarted.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	var cancellerThrew = new Semaphore(0);
	var cancellationReturned = new Semaphore(0);
	var canceller =
	new TestThread(
	() -> {
	assertThrows(InterruptedException.class, task::cancelAndAwait);
	cancellerThrew.release();
	// A retry interrupts the still-running task again and awaits quiescence.
	task.cancelAndAwait();
	cancellationReturned.release();
	});
	canceller.start();
	assertThat(taskInterruptions.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();

	canceller.interrupt();
	assertThat(cancellerThrew.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	assertThat(taskInterruptions.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	try {
	assertThat(cancellationReturned.tryAcquire(100, MILLISECONDS)).isFalse();
	assertThat(interruptionCount.get()).isEqualTo(2);
	} finally {
	taskMayFinish.release();
	}

	assertThat(cancellationReturned.tryAcquire(WAIT_TIMEOUT_SECONDS, SECONDS)).isTrue();
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	assertThat(interruptionCount.get()).isEqualTo(2);
	}

	@Test
	public void runAfterCancel_whileInterrupted_restoresInterruptBit() throws Exception {
	var task = new CancellableTask<>(() -> {});
	task.cancelAndAwait();
	Thread.currentThread().interrupt();

	assertThat(task.runIfNotCancelled()).isFalse();

	assertThat(Thread.interrupted()).isTrue();
	}

	@Test
	public void concurrentRunAndCancel_exactlyOneClaimWins() throws Exception {
	for (int i = 0; i < 1000; i++) {
	var bodyRuns = new AtomicInteger();
	var bodyRunning = new AtomicBoolean();
	var completionRuns = new AtomicInteger();
	var task =
	new CancellableTask<>(
	() -> {
	bodyRunning.set(true);
	bodyRuns.incrementAndGet();
	bodyRunning.set(false);
	},
	completionRuns::incrementAndGet);
	var barrier = new CyclicBarrier(3);
	var ran = new AtomicBoolean();
	var runner =
	new TestThread(
	() -> {
	barrier.await();
	ran.set(task.runIfNotCancelled());
	});
	var cancellers = new TestThread[2];
	for (int c = 0; c < cancellers.length; c++) {
	cancellers[c] =
	new TestThread(
	() -> {
	barrier.await();
	task.cancelAndAwait();
	// A normal return from cancellation guarantees that the body is not running
	// and that the completion action has finished.
	assertThat(bodyRunning.get()).isFalse();
	assertThat(completionRuns.get()).isEqualTo(1);
	});
	}
	// Rotate which thread arrives at the barrier last: that thread trips the barrier and
	// proceeds without parking while the others are still being woken, so it tends to win the
	// claim. Rotation makes both outcomes of the race common.
	var startOrder =
	switch (i % 3) {
	case 0 -> new TestThread[] {runner, cancellers[0], cancellers[1]};
	case 1 -> new TestThread[] {cancellers[0], cancellers[1], runner};
	default -> new TestThread[] {cancellers[0], runner, cancellers[1]};
	};
	for (var thread : startOrder) {
	thread.start();
	}
	runner.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	for (var canceller : cancellers) {
	canceller.joinAndAssertState(WAIT_TIMEOUT_MILLISECONDS);
	}

	// The claim is atomic: either the task ran to completion exactly once and no canceler
	// prevented it, or cancellation prevented it from running.
	assertThat(ran.get()).isEqualTo(bodyRuns.get() == 1);
	assertThat(bodyRuns.get()).isAtMost(1);
	assertThat(completionRuns.get()).isEqualTo(1);
	}
	}
	}