From a3636affa2851b2123f36de240e31d419363e6c4 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Mon, 2 Oct 2023 01:16:26 +0100
Subject: [PATCH] ResponseBodyEmitter allows complete after non-IOException

Closes gh-30687
---
 .../annotation/ResponseBodyEmitter.java       | 16 +++----
 .../annotation/ResponseBodyEmitterTests.java  | 46 ++++++++++++++++---
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java b/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java
index b6f534a380..055e5906df 100644
--- a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java
+++ b/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java
@@ -87,7 +87,7 @@ public class ResponseBodyEmitter {
 	 * that may come for example from an application try-catch block on the
 	 * thread of the I/O error.
 	 */
-	private boolean sendFailed;
+	private boolean ioErrorOnSend;
 
 	private final DefaultCallback timeoutCallback = new DefaultCallback();
 
@@ -198,11 +198,10 @@ public class ResponseBodyEmitter {
 				this.handler.send(object, mediaType);
 			}
 			catch (IOException ex) {
-				this.sendFailed = true;
+				this.ioErrorOnSend = true;
 				throw ex;
 			}
 			catch (Throwable ex) {
-				this.sendFailed = true;
 				throw new IllegalStateException("Failed to send " + object, ex);
 			}
 		}
@@ -235,11 +234,10 @@ public class ResponseBodyEmitter {
 				this.handler.send(items);
 			}
 			catch (IOException ex) {
-				this.sendFailed = true;
+				this.ioErrorOnSend = true;
 				throw ex;
 			}
 			catch (Throwable ex) {
-				this.sendFailed = true;
 				throw new IllegalStateException("Failed to send " + items, ex);
 			}
 		}
@@ -257,8 +255,8 @@ public class ResponseBodyEmitter {
 	 * related events such as an error while {@link #send(Object) sending}.
 	 */
 	public synchronized void complete() {
-		// Ignore, after send failure
-		if (this.sendFailed) {
+		// Ignore complete after IO failure on send
+		if (this.ioErrorOnSend) {
 			return;
 		}
 		this.complete = true;
@@ -279,8 +277,8 @@ public class ResponseBodyEmitter {
 	 * {@link #send(Object) sending}.
 	 */
 	public synchronized void completeWithError(Throwable ex) {
-		// Ignore, after send failure
-		if (this.sendFailed) {
+		// Ignore complete after IO failure on send
+		if (this.ioErrorOnSend) {
 			return;
 		}
 		this.complete = true;
diff --git a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java b/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java
index f691ae6461..dc98dc9ad9 100644
--- a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java
+++ b/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java
@@ -92,10 +92,9 @@ public class ResponseBodyEmitterTests {
 	}
 
 	@Test
-	void sendFailsAfterComplete() throws Exception {
+	void sendFailsAfterComplete() {
 		this.emitter.complete();
-		assertThatIllegalStateException().isThrownBy(() ->
-				this.emitter.send("foo"));
+		assertThatIllegalStateException().isThrownBy(() -> this.emitter.send("foo"));
 	}
 
 	@Test
@@ -143,14 +142,47 @@ public class ResponseBodyEmitterTests {
 		verify(this.handler).onCompletion(any());
 		verifyNoMoreInteractions(this.handler);
 
-		IOException failure = new IOException();
-		willThrow(failure).given(this.handler).send("foo", MediaType.TEXT_PLAIN);
-		assertThatIOException().isThrownBy(() ->
-				this.emitter.send("foo", MediaType.TEXT_PLAIN));
+		willThrow(new IOException()).given(this.handler).send("foo", MediaType.TEXT_PLAIN);
+		assertThatIOException().isThrownBy(() -> this.emitter.send("foo", MediaType.TEXT_PLAIN));
 		verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
 		verifyNoMoreInteractions(this.handler);
 	}
 
+	@Test // gh-30687
+	void completeIgnoredAfterIOException() throws Exception {
+		this.emitter.initialize(this.handler);
+		verify(this.handler).onTimeout(any());
+		verify(this.handler).onError(any());
+		verify(this.handler).onCompletion(any());
+		verifyNoMoreInteractions(this.handler);
+
+		willThrow(new IOException()).given(this.handler).send("foo", MediaType.TEXT_PLAIN);
+		assertThatIOException().isThrownBy(() -> this.emitter.send("foo", MediaType.TEXT_PLAIN));
+		verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
+		verifyNoMoreInteractions(this.handler);
+
+		this.emitter.complete();
+		verifyNoMoreInteractions(this.handler);
+	}
+
+	@Test // gh-30687
+	void completeAfterNonIOException() throws Exception {
+		this.emitter.initialize(this.handler);
+		verify(this.handler).onTimeout(any());
+		verify(this.handler).onError(any());
+		verify(this.handler).onCompletion(any());
+		verifyNoMoreInteractions(this.handler);
+
+		willThrow(new IllegalStateException()).given(this.handler).send("foo", MediaType.TEXT_PLAIN);
+		assertThatIllegalStateException().isThrownBy(() -> this.emitter.send("foo", MediaType.TEXT_PLAIN));
+		verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
+		verifyNoMoreInteractions(this.handler);
+
+		this.emitter.complete();
+		verify(this.handler).complete();
+		verifyNoMoreInteractions(this.handler);
+	}
+
 	@Test
 	void onTimeoutBeforeHandlerInitialized() throws Exception {
 		Runnable runnable = mock();