Polish 'Fix race condition in OutputCapture'

See gh-46685

Author: philwebb

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/system/OutputCapture.java

@@ -23,7 +23,6 @@
 import java.util.ArrayList;
 import java.util.Deque;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
@@ -51,14 +50,15 @@ class OutputCapture implements CapturedOutput {
 
 	private AnsiOutputState ansiOutputState;
 
-	private final AtomicLong outVersion = new AtomicLong();
-	private final AtomicReference<VersionedCacheResult> out = new AtomicReference<>(null);
+	private final AtomicReference<Object> out = new AtomicReference<>();
 
-	private final AtomicLong errVersion = new AtomicLong();
-	private final AtomicReference<VersionedCacheResult> err = new AtomicReference<>(null);
+	private final AtomicReference<Object> err = new AtomicReference<>();
 
-	private final AtomicLong allVersion = new AtomicLong();
-	private final AtomicReference<VersionedCacheResult> all = new AtomicReference<>(null);
+	private final AtomicReference<Object> all = new AtomicReference<>();
+
+	OutputCapture() {
+		clearExisting();
+	}
 
 	/**
 	 * Push a new system capture session onto the stack.
@@ -111,7 +111,7 @@ public String toString() {
 	 */
 	@Override
 	public String getAll() {
-		return get(this.all, this.allVersion, (type) -> true);
+		return get(this.all, (type) -> true);
 	}
 
 	/**
@@ -120,7 +120,7 @@ public String getAll() {
 	 */
 	@Override
 	public String getOut() {
-		return get(this.out, this.outVersion, Type.OUT::equals);
+		return get(this.out, Type.OUT::equals);
 	}
 
 	/**
@@ -129,7 +129,7 @@ public String getOut() {
 	 */
 	@Override
 	public String getErr() {
-		return get(this.err, this.errVersion, Type.ERR::equals);
+		return get(this.err, Type.ERR::equals);
 	}
 
 	/**
@@ -141,25 +141,21 @@ void reset() {
 	}
 
 	void clearExisting() {
-		this.outVersion.incrementAndGet();
-		this.out.set(null);
-		this.errVersion.incrementAndGet();
-		this.err.set(null);
-		this.allVersion.incrementAndGet();
-		this.all.set(null);
+		this.out.set(new NoOutput());
+		this.err.set(new NoOutput());
+		this.all.set(new NoOutput());
 	}
 
-	private String get(AtomicReference<VersionedCacheResult> resultCache, AtomicLong version, Predicate<Type> filter) {
+	private String get(AtomicReference<Object> existing, Predicate<Type> filter) {
 		Assert.state(!this.systemCaptures.isEmpty(),
 				"No system captures found. Please check your output capture registration.");
-		long currentVersion = version.get();
-		VersionedCacheResult cached = resultCache.get();
-		if (cached != null && cached.version == currentVersion) {
-			return cached.result;
+		Object existingOutput = existing.get();
+		if (existingOutput instanceof String) {
+			return (String) existingOutput;
 		}
-		String result = build(filter);
-		resultCache.compareAndSet(null, new VersionedCacheResult(result, currentVersion));
-		return result;
+		String builtOutput = build(filter);
+		existing.compareAndSet(existingOutput, builtOutput);
+		return builtOutput;
 	}
 
 	String build(Predicate<Type> filter) {
@@ -170,10 +166,6 @@ String build(Predicate<Type> filter) {
 		return builder.toString();
 	}
 
-	private record VersionedCacheResult(String result, long version) {
-
-	}
-
 	/**
 	 * A capture session that captures {@link System#out System.out} and {@link System#out
 	 * System.err}.
@@ -353,4 +345,8 @@ static AnsiOutputState saveAndDisable() {
 
 	}
 
+	static class NoOutput {
+
+	}
+
 }

spring-boot-project/spring-boot-test/src/test/java/org/springframework/boot/test/system/OutputCaptureTests.java

@@ -22,10 +22,10 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
 
-import org.jspecify.annotations.Nullable;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -196,41 +196,21 @@ void getErrUsesCache() {
 
 	@Test
 	void getOutCacheShouldNotReturnStaleDataWhenDataIsLoggedWhileReading() throws Exception {
-		this.output.push();
+		TestLatchedOutputCapture output = new TestLatchedOutputCapture();
+		output.push();
 		System.out.print("A");
-		this.output.waitAfterBuildLatch = new CountDownLatch(1);
-
-		ExecutorService executorService = null;
+		ExecutorService executor = Executors.newFixedThreadPool(2);
 		try {
-			executorService = Executors.newFixedThreadPool(2);
-			var readingThreadFuture = executorService.submit(() -> {
-				// this will release the releaseAfterBuildLatch and block on the waitAfterBuildLatch
-				assertThat(this.output.getOut()).isEqualTo("A");
-			});
-			var writingThreadFuture = executorService.submit(() -> {
-				// wait until we finished building the first result (but did not yet update the cache)
-				try {
-					this.output.releaseAfterBuildLatch.await();
-				}
-				catch (InterruptedException e) {
-					throw new RuntimeException(e);
-				}
-				// print something else and then release the latch, for the other thread to continue
-				System.out.print("B");
-				this.output.waitAfterBuildLatch.countDown();
-			});
-			readingThreadFuture.get();
-			writingThreadFuture.get();
+			Future<?> reader = executor.submit(output::releaseAfterBuildAndAssertResultIsA);
+			Future<?> writer = executor.submit(output::awaitReleaseAfterBuildThenWriteBAndRelease);
+			reader.get();
+			writer.get();
 		}
 		finally {
-			if (executorService != null) {
-				executorService.shutdown();
-				executorService.awaitTermination(10, TimeUnit.SECONDS);
-			}
+			executor.shutdown();
+			executor.awaitTermination(10, TimeUnit.SECONDS);
 		}
-
-		// If not synchronized correctly this will fail, because the second print did not clear the cache and the cache will return stale data.
-		assertThat(this.output.getOut()).isEqualTo("AB");
+		assertThat(output.getOut()).isEqualTo("AB");
 	}
 
 	private void pushAndPrint() {
@@ -257,28 +237,48 @@ static class TestOutputCapture extends OutputCapture {
 
 		int buildCount;
 
-		@Nullable
-		CountDownLatch waitAfterBuildLatch = null;
+		@Override
+		String build(Predicate<Type> filter) {
+			this.buildCount++;
+			return super.build(filter);
+		}
+
+	}
+
+	static class TestLatchedOutputCapture extends OutputCapture {
 
-		CountDownLatch releaseAfterBuildLatch = new CountDownLatch(1);
+		private final CountDownLatch waitAfterBuild = new CountDownLatch(1);
+
+		private final CountDownLatch releaseAfterBuild = new CountDownLatch(1);
 
 		@Override
 		String build(Predicate<Type> filter) {
-			this.buildCount++;
 			var result = super.build(filter);
-			this.releaseAfterBuildLatch.countDown();
-			if (this.waitAfterBuildLatch != null) {
-				try {
-					this.waitAfterBuildLatch.await();
-				}
-				catch (InterruptedException e) {
-					Thread.currentThread().interrupt();
-					throw new RuntimeException(e);
-				}
-			}
+			this.releaseAfterBuild.countDown();
+			await(this.waitAfterBuild);
 			return result;
 		}
 
+		void releaseAfterBuildAndAssertResultIsA() {
+			assertThat(getOut()).isEqualTo("A");
+		}
+
+		void awaitReleaseAfterBuildThenWriteBAndRelease() {
+			await(this.releaseAfterBuild);
+			System.out.print("B");
+			this.waitAfterBuild.countDown();
+		}
+
+		private void await(CountDownLatch latch) {
+			try {
+				latch.await();
+			}
+			catch (InterruptedException ex) {
+				Thread.currentThread().interrupt();
+				throw new RuntimeException(ex);
+			}
+		}
+
 	}
 
 }