Avoid thread pinning in SseEmitter, ResponseBodyEmitter

Closes gh-35423

Signed-off-by: Taeik Lim <[email protected]>

Author: acktsap

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

@@ -21,6 +21,8 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 
 import org.springframework.http.MediaType;
@@ -62,6 +64,7 @@
  * @author Rossen Stoyanchev
  * @author Juergen Hoeller
  * @author Brian Clozel
+ * @author Taeik Lim
  * @since 4.2
  */
 public class ResponseBodyEmitter {
@@ -88,6 +91,8 @@ public class ResponseBodyEmitter {
 
 	private final DefaultCallback completionCallback = new DefaultCallback();
 
+	/** Guards access to write operations on the response. */
+	protected final Lock writeLock = new ReentrantLock();
 
 	/**
 	 * Create a new ResponseBodyEmitter instance.
@@ -117,36 +122,48 @@ public Long getTimeout() {
 	}
 
 
-	synchronized void initialize(Handler handler) throws IOException {
-		this.handler = handler;
-
+	void initialize(Handler handler) throws IOException {
+		this.writeLock.lock();
 		try {
-			sendInternal(this.earlySendAttempts);
-		}
-		finally {
-			this.earlySendAttempts.clear();
-		}
+			this.handler = handler;
+
+			try {
+				sendInternal(this.earlySendAttempts);
+			}
+			finally {
+				this.earlySendAttempts.clear();
+			}
 
-		if (this.complete) {
-			if (this.failure != null) {
-				this.handler.completeWithError(this.failure);
+			if (this.complete) {
+				if (this.failure != null) {
+					this.handler.completeWithError(this.failure);
+				}
+				else {
+					this.handler.complete();
+				}
 			}
 			else {
-				this.handler.complete();
+				this.handler.onTimeout(this.timeoutCallback);
+				this.handler.onError(this.errorCallback);
+				this.handler.onCompletion(this.completionCallback);
 			}
 		}
-		else {
-			this.handler.onTimeout(this.timeoutCallback);
-			this.handler.onError(this.errorCallback);
-			this.handler.onCompletion(this.completionCallback);
+		finally {
+			this.writeLock.unlock();
 		}
 	}
 
-	synchronized void initializeWithError(Throwable ex) {
-		this.complete = true;
-		this.failure = ex;
-		this.earlySendAttempts.clear();
-		this.errorCallback.accept(ex);
+	void initializeWithError(Throwable ex) {
+		this.writeLock.lock();
+		try {
+			this.complete = true;
+			this.failure = ex;
+			this.earlySendAttempts.clear();
+			this.errorCallback.accept(ex);
+		}
+		finally {
+			this.writeLock.unlock();
+		}
 	}
 
 	/**
@@ -183,22 +200,28 @@ public void send(Object object) throws IOException {
 	 * @throws IOException raised when an I/O error occurs
 	 * @throws java.lang.IllegalStateException wraps any other errors
 	 */
-	public synchronized void send(Object object, @Nullable MediaType mediaType) throws IOException {
+	public void send(Object object, @Nullable MediaType mediaType) throws IOException {
 		Assert.state(!this.complete, () -> "ResponseBodyEmitter has already completed" +
 				(this.failure != null ? " with error: " + this.failure : ""));
-		if (this.handler != null) {
-			try {
-				this.handler.send(object, mediaType);
-			}
-			catch (IOException ex) {
-				throw ex;
+		this.writeLock.lock();
+		try {
+			if (this.handler != null) {
+				try {
+					this.handler.send(object, mediaType);
+				}
+				catch (IOException ex) {
+					throw ex;
+				}
+				catch (Throwable ex) {
+					throw new IllegalStateException("Failed to send " + object, ex);
+				}
 			}
-			catch (Throwable ex) {
-				throw new IllegalStateException("Failed to send " + object, ex);
+			else {
+				this.earlySendAttempts.add(new DataWithMediaType(object, mediaType));
 			}
 		}
-		else {
-			this.earlySendAttempts.add(new DataWithMediaType(object, mediaType));
+		finally {
+			this.writeLock.unlock();
 		}
 	}
 
@@ -211,10 +234,16 @@ public synchronized void send(Object object, @Nullable MediaType mediaType) thro
 	 * @throws java.lang.IllegalStateException wraps any other errors
 	 * @since 6.0.12
 	 */
-	public synchronized void send(Set<DataWithMediaType> items) throws IOException {
+	public void send(Set<DataWithMediaType> items) throws IOException {
 		Assert.state(!this.complete, () -> "ResponseBodyEmitter has already completed" +
 				(this.failure != null ? " with error: " + this.failure : ""));
-		sendInternal(items);
+		this.writeLock.lock();
+		try {
+			sendInternal(items);
+		}
+		finally {
+			this.writeLock.unlock();
+		}
 	}
 
 	private void sendInternal(Set<DataWithMediaType> items) throws IOException {
@@ -245,10 +274,16 @@ private void sendInternal(Set<DataWithMediaType> items) throws IOException {
 	 * to complete request processing. It should not be used after container
 	 * related events such as an error while {@link #send(Object) sending}.
 	 */
-	public synchronized void complete() {
-		this.complete = true;
-		if (this.handler != null) {
-			this.handler.complete();
+	public void complete() {
+		this.writeLock.lock();
+		try {
+			this.complete = true;
+			if (this.handler != null) {
+				this.handler.complete();
+			}
+		}
+		finally {
+			this.writeLock.unlock();
 		}
 	}
 
@@ -263,11 +298,17 @@ public synchronized void complete() {
 	 * container related events such as an error while
 	 * {@link #send(Object) sending}.
 	 */
-	public synchronized void completeWithError(Throwable ex) {
-		this.complete = true;
-		this.failure = ex;
-		if (this.handler != null) {
-			this.handler.completeWithError(ex);
+	public void completeWithError(Throwable ex) {
+		this.writeLock.lock();
+		try {
+			this.complete = true;
+			this.failure = ex;
+			if (this.handler != null) {
+				this.handler.completeWithError(ex);
+			}
+		}
+		finally {
+			this.writeLock.unlock();
 		}
 	}
 
@@ -276,8 +317,14 @@ public synchronized void completeWithError(Throwable ex) {
 	 * called from a container thread when an async request times out.
 	 * <p>As of 6.2, one can register multiple callbacks for this event.
 	 */
-	public synchronized void onTimeout(Runnable callback) {
-		this.timeoutCallback.addDelegate(callback);
+	public void onTimeout(Runnable callback) {
+		this.writeLock.lock();
+		try {
+			this.timeoutCallback.addDelegate(callback);
+		}
+		finally {
+			this.writeLock.unlock();
+		}
 	}
 
 	/**
@@ -287,8 +334,14 @@ public synchronized void onTimeout(Runnable callback) {
 	 * <p>As of 6.2, one can register multiple callbacks for this event.
 	 * @since 5.0
 	 */
-	public synchronized void onError(Consumer<Throwable> callback) {
-		this.errorCallback.addDelegate(callback);
+	public void onError(Consumer<Throwable> callback) {
+		this.writeLock.lock();
+		try {
+			this.errorCallback.addDelegate(callback);
+		}
+		finally {
+			this.writeLock.unlock();
+		}
 	}
 
 	/**
@@ -298,8 +351,14 @@ public synchronized void onError(Consumer<Throwable> callback) {
 	 * detecting that a {@code ResponseBodyEmitter} instance is no longer usable.
 	 * <p>As of 6.2, one can register multiple callbacks for this event.
 	 */
-	public synchronized void onCompletion(Runnable callback) {
-		this.completionCallback.addDelegate(callback);
+	public void onCompletion(Runnable callback) {
+		this.writeLock.lock();
+		try {
+			this.completionCallback.addDelegate(callback);
+		}
+		finally {
+			this.writeLock.unlock();
+		}
 	}
 
 

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

@@ -21,8 +21,6 @@
 import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.Set;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
@@ -46,10 +44,6 @@ public class SseEmitter extends ResponseBodyEmitter {
 
 	private static final MediaType TEXT_PLAIN = new MediaType("text", "plain", StandardCharsets.UTF_8);
 
-	/** Guards access to write operations on the response. */
-	private final Lock writeLock = new ReentrantLock();
-
-
 	/**
 	 * Create a new SseEmitter instance.
 	 */