Revert "Fix handling of timeout in SseEmitter"

This reverts commit f92f9c1.

See gh-34762

Author: rstoyanchev

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2025 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,7 +21,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 
 import org.springframework.http.MediaType;
@@ -73,20 +73,21 @@ public class ResponseBodyEmitter {
 	@Nullable
 	private Handler handler;
 
-	private final AtomicReference<State> state = new AtomicReference<>(State.START);
-
 	/** Store send data before handler is initialized. */
 	private final Set<DataWithMediaType> earlySendAttempts = new LinkedHashSet<>(8);
 
+	/** Store successful completion before the handler is initialized. */
+	private final AtomicBoolean complete = new AtomicBoolean();
+
 	/** Store an error before the handler is initialized. */
 	@Nullable
 	private Throwable failure;
 
-	private final TimeoutCallback timeoutCallback = new TimeoutCallback();
+	private final DefaultCallback timeoutCallback = new DefaultCallback();
 
 	private final ErrorCallback errorCallback = new ErrorCallback();
 
-	private final CompletionCallback completionCallback = new CompletionCallback();
+	private final DefaultCallback completionCallback = new DefaultCallback();
 
 
 	/**
@@ -127,7 +128,7 @@ synchronized void initialize(Handler handler) throws IOException {
 			this.earlySendAttempts.clear();
 		}
 
-		if (this.state.get() == State.COMPLETE) {
+		if (this.complete.get()) {
 			if (this.failure != null) {
 				this.handler.completeWithError(this.failure);
 			}
@@ -143,7 +144,7 @@ synchronized void initialize(Handler handler) throws IOException {
 	}
 
 	void initializeWithError(Throwable ex) {
-		if (this.state.compareAndSet(State.START, State.COMPLETE)) {
+		if (this.complete.compareAndSet(false, true)) {
 			this.failure = ex;
 			this.earlySendAttempts.clear();
 			this.errorCallback.accept(ex);
@@ -185,7 +186,8 @@ public void send(Object object) throws IOException {
 	 * @throws java.lang.IllegalStateException wraps any other errors
 	 */
 	public synchronized void send(Object object, @Nullable MediaType mediaType) throws IOException {
-		assertNotComplete();
+		Assert.state(!this.complete.get(), () -> "ResponseBodyEmitter has already completed" +
+				(this.failure != null ? " with error: " + this.failure : ""));
 		if (this.handler != null) {
 			try {
 				this.handler.send(object, mediaType);
@@ -212,13 +214,9 @@ public synchronized void send(Object object, @Nullable MediaType mediaType) thro
 	 * @since 6.0.12
 	 */
 	public synchronized void send(Set<DataWithMediaType> items) throws IOException {
-		assertNotComplete();
-		sendInternal(items);
-	}
-
-	private void assertNotComplete() {
-		Assert.state(this.state.get() == State.START, () -> "ResponseBodyEmitter has already completed" +
+		Assert.state(!this.complete.get(), () -> "ResponseBodyEmitter has already completed" +
 				(this.failure != null ? " with error: " + this.failure : ""));
+		sendInternal(items);
 	}
 
 	private void sendInternal(Set<DataWithMediaType> items) throws IOException {
@@ -250,7 +248,7 @@ private void sendInternal(Set<DataWithMediaType> items) throws IOException {
 	 * related events such as an error while {@link #send(Object) sending}.
 	 */
 	public void complete() {
-		if (trySetComplete() && this.handler != null) {
+		if (this.complete.compareAndSet(false, true) && this.handler != null) {
 			this.handler.complete();
 		}
 	}
@@ -267,19 +265,14 @@ public void complete() {
 	 * {@link #send(Object) sending}.
 	 */
 	public void completeWithError(Throwable ex) {
-		if (trySetComplete()) {
+		if (this.complete.compareAndSet(false, true)) {
 			this.failure = ex;
 			if (this.handler != null) {
 				this.handler.completeWithError(ex);
 			}
 		}
 	}
 
-	private boolean trySetComplete() {
-		return (this.state.compareAndSet(State.START, State.COMPLETE) ||
-				(this.state.compareAndSet(State.TIMEOUT, State.COMPLETE)));
-	}
-
 	/**
 	 * Register code to invoke when the async request times out. This method is
 	 * called from a container thread when an async request times out.
@@ -376,7 +369,7 @@ public MediaType getMediaType() {
 	}
 
 
-	private class TimeoutCallback implements Runnable {
+	private class DefaultCallback implements Runnable {
 
 		private final List<Runnable> delegates = new ArrayList<>(1);
 
@@ -386,10 +379,9 @@ public synchronized void addDelegate(Runnable delegate) {
 
 		@Override
 		public void run() {
-			if (ResponseBodyEmitter.this.state.compareAndSet(State.START, State.TIMEOUT)) {
-				for (Runnable delegate : this.delegates) {
-					delegate.run();
-				}
+			ResponseBodyEmitter.this.complete.compareAndSet(false, true);
+			for (Runnable delegate : this.delegates) {
+				delegate.run();
 			}
 		}
 	}
@@ -405,51 +397,11 @@ public synchronized void addDelegate(Consumer<Throwable> callback) {
 
 		@Override
 		public void accept(Throwable t) {
-			if (ResponseBodyEmitter.this.state.compareAndSet(State.START, State.COMPLETE)) {
-				for (Consumer<Throwable> delegate : this.delegates) {
-					delegate.accept(t);
-				}
-			}
-		}
-	}
-
-
-	private class CompletionCallback implements Runnable {
-
-		private final List<Runnable> delegates = new ArrayList<>(1);
-
-		public synchronized void addDelegate(Runnable delegate) {
-			this.delegates.add(delegate);
-		}
-
-		@Override
-		public void run() {
-			if (ResponseBodyEmitter.this.state.compareAndSet(State.START, State.COMPLETE)) {
-				for (Runnable delegate : this.delegates) {
-					delegate.run();
-				}
+			ResponseBodyEmitter.this.complete.compareAndSet(false, true);
+			for(Consumer<Throwable> delegate : this.delegates) {
+				delegate.accept(t);
 			}
 		}
 	}
 
-
-	/**
-	 * Represents a state for {@link ResponseBodyEmitter}.
-	 * <p><pre>
-	 *     START ----+
-	 *       |       |
-	 *       v       |
-	 *    TIMEOUT    |
-	 *       |       |
-	 *       v       |
-	 *   COMPLETE <--+
-	 * </pre>
-	 * @since 6.2.4
-	 */
-	private enum State {
-		START,
-		TIMEOUT, // handling a timeout
-		COMPLETE
-	}
-
 }