Allow multiple listeners on ResponseBodyEmitter

Prior to this commit, `ResponseBodyEmitter` woud accept a single
`Runnable` callback on each of its `onTimeout`, `onError` or
`onCompletion` methods. This would limit the developers' ability to
register multiple sets of callbacks: one for managing the publication of
streaming values, another one for managing other concerns like
keep-alive signals to maintain the connection.

This commit now allows multiple calls to `onTimeout`, `onError` and
`onCompletion` and will register all callbacks accordingly.

Closes gh-33356

Author: bclozel

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 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.
@@ -17,7 +17,9 @@
 package org.springframework.web.servlet.mvc.method.annotation;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Set;
 import java.util.function.Consumer;
 
@@ -59,6 +61,7 @@
  *
  * @author Rossen Stoyanchev
  * @author Juergen Hoeller
+ * @author Brian Clozel
  * @since 4.2
  */
 public class ResponseBodyEmitter {
@@ -271,29 +274,32 @@ public synchronized void completeWithError(Throwable ex) {
 	/**
 	 * Register code to invoke when the async request times out. This method is
 	 * 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.setDelegate(callback);
+		this.timeoutCallback.addDelegate(callback);
 	}
 
 	/**
 	 * Register code to invoke for an error during async request processing.
 	 * This method is called from a container thread when an error occurred
 	 * while processing an async request.
+	 * <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.setDelegate(callback);
+		this.errorCallback.addDelegate(callback);
 	}
 
 	/**
 	 * Register code to invoke when the async request completes. This method is
 	 * called from a container thread when an async request completed for any
 	 * reason including timeout and network error. This method is useful for
 	 * 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.setDelegate(callback);
+		this.completionCallback.addDelegate(callback);
 	}
 
 
@@ -363,37 +369,35 @@ public MediaType getMediaType() {
 
 	private class DefaultCallback implements Runnable {
 
-		@Nullable
-		private Runnable delegate;
+		private List<Runnable> delegates = new ArrayList<>(1);
 
-		public void setDelegate(Runnable delegate) {
-			this.delegate = delegate;
+		public void addDelegate(Runnable delegate) {
+			this.delegates.add(delegate);
 		}
 
 		@Override
 		public void run() {
 			ResponseBodyEmitter.this.complete = true;
-			if (this.delegate != null) {
-				this.delegate.run();
+			for (Runnable delegate : this.delegates) {
+				delegate.run();
 			}
 		}
 	}
 
 
 	private class ErrorCallback implements Consumer<Throwable> {
 
-		@Nullable
-		private Consumer<Throwable> delegate;
+		private List<Consumer<Throwable>> delegates = new ArrayList<>(1);
 
-		public void setDelegate(Consumer<Throwable> callback) {
-			this.delegate = callback;
+		public void addDelegate(Consumer<Throwable> callback) {
+			this.delegates.add(callback);
 		}
 
 		@Override
 		public void accept(Throwable t) {
 			ResponseBodyEmitter.this.complete = true;
-			if (this.delegate != null) {
-				this.delegate.accept(t);
+			for(Consumer<Throwable> delegate : this.delegates) {
+				delegate.accept(t);
 			}
 		}
 	}

Diff for: spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java

@@ -17,6 +17,7 @@
 package org.springframework.web.servlet.mvc.method.annotation;
 
 import java.io.IOException;
+import java.util.function.Consumer;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -31,6 +32,7 @@
 import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anySet;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.BDDMockito.willThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -41,6 +43,7 @@
  *
  * @author Rossen Stoyanchev
  * @author Tomasz Nurkiewicz
+ * @author Brian Clozel
  */
 @ExtendWith(MockitoExtension.class)
 public class ResponseBodyEmitterTests {
@@ -197,6 +200,25 @@ void onTimeoutAfterHandlerInitialized() throws Exception {
 		verify(runnable).run();
 	}
 
+	@Test
+	void multipleOnTimeoutCallbacks() throws Exception {
+		this.emitter.initialize(this.handler);
+
+		ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+		verify(this.handler).onTimeout(captor.capture());
+		verify(this.handler).onCompletion(any());
+
+		Runnable first = mock();
+		Runnable second = mock();
+		this.emitter.onTimeout(first);
+		this.emitter.onTimeout(second);
+
+		assertThat(captor.getValue()).isNotNull();
+		captor.getValue().run();
+		verify(first).run();
+		verify(second).run();
+	}
+
 	@Test
 	void onCompletionBeforeHandlerInitialized() throws Exception {
 		Runnable runnable = mock();
@@ -228,4 +250,42 @@ void onCompletionAfterHandlerInitialized() throws Exception {
 		verify(runnable).run();
 	}
 
+	@Test
+	void multipleOnCompletionCallbacks() throws Exception {
+		this.emitter.initialize(this.handler);
+
+		ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+		verify(this.handler).onTimeout(any());
+		verify(this.handler).onCompletion(captor.capture());
+
+		Runnable first = mock();
+		Runnable second = mock();
+		this.emitter.onCompletion(first);
+		this.emitter.onCompletion(second);
+
+		assertThat(captor.getValue()).isNotNull();
+		captor.getValue().run();
+		verify(first).run();
+		verify(second).run();
+	}
+
+	@Test
+	void multipleOnErrorCallbacks() throws Exception {
+		this.emitter.initialize(this.handler);
+
+		ArgumentCaptor<Consumer<Throwable>> captor = ArgumentCaptor.<Consumer<Throwable>, Consumer>forClass(Consumer.class);
+		verify(this.handler).onError(captor.capture());
+
+		Consumer<Throwable> first = mock();
+		Consumer<Throwable> second = mock();
+		this.emitter.onError(first);
+		this.emitter.onError(second);
+
+		assertThat(captor.getValue()).isNotNull();
+		IllegalStateException illegalStateException = new IllegalStateException();
+		captor.getValue().accept(illegalStateException);
+		verify(first).accept(eq(illegalStateException));
+		verify(second).accept(eq(illegalStateException));
+	}
+
 }